vortex_btrblocks/schemes/
temporal.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Temporal compression scheme using datetime-part decomposition.
5
6use vortex_array::ArrayRef;
7use vortex_array::Canonical;
8use vortex_array::IntoArray;
9use vortex_array::ToCanonical;
10use vortex_array::aggregate_fn::fns::is_constant::is_constant;
11use vortex_array::arrays::ConstantArray;
12use vortex_array::arrays::TemporalArray;
13use vortex_array::arrays::primitive::PrimitiveArrayExt;
14use vortex_array::dtype::extension::Matcher;
15use vortex_array::extension::datetime::AnyTemporal;
16use vortex_array::extension::datetime::TemporalMetadata;
17use vortex_compressor::estimate::CompressionEstimate;
18use vortex_datetime_parts::DateTimeParts;
19use vortex_datetime_parts::TemporalParts;
20use vortex_datetime_parts::split_temporal;
21use vortex_error::VortexResult;
22
23use crate::ArrayAndStats;
24use crate::CascadingCompressor;
25use crate::CompressorContext;
26use crate::Scheme;
27use crate::SchemeExt;
28
29/// Compression scheme for temporal timestamp arrays via datetime-part decomposition.
30///
31/// Splits timestamps into days, seconds, and subseconds components, compresses each
32/// independently, and wraps the result in a `DateTimePartsArray`.
33#[derive(Debug, Copy, Clone, PartialEq, Eq)]
34pub struct TemporalScheme;
35
36impl Scheme for TemporalScheme {
37    fn scheme_name(&self) -> &'static str {
38        "vortex.ext.temporal"
39    }
40
41    fn matches(&self, canonical: &Canonical) -> bool {
42        let Canonical::Extension(ext) = canonical else {
43            return false;
44        };
45
46        let ext_dtype = ext.ext_dtype();
47
48        matches!(
49            AnyTemporal::try_match(ext_dtype),
50            Some(TemporalMetadata::Timestamp(..))
51        )
52    }
53
54    /// Children: days=0, seconds=1, subseconds=2.
55    fn num_children(&self) -> usize {
56        3
57    }
58
59    fn expected_compression_ratio(
60        &self,
61        _data: &mut ArrayAndStats,
62        _ctx: CompressorContext,
63    ) -> CompressionEstimate {
64        // Temporal compression (splitting into parts) is almost always beneficial.
65        CompressionEstimate::AlwaysUse
66    }
67
68    fn compress(
69        &self,
70        compressor: &CascadingCompressor,
71        data: &mut ArrayAndStats,
72        ctx: CompressorContext,
73    ) -> VortexResult<ArrayRef> {
74        let array = data.array().clone();
75        let ext_array = array.to_extension();
76        let temporal_array = TemporalArray::try_from(ext_array.clone().into_array())?;
77
78        // Check for constant array and return early if so.
79        let is_constant = is_constant(
80            &ext_array.clone().into_array(),
81            &mut compressor.execution_ctx(),
82        )?;
83
84        if is_constant {
85            return Ok(ConstantArray::new(ext_array.scalar_at(0)?, ext_array.len()).into_array());
86        }
87
88        let dtype = temporal_array.dtype().clone();
89        let TemporalParts {
90            days,
91            seconds,
92            subseconds,
93        } = split_temporal(temporal_array)?;
94
95        let days = compressor.compress_child(
96            &days.to_primitive().narrow()?.into_array(),
97            &ctx,
98            self.id(),
99            0,
100        )?;
101        let seconds = compressor.compress_child(
102            &seconds.to_primitive().narrow()?.into_array(),
103            &ctx,
104            self.id(),
105            1,
106        )?;
107        let subseconds = compressor.compress_child(
108            &subseconds.to_primitive().narrow()?.into_array(),
109            &ctx,
110            self.id(),
111            2,
112        )?;
113
114        Ok(DateTimeParts::try_new(dtype, days, seconds, subseconds)?.into_array())
115    }
116}
vortex_btrblocks/schemes/temporal.rs

vortex_btrblocks/schemes/
temporal.rs
