vortex_btrblocks/schemes/integer/
bitpacking.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! BitPacking integer encoding.
5
6use itertools::Itertools;
7use vortex_array::ArrayId;
8use vortex_array::ArrayRef;
9use vortex_array::Canonical;
10use vortex_array::ExecutionCtx;
11use vortex_array::IntoArray;
12use vortex_array::VTable;
13use vortex_array::arrays::Patched;
14use vortex_array::arrays::PrimitiveArray;
15use vortex_array::arrays::patched::use_experimental_patches;
16use vortex_array::arrays::primitive::PrimitiveArrayExt;
17use vortex_compressor::scheme::CompressionEstimate;
18use vortex_compressor::scheme::DeferredEstimate;
19use vortex_compressor::scheme::EstimateVerdict;
20use vortex_error::VortexResult;
21use vortex_fastlanes::BitPacked;
22use vortex_fastlanes::BitPackedArray;
23use vortex_fastlanes::BitWidths;
24use vortex_fastlanes::bitpack_compress::bit_width_histogram;
25use vortex_fastlanes::bitpack_compress::bitpack_encode;
26use vortex_fastlanes::bitpack_compress::bitpack_encode_blocked;
27use vortex_fastlanes::bitpack_compress::find_best_bit_width;
28use vortex_fastlanes::bitpack_compress::find_best_bit_widths;
29use vortex_fastlanes::bitpacked_v1_id;
30use vortex_fastlanes::bitpacked_v2_id;
31
32use crate::ArrayAndStats;
33use crate::CascadingCompressor;
34use crate::CompressorContext;
35use crate::Scheme;
36use crate::SchemeExt;
37use crate::compress_patches;
38
39#[derive(Debug, Copy, Clone, PartialEq, Eq)]
40enum BitPackingSchemeMode {
41    V1,
42    V2,
43}
44
45/// The BitPacking scheme that only produces arrays with one bit width.
46pub(crate) static BITPACKING_V1: BitPackingScheme = BitPackingScheme::v1();
47
48/// The BitPacking scheme that may give each 1024-element block its own bit width.
49pub(crate) static BITPACKING_V2: BitPackingScheme = BitPackingScheme::v2();
50
51/// BitPacking encoding for non-negative integers.
52///
53/// The v1 mode packs every value at one bit width, and serializes as `fastlanes.bitpacked`. The v2
54/// mode chooses a width for each 1024-element block. Uniform widths use the v1 format; varying
55/// widths serialize as `fastlanes.bitpacked.v2`.
56///
57/// The default uses v1. [`refine`](Scheme::refine) picks v2 when the v2 ID is allowed and v1
58/// otherwise.
59#[derive(Debug, Copy, Clone, PartialEq, Eq)]
60pub struct BitPackingScheme {
61    mode: BitPackingSchemeMode,
62}
63
64impl BitPackingScheme {
65    /// Creates a BitPacking scheme configured for v1, which uses one bit width.
66    pub const fn v1() -> Self {
67        Self {
68            mode: BitPackingSchemeMode::V1,
69        }
70    }
71
72    /// Creates a BitPacking scheme configured for v2, which may use one bit width per block.
73    pub const fn v2() -> Self {
74        Self {
75            mode: BitPackingSchemeMode::V2,
76        }
77    }
78}
79
80impl Default for BitPackingScheme {
81    fn default() -> Self {
82        Self::v1()
83    }
84}
85
86impl Scheme for BitPackingScheme {
87    fn scheme_name(&self) -> &'static str {
88        "vortex.int.bitpacking"
89    }
90
91    fn matches(&self, canonical: &Canonical) -> bool {
92        canonical.dtype().is_int()
93    }
94
95    fn produced_encodings(&self) -> Vec<ArrayId> {
96        let mut encodings = match self.mode {
97            // Global-width arrays serialize under the frozen v1 ID.
98            BitPackingSchemeMode::V1 => vec![bitpacked_v1_id()],
99            BitPackingSchemeMode::V2 => vec![bitpacked_v1_id(), bitpacked_v2_id()],
100        };
101        if use_experimental_patches() {
102            encodings.push(Patched.id());
103        }
104        encodings
105    }
106
107    fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme {
108        if allowed(&bitpacked_v2_id()) {
109            &BITPACKING_V2
110        } else {
111            &BITPACKING_V1
112        }
113    }
114
115    /// Children: block offsets=0 in v2 mode.
116    fn num_children(&self) -> usize {
117        match self.mode {
118            BitPackingSchemeMode::V1 => 0,
119            BitPackingSchemeMode::V2 => 1,
120        }
121    }
122
123    fn expected_compression_ratio(
124        &self,
125        data: &ArrayAndStats,
126        _compress_ctx: CompressorContext,
127        exec_ctx: &mut ExecutionCtx,
128    ) -> CompressionEstimate {
129        let stats = data.integer_stats(exec_ctx);
130
131        // BitPacking only works for non-negative values.
132        if stats.erased().min_is_negative() {
133            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
134        }
135
136        CompressionEstimate::Deferred(DeferredEstimate::Sample)
137    }
138
139    fn compress(
140        &self,
141        compressor: &CascadingCompressor,
142        data: &ArrayAndStats,
143        compress_ctx: CompressorContext,
144        exec_ctx: &mut ExecutionCtx,
145    ) -> VortexResult<ArrayRef> {
146        let primitive_array = data.array_as_primitive().into_owned();
147        let packed = match self.mode {
148            BitPackingSchemeMode::V1 => compress_global(&primitive_array, exec_ctx)?,
149            BitPackingSchemeMode::V2 => compress_blocked(&primitive_array, exec_ctx)?,
150        };
151        let Some(packed) = packed else {
152            return Ok(primitive_array.into_array());
153        };
154
155        let packed_stats = packed.statistics().to_owned();
156        let ptype = packed.dtype().as_ptype();
157        let mut parts = BitPacked::into_parts(packed);
158
159        let patches = if use_experimental_patches() {
160            parts.patches.take()
161        } else {
162            parts.patches = parts
163                .patches
164                .take()
165                .map(|p| compress_patches(p, exec_ctx))
166                .transpose()?;
167            None
168        };
169
170        let array = match parts.bit_widths {
171            BitWidths::Global(bw) => BitPacked::try_new(
172                parts.packed,
173                ptype,
174                parts.validity,
175                parts.patches,
176                bw,
177                parts.len,
178                parts.offset,
179            )?,
180            BitWidths::Blocked(block_offsets) => {
181                let block_offsets = compressor.compress_child(
182                    &block_offsets,
183                    &compress_ctx,
184                    self.id(),
185                    0,
186                    exec_ctx,
187                )?;
188                BitPacked::try_new_with_block_offsets(
189                    parts.packed,
190                    ptype,
191                    parts.validity,
192                    parts.patches,
193                    block_offsets,
194                    parts.len,
195                    parts.offset,
196                )?
197            }
198        };
199
200        match patches {
201            None => Ok(array.with_stats_set(packed_stats).into_array()),
202            Some(p) => Ok(
203                Patched::from_array_and_patches(array.into_array(), &p, exec_ctx)?
204                    .with_stats_set(packed_stats)
205                    .into_array(),
206            ),
207        }
208    }
209}
210
211/// Returns `None` when the best width is the native width.
212fn compress_global(
213    primitive_array: &PrimitiveArray,
214    exec_ctx: &mut ExecutionCtx,
215) -> VortexResult<Option<BitPackedArray>> {
216    let histogram = bit_width_histogram(primitive_array.as_view(), exec_ctx)?;
217    let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?;
218    if usize::from(bw) == primitive_array.ptype().bit_width() {
219        return Ok(None);
220    }
221    bitpack_encode(primitive_array, bw, Some(&histogram), exec_ctx).map(Some)
222}
223
224/// Uses global packing for uniform widths. Returns `None` for empty input or all-native widths.
225fn compress_blocked(
226    primitive_array: &PrimitiveArray,
227    exec_ctx: &mut ExecutionCtx,
228) -> VortexResult<Option<BitPackedArray>> {
229    let (bit_widths, num_exceptions) = find_best_bit_widths(primitive_array, exec_ctx)?;
230    if bit_widths.is_empty() {
231        return Ok(None);
232    }
233
234    match bit_widths.iter().copied().all_equal_value() {
235        Ok(bw) if usize::from(bw) == primitive_array.ptype().bit_width() => Ok(None),
236        Ok(bw) => bitpack_encode(primitive_array, bw, None, exec_ctx).map(Some),
237        Err(_) => {
238            bitpack_encode_blocked(primitive_array, &bit_widths, Some(num_exceptions), exec_ctx)
239                .map(Some)
240        }
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use rstest::rstest;
247    use vortex_array::ArrayId;
248    use vortex_fastlanes::bitpacked_v1_id;
249    use vortex_fastlanes::bitpacked_v2_id;
250
251    use super::BITPACKING_V1;
252    use super::BITPACKING_V2;
253    use crate::Scheme;
254    use crate::SchemeExt;
255
256    /// Both variants refine to v2 exactly when the v2 ID is allowed.
257    #[rstest]
258    #[case::neither(false, false, false)]
259    #[case::v1(true, false, false)]
260    #[case::v2_only(false, true, true)]
261    #[case::both(true, true, true)]
262    fn refine_picks_v2_when_v2_id_is_allowed(
263        #[case] allow_v1: bool,
264        #[case] allow_v2: bool,
265        #[case] expect_v2: bool,
266        #[values(&BITPACKING_V1, &BITPACKING_V2)] scheme: &'static dyn Scheme,
267    ) {
268        let allowed = |id: &ArrayId| {
269            (allow_v1 && *id == bitpacked_v1_id()) || (allow_v2 && *id == bitpacked_v2_id())
270        };
271        let refined = scheme.refine(&allowed);
272        assert_eq!(refined.id(), scheme.id());
273        let expected: &dyn Scheme = if expect_v2 {
274            &BITPACKING_V2
275        } else {
276            &BITPACKING_V1
277        };
278        assert_eq!(refined.produced_encodings(), expected.produced_encodings());
279    }
280}
vortex_btrblocks/schemes/integer/bitpacking.rs

vortex_btrblocks/schemes/integer/
bitpacking.rs
