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 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::Patched;
13use vortex_array::arrays::patched::use_experimental_patches;
14use vortex_array::arrays::primitive::PrimitiveArrayExt;
15use vortex_compressor::scheme::CompressionEstimate;
16use vortex_compressor::scheme::DeferredEstimate;
17use vortex_compressor::scheme::EstimateVerdict;
18use vortex_error::VortexResult;
19use vortex_fastlanes::BitPacked;
20use vortex_fastlanes::bitpack_compress::bit_width_histogram;
21use vortex_fastlanes::bitpack_compress::bitpack_encode;
22use vortex_fastlanes::bitpack_compress::find_best_bit_width;
23
24use crate::ArrayAndStats;
25use crate::CascadingCompressor;
26use crate::CompressorContext;
27use crate::Scheme;
28use crate::compress_patches;
29
30/// BitPacking encoding for non-negative integers.
31#[derive(Debug, Copy, Clone, PartialEq, Eq)]
32pub struct BitPackingScheme;
33
34impl Scheme for BitPackingScheme {
35    fn scheme_name(&self) -> &'static str {
36        "vortex.int.bitpacking"
37    }
38
39    fn matches(&self, canonical: &Canonical) -> bool {
40        canonical.dtype().is_int()
41    }
42
43    fn produced_encodings(&self) -> Vec<ArrayId> {
44        let mut encodings = vec![BitPacked.id()];
45        if use_experimental_patches() {
46            encodings.push(Patched.id());
47        }
48        encodings
49    }
50
51    fn expected_compression_ratio(
52        &self,
53        data: &ArrayAndStats,
54        _compress_ctx: CompressorContext,
55        exec_ctx: &mut ExecutionCtx,
56    ) -> CompressionEstimate {
57        let stats = data.integer_stats(exec_ctx);
58
59        // BitPacking only works for non-negative values.
60        if stats.erased().min_is_negative() {
61            return CompressionEstimate::Verdict(EstimateVerdict::Skip);
62        }
63
64        CompressionEstimate::Deferred(DeferredEstimate::Sample)
65    }
66
67    fn compress(
68        &self,
69        _compressor: &CascadingCompressor,
70        data: &ArrayAndStats,
71        _compress_ctx: CompressorContext,
72        exec_ctx: &mut ExecutionCtx,
73    ) -> VortexResult<ArrayRef> {
74        let primitive_array = data.array_as_primitive();
75
76        let histogram = bit_width_histogram(primitive_array, exec_ctx)?;
77        let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?;
78
79        // If best bw is determined to be the current bit-width, return the original array.
80        if bw as usize == primitive_array.ptype().bit_width() {
81            return Ok(primitive_array.array().clone());
82        }
83
84        // Otherwise we can bitpack the array.
85        let primitive_array = primitive_array.into_owned();
86        let packed = bitpack_encode(&primitive_array, bw, Some(&histogram), exec_ctx)?;
87
88        let packed_stats = packed.statistics().to_owned();
89        let ptype = packed.dtype().as_ptype();
90        let mut parts = BitPacked::into_parts(packed);
91
92        let array = if use_experimental_patches() {
93            let patches = parts.patches.take();
94            // Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray.
95            let array = BitPacked::try_new(
96                parts.packed,
97                ptype,
98                parts.validity,
99                None,
100                parts.bit_width,
101                parts.len,
102                parts.offset,
103            )?
104            .into_array();
105
106            match patches {
107                None => array,
108                Some(p) => Patched::from_array_and_patches(array, &p, exec_ctx)?
109                    .with_stats_set(packed_stats)
110                    .into_array(),
111            }
112        } else {
113            // Compress patches and place back into BitPackedArray.
114            let patches = parts
115                .patches
116                .take()
117                .map(|p| compress_patches(p, exec_ctx))
118                .transpose()?;
119            parts.patches = patches;
120            BitPacked::try_new(
121                parts.packed,
122                ptype,
123                parts.validity,
124                parts.patches,
125                parts.bit_width,
126                parts.len,
127                parts.offset,
128            )?
129            .with_stats_set(packed_stats)
130            .into_array()
131        };
132
133        Ok(array)
134    }
135}
vortex_btrblocks/schemes/integer/bitpacking.rs

vortex_btrblocks/schemes/integer/
bitpacking.rs
