vortex_btrblocks/schemes/
decimal.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Decimal compression scheme using byte-part decomposition.
5
6use vortex_array::ArrayId;
7use vortex_array::ArrayRef;
8use vortex_array::Canonical;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::arrays::DecimalArray;
12use vortex_array::arrays::decimal::narrowed_decimal;
13use vortex_array::dtype::DecimalType;
14use vortex_compressor::scheme::CompressionEstimate;
15use vortex_compressor::scheme::EstimateVerdict;
16use vortex_decimal_byte_parts::DecimalByteParts;
17use vortex_decimal_byte_parts::DecimalBytePartsSlots;
18use vortex_decimal_byte_parts::decimal_byte_parts_v1_id;
19use vortex_decimal_byte_parts::decimal_byte_parts_v2_id;
20use vortex_decimal_byte_parts::split_decimal;
21use vortex_error::VortexResult;
22
23use crate::ArrayAndStats;
24use crate::CascadingCompressor;
25use crate::CompressorContext;
26use crate::Scheme;
27use crate::SchemeExt;
28
29#[derive(Debug, Copy, Clone, PartialEq, Eq)]
30enum DecimalSchemeMode {
31    V1,
32    V2,
33}
34
35/// The decimal scheme that only produces v1-compatible DBP arrays (does not compress wide
36/// decimals).
37pub(crate) static DECIMAL_V1: DecimalScheme = DecimalScheme::v1();
38
39/// The decimal scheme that produces both v1 and v2 compatible DBP arrays (compresses wide
40/// decimals).
41pub(crate) static DECIMAL_V2: DecimalScheme = DecimalScheme::v2();
42
43/// Compression scheme for decimal arrays via byte-part decomposition.
44///
45/// Narrows the decimal to the smallest integer type and compresses its byte parts independently.
46/// The v1 mode leaves values wider than `i64` canonical; v2 splits them into a signed most
47/// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1
48/// in either mode, while arrays with lower parts serialize as v2.
49///
50/// The default uses v1. [`refine`](Scheme::refine) picks v2 when the v2 ID is allowed and v1
51/// otherwise.
52#[derive(Debug, Copy, Clone, PartialEq, Eq)]
53pub struct DecimalScheme {
54    mode: DecimalSchemeMode,
55}
56
57impl DecimalScheme {
58    /// Creates a decimal scheme configured for v1, disallowing splitting of wide decimals.
59    ///
60    /// Values that remain wider than `i64` after narrowing stay canonical.
61    pub const fn v1() -> Self {
62        Self {
63            mode: DecimalSchemeMode::V1,
64        }
65    }
66
67    /// Creates a decimal scheme configured for v2, allowing splitting of wide decimals.
68    pub const fn v2() -> Self {
69        Self {
70            mode: DecimalSchemeMode::V2,
71        }
72    }
73}
74
75impl Default for DecimalScheme {
76    fn default() -> Self {
77        Self::v1()
78    }
79}
80
81impl Scheme for DecimalScheme {
82    fn scheme_name(&self) -> &'static str {
83        "vortex.decimal.byte_parts"
84    }
85
86    fn matches(&self, canonical: &Canonical) -> bool {
87        matches!(canonical, Canonical::Decimal(_))
88    }
89
90    fn produced_encodings(&self) -> Vec<ArrayId> {
91        match self.mode {
92            DecimalSchemeMode::V1 => vec![decimal_byte_parts_v1_id()],
93            DecimalSchemeMode::V2 => {
94                vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]
95            }
96        }
97    }
98
99    fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme {
100        if allowed(&decimal_byte_parts_v2_id()) {
101            &DECIMAL_V2
102        } else {
103            &DECIMAL_V1
104        }
105    }
106
107    /// Children: msp=0, then up to three lower parts in v2 mode.
108    fn num_children(&self) -> usize {
109        match self.mode {
110            DecimalSchemeMode::V1 => 1,
111            DecimalSchemeMode::V2 => 4,
112        }
113    }
114
115    fn expected_compression_ratio(
116        &self,
117        _data: &ArrayAndStats,
118        _compress_ctx: CompressorContext,
119        _exec_ctx: &mut ExecutionCtx,
120    ) -> CompressionEstimate {
121        // Decimal compression is almost always beneficial (narrowing + primitive compression).
122        CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse)
123    }
124
125    fn compress(
126        &self,
127        compressor: &CascadingCompressor,
128        data: &ArrayAndStats,
129        compress_ctx: CompressorContext,
130        exec_ctx: &mut ExecutionCtx,
131    ) -> VortexResult<ArrayRef> {
132        let decimal = data.array().clone().execute::<DecimalArray>(exec_ctx)?;
133        let decimal = narrowed_decimal(decimal);
134        if self.mode == DecimalSchemeMode::V1
135            && matches!(decimal.values_type(), DecimalType::I128 | DecimalType::I256)
136        {
137            return Ok(decimal.into_array());
138        }
139
140        let parts = split_decimal(&decimal, exec_ctx)?;
141        let msp = compressor.compress_child(
142            &parts.msp,
143            &compress_ctx,
144            self.id(),
145            DecimalBytePartsSlots::MSP,
146            exec_ctx,
147        )?;
148        let lower_parts = parts
149            .lower_parts
150            .iter()
151            .enumerate()
152            .map(|(idx, part)| {
153                compressor.compress_child(
154                    part,
155                    &compress_ctx,
156                    self.id(),
157                    DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx,
158                    exec_ctx,
159                )
160            })
161            .collect::<VortexResult<Vec<_>>>()?;
162
163        DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal.decimal_dtype())
164            .map(IntoArray::into_array)
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use rstest::rstest;
171    use vortex_array::ArrayId;
172    use vortex_decimal_byte_parts::decimal_byte_parts_v1_id;
173    use vortex_decimal_byte_parts::decimal_byte_parts_v2_id;
174
175    use super::DECIMAL_V1;
176    use super::DECIMAL_V2;
177    use crate::Scheme;
178    use crate::SchemeExt;
179
180    /// Both variants refine to v2 exactly when the v2 ID is allowed.
181    #[rstest]
182    #[case::neither(false, false, false)]
183    #[case::v1(true, false, false)]
184    #[case::v2_only(false, true, true)]
185    #[case::both(true, true, true)]
186    fn refine_picks_v2_when_v2_id_is_allowed(
187        #[case] allow_v1: bool,
188        #[case] allow_v2: bool,
189        #[case] expect_v2: bool,
190        #[values(&DECIMAL_V1, &DECIMAL_V2)] scheme: &'static dyn Scheme,
191    ) {
192        let allowed = |id: &ArrayId| {
193            (allow_v1 && *id == decimal_byte_parts_v1_id())
194                || (allow_v2 && *id == decimal_byte_parts_v2_id())
195        };
196        let refined = scheme.refine(&allowed);
197        assert_eq!(refined.id(), scheme.id());
198        let expected = if expect_v2 {
199            vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]
200        } else {
201            vec![decimal_byte_parts_v1_id()]
202        };
203        assert_eq!(refined.produced_encodings(), expected);
204    }
205}
vortex_btrblocks/schemes/decimal.rs

vortex_btrblocks/schemes/
decimal.rs
