vortex_fastlanes/bitpacking/array/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::mem::MaybeUninit;
7
8use fastlanes::BitPacking;
9use vortex_array::ArrayRef;
10use vortex_array::ExecutionCtx;
11use vortex_array::TypedArrayRef;
12use vortex_array::array_slots;
13use vortex_array::arrays::Primitive;
14use vortex_array::arrays::PrimitiveArray;
15use vortex_array::buffer::BufferHandle;
16use vortex_array::dtype::DType;
17use vortex_array::dtype::NativePType;
18use vortex_array::dtype::PType;
19use vortex_array::patches::PatchSlotIndices;
20use vortex_array::patches::Patches;
21use vortex_array::patches::PatchesData;
22use vortex_array::validity::Validity;
23use vortex_array::vtable::child_to_validity;
24use vortex_error::VortexResult;
25use vortex_error::vortex_ensure;
26use vortex_error::vortex_err;
27
28pub mod bitpack_compress;
29pub mod bitpack_decompress;
30pub mod unpack_iter;
31
32use crate::BitPackedArray;
33use crate::FL_CHUNK_SIZE;
34use crate::bitpack_compress::bitpack_encode;
35use crate::unpack_iter::BitPacked as BitPackedIter;
36use crate::unpack_iter::BitUnpackedChunks;
37
38#[array_slots(crate::BitPacked)]
39pub struct BitPackedSlots {
40    /// The indices of exception values that don't fit in the bit-packed representation.
41    #[slot(0)]
42    pub patch_indices: Option<ArrayRef>,
43    /// The exception values that don't fit in the bit-packed representation.
44    #[slot(1)]
45    pub patch_values: Option<ArrayRef>,
46    /// Chunk offsets for the patch indices/values.
47    #[slot(2)]
48    pub patch_chunk_offsets: Option<ArrayRef>,
49    /// The validity bitmap indicating which elements are non-null.
50    #[slot(3)]
51    pub validity_child: Option<ArrayRef>,
52}
53
54pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices {
55    indices: BitPackedSlots::PATCH_INDICES,
56    values: BitPackedSlots::PATCH_VALUES,
57    chunk_offsets: BitPackedSlots::PATCH_CHUNK_OFFSETS,
58};
59
60pub struct BitPackedDataParts {
61    pub offset: u16,
62    pub bit_width: u8,
63    pub len: usize,
64    pub packed: BufferHandle,
65    pub patches: Option<Patches>,
66    pub validity: Validity,
67}
68
69#[derive(Clone, Debug)]
70pub struct BitPackedData {
71    /// The offset within the first block (created with a slice).
72    /// 0 <= offset < 1024
73    pub(super) offset: u16,
74    pub(super) bit_width: u8,
75    pub(super) packed: BufferHandle,
76    /// Patch metadata for reconstructing Patches from slots.
77    pub(super) patches_data: Option<PatchesData>,
78}
79
80impl Display for BitPackedData {
81    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
82        write!(f, "bit_width: {}, offset: {}", self.bit_width, self.offset)
83    }
84}
85
86impl BitPackedData {
87    /// Create a new bitpacked array using a buffer of packed data.
88    ///
89    /// The packed data should be interpreted as a sequence of values with size `bit_width`.
90    ///
91    /// # Errors
92    ///
93    /// This method returns errors if any of the metadata is inconsistent, for example the packed
94    /// buffer provided does not have the right size according to the supplied length and target
95    /// PType.
96    ///
97    /// # Safety
98    ///
99    /// For signed arrays, it is the caller's responsibility to ensure that there are no values
100    /// that can be interpreted once unpacked to the provided PType.
101    ///
102    /// This invariant is upheld by the compressor, but callers must ensure this if they wish to
103    /// construct a new `BitPackedArray` from parts.
104    ///
105    /// See also the [`encode`][Self::encode] method on this type for a safe path to create a new
106    /// bit-packed array.
107    /// A safe constructor for a `BitPackedArray` from its components:
108    ///
109    /// * `packed` is ByteBuffer holding the compressed data that was packed with FastLanes
110    ///   bit-packing to a `bit_width` bits per value. `length` is the length of the original
111    ///   vector. Note that the packed is padded with zeros to the next multiple of 1024 elements
112    ///   if `length` is not divisible by 1024.
113    /// * `ptype` of the original data
114    /// * `validity` to track any nulls
115    /// * `patches` optionally provided for values that did not pack
116    ///
117    /// Any failure in validation will result in an error.
118    ///
119    /// # Validation
120    ///
121    /// * The `ptype` must be an integer
122    /// * `validity` must have `length` len
123    /// * Any patches must have any `array_len` equal to `length`
124    /// * The `packed` buffer must be exactly sized to hold `length` values of `bit_width` rounded
125    ///   up to the next multiple of 1024.
126    ///
127    /// Any violation of these preconditions will result in an error.
128    pub fn try_new(
129        packed: BufferHandle,
130        patches: Option<Patches>,
131        bit_width: u8,
132        offset: u16,
133    ) -> VortexResult<Self> {
134        vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}");
135        vortex_ensure!(
136            offset < 1024,
137            "Offset must be less than the full block i.e., 1024, got {offset}"
138        );
139
140        Ok(Self {
141            offset,
142            bit_width,
143            packed,
144            patches_data: patches.as_ref().map(PatchesData::from_patches),
145        })
146    }
147
148    pub(crate) fn validate(
149        packed: &BufferHandle,
150        ptype: PType,
151        validity: &Validity,
152        patches: Option<&Patches>,
153        bit_width: u8,
154        length: usize,
155        offset: u16,
156    ) -> VortexResult<()> {
157        vortex_ensure!(ptype.is_int(), MismatchedTypes: "integer", ptype);
158        vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}");
159
160        if let Some(validity_len) = validity.maybe_len() {
161            vortex_ensure!(
162                validity_len == length,
163                "BitPackedArray validity length {validity_len} != array length {length}",
164            );
165        }
166
167        // Validate patches
168        if let Some(patches) = patches {
169            Self::validate_patches(patches, ptype, length)?;
170        }
171
172        // Validate packed buffer
173        let expected_packed_len =
174            (length + offset as usize).div_ceil(1024) * (128 * bit_width as usize);
175        vortex_ensure!(
176            packed.len() == expected_packed_len,
177            "Expected {} packed bytes, got {}",
178            expected_packed_len,
179            packed.len()
180        );
181
182        Ok(())
183    }
184
185    fn validate_patches(patches: &Patches, ptype: PType, len: usize) -> VortexResult<()> {
186        // Ensure that array and patches have same ptype
187        vortex_ensure!(
188            patches.dtype().eq_ignore_nullability(ptype.into()),
189            "Patches DType {} does not match BitPackedArray dtype {}",
190            patches.dtype().as_nonnullable(),
191            ptype
192        );
193
194        vortex_ensure!(
195            patches.array_len() == len,
196            "BitPackedArray patches length {} != expected {len}",
197            patches.array_len(),
198        );
199
200        Ok(())
201    }
202
203    pub fn ptype(&self, dtype: &DType) -> PType {
204        dtype.as_ptype()
205    }
206
207    /// Underlying bit packed values as byte array
208    #[inline]
209    pub fn packed(&self) -> &BufferHandle {
210        &self.packed
211    }
212
213    /// Access the slice of packed values as an array of `T`
214    #[inline]
215    pub fn packed_slice<T: NativePType + BitPacking>(&self) -> &[T] {
216        let packed_bytes = self.packed().as_host();
217        let packed_ptr: *const T = packed_bytes.as_ptr().cast();
218        // Return number of elements of type `T` packed in the buffer
219        let packed_len = packed_bytes.len() / size_of::<T>();
220
221        // SAFETY: as_slice points to buffer memory that outlives the lifetime of `self`.
222        //  Unfortunately Rust cannot understand this, so we reconstruct the slice from raw parts
223        //  to get it to reinterpret the lifetime.
224        unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) }
225    }
226
227    /// Accessor for bit unpacked chunks
228    pub fn unpacked_chunks<'a, T: BitPackedIter>(
229        &'a self,
230        dtype: &DType,
231        len: usize,
232        scratch: &'a mut [MaybeUninit<T>; FL_CHUNK_SIZE],
233    ) -> VortexResult<BitUnpackedChunks<'a, T>> {
234        assert_eq!(
235            T::PTYPE,
236            self.ptype(dtype),
237            "Requested type doesn't match the array ptype"
238        );
239        BitUnpackedChunks::try_new(self, len, scratch)
240    }
241
242    /// Bit-width of the packed values
243    #[inline]
244    pub fn bit_width(&self) -> u8 {
245        self.bit_width
246    }
247
248    #[inline]
249    pub fn offset(&self) -> u16 {
250        self.offset
251    }
252
253    /// Bit-pack an array of primitive integers down to the target bit-width using the FastLanes
254    /// SIMD-accelerated packing kernels.
255    ///
256    /// # Errors
257    ///
258    /// If the provided array is not an integer type, an error will be returned.
259    ///
260    /// If the provided array contains negative values, an error will be returned.
261    ///
262    /// If the requested bit-width for packing is larger than the array's native width, an
263    /// error will be returned.
264    pub fn encode(
265        array: &ArrayRef,
266        bit_width: u8,
267        ctx: &mut ExecutionCtx,
268    ) -> VortexResult<BitPackedArray> {
269        let parray: PrimitiveArray = array
270            .clone()
271            .try_downcast::<Primitive>()
272            .map_err(|a| vortex_err!(InvalidArgument: "Bitpacking can only encode primitive arrays, got {}", a.encoding_id()))?;
273        bitpack_encode(&parray, bit_width, None, ctx)
274    }
275
276    /// Calculate the maximum value that **can** be contained by this array, given its bit-width.
277    ///
278    /// Note that this value need not actually be present in the array.
279    #[inline]
280    pub fn max_packed_value(&self) -> usize {
281        (1 << self.bit_width()) - 1
282    }
283}
284
285pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
286    #[inline]
287    fn packed(&self) -> &BufferHandle {
288        BitPackedData::packed(self)
289    }
290
291    #[inline]
292    fn bit_width(&self) -> u8 {
293        BitPackedData::bit_width(self)
294    }
295
296    #[inline]
297    fn offset(&self) -> u16 {
298        BitPackedData::offset(self)
299    }
300
301    #[inline]
302    fn patches(&self) -> Option<Patches> {
303        PatchesData::patches_from_slots(
304            self.patches_data.as_ref(),
305            self.as_ref().len(),
306            self.as_ref().slots(),
307            PATCH_SLOTS,
308        )
309    }
310
311    #[inline]
312    fn validity(&self) -> Validity {
313        child_to_validity(self.validity_child(), self.as_ref().dtype().nullability())
314    }
315
316    #[inline]
317    fn packed_slice<T: NativePType + BitPacking>(&self) -> &[T] {
318        BitPackedData::packed_slice::<T>(self)
319    }
320
321    #[inline]
322    fn unpacked_chunks<'a, T: BitPackedIter>(
323        &'a self,
324        scratch: &'a mut [MaybeUninit<T>; FL_CHUNK_SIZE],
325    ) -> VortexResult<BitUnpackedChunks<'a, T>> {
326        BitPackedData::unpacked_chunks::<T>(
327            self,
328            self.as_ref().dtype(),
329            self.as_ref().len(),
330            scratch,
331        )
332    }
333}
334
335impl<T: TypedArrayRef<crate::BitPacked>> BitPackedArrayExt for T {}
336
337#[cfg(test)]
338mod test {
339    use std::sync::LazyLock;
340
341    use vortex_array::IntoArray;
342    use vortex_array::VortexSessionExecute;
343    use vortex_array::arrays::PrimitiveArray;
344    use vortex_array::assert_arrays_eq;
345    use vortex_buffer::Buffer;
346    use vortex_session::VortexSession;
347
348    use crate::BitPackedData;
349    use crate::bitpacking::array::BitPackedArrayExt;
350
351    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
352        let session = vortex_array::array_session();
353        crate::initialize(&session);
354        session
355    });
356
357    #[test]
358    fn test_encode() {
359        let mut ctx = SESSION.create_execution_ctx();
360        let values = [
361            Some(1u64),
362            None,
363            Some(1),
364            None,
365            Some(1),
366            None,
367            Some(u64::MAX),
368        ];
369        let uncompressed = PrimitiveArray::from_option_iter(values);
370        let packed = BitPackedData::encode(&uncompressed.into_array(), 1, &mut ctx).unwrap();
371        let expected = PrimitiveArray::from_option_iter(values);
372        let packed_primitive = packed
373            .as_array()
374            .clone()
375            .execute::<PrimitiveArray>(&mut ctx)
376            .unwrap();
377        assert_arrays_eq!(packed_primitive, expected, &mut ctx);
378    }
379
380    #[test]
381    fn test_encode_too_wide() {
382        let mut ctx = SESSION.create_execution_ctx();
383        let values = [Some(1u8), None, Some(1), None, Some(1), None];
384        let uncompressed = PrimitiveArray::from_option_iter(values);
385        let _packed = BitPackedData::encode(&uncompressed.clone().into_array(), 8, &mut ctx)
386            .expect_err("Cannot pack value into the same width");
387        let _packed = BitPackedData::encode(&uncompressed.into_array(), 9, &mut ctx)
388            .expect_err("Cannot pack value into larger width");
389    }
390
391    #[test]
392    fn signed_with_patches() {
393        let mut ctx = SESSION.create_execution_ctx();
394        let values: Buffer<i32> = (0i32..=512).collect();
395        let parray = values.clone().into_array();
396
397        let packed_with_patches = BitPackedData::encode(&parray, 9, &mut ctx).unwrap();
398        assert!(packed_with_patches.patches().is_some());
399        let packed_primitive = packed_with_patches
400            .as_array()
401            .clone()
402            .execute::<PrimitiveArray>(&mut ctx)
403            .unwrap();
404        assert_arrays_eq!(
405            packed_primitive,
406            PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable),
407            &mut ctx
408        );
409    }
410}
vortex_fastlanes/bitpacking/array/mod.rs

vortex_fastlanes/bitpacking/array/
mod.rs
