vortex_btrblocks/schemes/integer/
for_.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Frame of Reference integer encoding.
5
6use vortex_array::ArrayId;
7use vortex_array::ArrayRef;
8use vortex_array::Canonical;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::VTable;
12use vortex_array::arrays::PrimitiveArray;
13use vortex_compressor::builtins::BinaryDictScheme;
14use vortex_compressor::builtins::FloatDictScheme;
15use vortex_compressor::builtins::IntDictScheme;
16use vortex_compressor::builtins::StringDictScheme;
17use vortex_compressor::scheme::AncestorExclusion;
18use vortex_compressor::scheme::ChildSelection;
19use vortex_compressor::scheme::CompressionEstimate;
20use vortex_compressor::scheme::EstimateVerdict;
21use vortex_error::VortexExpect;
22use vortex_error::VortexResult;
23use vortex_fastlanes::FoR;
24use vortex_fastlanes::FoRArrayExt;
25use vortex_fastlanes::FoRArraySlotsExt;
26
27use super::BitPackingScheme;
28use crate::ArrayAndStats;
29use crate::CascadingCompressor;
30use crate::CompressorContext;
31use crate::Scheme;
32use crate::SchemeExt;
33
34/// Frame of Reference encoding.
35#[derive(Debug, Copy, Clone, PartialEq, Eq)]
36pub struct FoRScheme;
37
38impl Scheme for FoRScheme {
39    fn scheme_name(&self) -> &'static str {
40        "vortex.int.for"
41    }
42
43    fn matches(&self, canonical: &Canonical) -> bool {
44        canonical.dtype().is_int()
45    }
46
47    fn produced_encodings(&self) -> Vec<ArrayId> {
48        vec![FoR.id()]
49    }
50
51    /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op.
52    fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
53        vec![
54            AncestorExclusion {
55                ancestor: IntDictScheme.id(),
56                children: ChildSelection::One(1),
57            },
58            AncestorExclusion {
59                ancestor: FloatDictScheme.id(),
60                children: ChildSelection::One(1),
61            },
62            AncestorExclusion {
63                ancestor: StringDictScheme.id(),
64                children: ChildSelection::One(1),
65            },
66            AncestorExclusion {
67                ancestor: BinaryDictScheme.id(),
68                children: ChildSelection::One(1),
69            },
70        ]
71    }
72
73    fn expected_compression_ratio(
74        &self,
75        data: &ArrayAndStats,
76        compress_ctx: CompressorContext,
77        exec_ctx: &mut ExecutionCtx,
78    ) -> CompressionEstimate {
79        // FoR only subtracts the min. Without further compression (e.g. BitPacking), the output is
80        // the same size.
81        if compress_ctx.finished_cascading() {
82            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
83        }
84        let stats = data.integer_stats(exec_ctx);
85
86        // Only apply when the min is not already zero.
87        if stats.erased().min_is_zero() {
88            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
89        }
90
91        // Difference between max and min.
92        let for_bitwidth = match stats.erased().max_minus_min().checked_ilog2() {
93            Some(l) => l + 1,
94            // If max-min == 0, the we should be compressing this as a constant array.
95            None => return CompressionEstimate::Verdict(EstimateVerdict::Skip),
96        };
97
98        // If BitPacking can be applied (only non-negative values) and FoR doesn't reduce bit width
99        // compared to BitPacking, don't use FoR since it has a small amount of overhead (storing
100        // the reference) for effectively no benefits.
101        if let Some(max_log) = stats
102            .erased()
103            .max_ilog2()
104            // Only skip FoR when min >= 0, otherwise BitPacking can't be applied without ZigZag.
105            .filter(|_| !stats.erased().min_is_negative())
106        {
107            let bitpack_bitwidth = max_log + 1;
108            if for_bitwidth >= bitpack_bitwidth {
109                return CompressionEstimate::Verdict(EstimateVerdict::Skip);
110            }
111        }
112
113        let full_width: u32 = data
114            .array_as_primitive()
115            .ptype()
116            .bit_width()
117            .try_into()
118            .vortex_expect("bit width must fit in u32");
119
120        CompressionEstimate::Verdict(EstimateVerdict::Ratio(
121            full_width as f64 / for_bitwidth as f64,
122        ))
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 primitive = data.array().clone().execute::<PrimitiveArray>(exec_ctx)?;
133        let for_array = FoR::encode(primitive, exec_ctx)?;
134        let biased = for_array
135            .encoded()
136            .clone()
137            .execute::<PrimitiveArray>(exec_ctx)?;
138
139        // Immediately bitpack. If any other scheme was preferable, it would be chosen instead
140        // of bitpacking.
141        // NOTE: we could delegate in the future if we had another downstream codec that performs
142        //  as well.
143        let leaf_ctx = compress_ctx.clone().as_leaf();
144        let biased_data =
145            ArrayAndStats::new(biased.into_array(), compress_ctx.merged_stats_options());
146        let compressed = BitPackingScheme.compress(compressor, &biased_data, leaf_ctx, exec_ctx)?;
147
148        // TODO(connor): This should really be `new_unchecked`.
149        let for_compressed = FoR::try_new(compressed, for_array.reference_scalar().clone())?;
150        for_compressed
151            .as_ref()
152            .statistics()
153            .inherit_from(for_array.as_ref().statistics());
154
155        Ok(for_compressed.into_array())
156    }
157}
vortex_btrblocks/schemes/integer/for_.rs

vortex_btrblocks/schemes/integer/
for_.rs
