|
| 1 | +// SPDX-License-Identifier: Apache-2.0 |
| 2 | +// SPDX-FileCopyrightText: Copyright the Vortex contributors |
| 3 | + |
| 4 | +//! BitPacking with a bit width per FastLanes chunk. |
| 5 | +
|
| 6 | +use vortex_array::ArrayId; |
| 7 | +use vortex_array::ArrayRef; |
| 8 | +use vortex_array::Canonical; |
| 9 | +use vortex_array::ExecutionCtx; |
| 10 | +use vortex_array::IntoArray; |
| 11 | +use vortex_array::VTable; |
| 12 | +use vortex_compressor::scheme::CompressionEstimate; |
| 13 | +use vortex_compressor::scheme::DeferredEstimate; |
| 14 | +use vortex_compressor::scheme::EstimateVerdict; |
| 15 | +use vortex_error::VortexResult; |
| 16 | +use vortex_fastlanes::BitPacked; |
| 17 | +use vortex_fastlanes::BitPackedArrayExt; |
| 18 | +use vortex_fastlanes::bitpack_compress::bitpack_to_best_chunk_widths; |
| 19 | +use vortex_fastlanes::bitpacked_v2_id; |
| 20 | + |
| 21 | +use crate::ArrayAndStats; |
| 22 | +use crate::CascadingCompressor; |
| 23 | +use crate::CompressorContext; |
| 24 | +use crate::Scheme; |
| 25 | +use crate::compress_patches; |
| 26 | + |
| 27 | +/// BitPacking for non-negative integers with a width chosen per 1024-element chunk. |
| 28 | +/// |
| 29 | +/// The output is an ordinary [`BitPacked`] array, but one whose chunks may differ in width, which |
| 30 | +/// serializes as `fastlanes.bitpacked_v2`. The writer therefore keeps this scheme only when an |
| 31 | +/// enabled edition permits that format; while it is present, [`super::BitPackingScheme`] defers |
| 32 | +/// to it so bit-packed residuals get per-chunk widths too. |
| 33 | +#[derive(Debug, Copy, Clone, PartialEq, Eq)] |
| 34 | +pub struct BitPackingV2Scheme; |
| 35 | + |
| 36 | +impl Scheme for BitPackingV2Scheme { |
| 37 | + fn scheme_name(&self) -> &'static str { |
| 38 | + "vortex.int.bitpacking_v2" |
| 39 | + } |
| 40 | + |
| 41 | + fn matches(&self, canonical: &Canonical) -> bool { |
| 42 | + canonical.dtype().is_int() |
| 43 | + } |
| 44 | + |
| 45 | + fn produced_encodings(&self) -> Vec<ArrayId> { |
| 46 | + vec![BitPacked.id()] |
| 47 | + } |
| 48 | + |
| 49 | + fn produced_serialized_ids(&self) -> Vec<ArrayId> { |
| 50 | + vec![bitpacked_v2_id()] |
| 51 | + } |
| 52 | + |
| 53 | + fn expected_compression_ratio( |
| 54 | + &self, |
| 55 | + data: &ArrayAndStats, |
| 56 | + _compress_ctx: CompressorContext, |
| 57 | + exec_ctx: &mut ExecutionCtx, |
| 58 | + ) -> CompressionEstimate { |
| 59 | + let stats = data.integer_stats(exec_ctx); |
| 60 | + |
| 61 | + // BitPacking only works for non-negative values. |
| 62 | + if stats.erased().min_is_negative() { |
| 63 | + return CompressionEstimate::Verdict(EstimateVerdict::Skip); |
| 64 | + } |
| 65 | + |
| 66 | + CompressionEstimate::Deferred(DeferredEstimate::Sample) |
| 67 | + } |
| 68 | + |
| 69 | + fn compress( |
| 70 | + &self, |
| 71 | + _compressor: &CascadingCompressor, |
| 72 | + data: &ArrayAndStats, |
| 73 | + _compress_ctx: CompressorContext, |
| 74 | + exec_ctx: &mut ExecutionCtx, |
| 75 | + ) -> VortexResult<ArrayRef> { |
| 76 | + let primitive_array = data.array_as_primitive().into_owned(); |
| 77 | + let packed = bitpack_to_best_chunk_widths(&primitive_array, exec_ctx)?; |
| 78 | + |
| 79 | + // If every chunk needs the full bit-width, return the original array. |
| 80 | + if packed.chunk_widths().uniform_width().map(usize::from) |
| 81 | + == Some(primitive_array.ptype().bit_width()) |
| 82 | + { |
| 83 | + return Ok(primitive_array.into_array()); |
| 84 | + } |
| 85 | + |
| 86 | + let Some(patches) = packed.patches() else { |
| 87 | + return Ok(packed.into_array()); |
| 88 | + }; |
| 89 | + |
| 90 | + // Compress the patches and place them back into the array. |
| 91 | + let packed_stats = packed.statistics().to_owned(); |
| 92 | + let ptype = packed.dtype().as_ptype(); |
| 93 | + let mut parts = BitPacked::into_parts(packed); |
| 94 | + parts.patches = Some(compress_patches(patches, exec_ctx)?); |
| 95 | + Ok(BitPacked::try_new( |
| 96 | + parts.packed, |
| 97 | + ptype, |
| 98 | + parts.validity, |
| 99 | + parts.patches, |
| 100 | + parts.widths, |
| 101 | + parts.len, |
| 102 | + parts.offset, |
| 103 | + )? |
| 104 | + .with_stats_set(packed_stats) |
| 105 | + .into_array()) |
| 106 | + } |
| 107 | +} |
0 commit comments