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

vortex_btrblocks/schemes/integer/
for_.rs
