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

vortex_btrblocks/schemes/
temporal.rs
