|
| 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::BitPackedArray; |
| 18 | +use vortex_fastlanes::BitPackedArrayExt; |
| 19 | +use vortex_fastlanes::BitPackedArraySlotsExt; |
| 20 | +use vortex_fastlanes::BitPackedSlots; |
| 21 | +use vortex_fastlanes::bitpack_compress::bitpack_to_best_chunk_widths; |
| 22 | +use vortex_fastlanes::bitpacked_v2_id; |
| 23 | + |
| 24 | +use crate::ArrayAndStats; |
| 25 | +use crate::CascadingCompressor; |
| 26 | +use crate::CompressorContext; |
| 27 | +use crate::Scheme; |
| 28 | +use crate::SchemeExt; |
| 29 | +use crate::compress_patches; |
| 30 | + |
| 31 | +/// BitPacking for non-negative integers with a width chosen per 1024-element chunk. |
| 32 | +/// |
| 33 | +/// The output is an ordinary [`BitPacked`] array, but one whose chunks may differ in width, which |
| 34 | +/// serializes as `fastlanes.bitpacked_v2`. The writer therefore keeps this scheme only when an |
| 35 | +/// enabled edition permits that format; while it is present, [`super::BitPackingScheme`] defers |
| 36 | +/// to it so bit-packed residuals get per-chunk widths too. |
| 37 | +#[derive(Debug, Copy, Clone, PartialEq, Eq)] |
| 38 | +pub struct BitPackingV2Scheme; |
| 39 | + |
| 40 | +impl Scheme for BitPackingV2Scheme { |
| 41 | + fn scheme_name(&self) -> &'static str { |
| 42 | + "vortex.int.bitpacking_v2" |
| 43 | + } |
| 44 | + |
| 45 | + fn matches(&self, canonical: &Canonical) -> bool { |
| 46 | + canonical.dtype().is_int() |
| 47 | + } |
| 48 | + |
| 49 | + fn produced_encodings(&self) -> Vec<ArrayId> { |
| 50 | + vec![BitPacked.id()] |
| 51 | + } |
| 52 | + |
| 53 | + fn produced_serialized_ids(&self) -> Vec<ArrayId> { |
| 54 | + vec![bitpacked_v2_id()] |
| 55 | + } |
| 56 | + |
| 57 | + fn expected_compression_ratio( |
| 58 | + &self, |
| 59 | + data: &ArrayAndStats, |
| 60 | + _compress_ctx: CompressorContext, |
| 61 | + exec_ctx: &mut ExecutionCtx, |
| 62 | + ) -> CompressionEstimate { |
| 63 | + let stats = data.integer_stats(exec_ctx); |
| 64 | + |
| 65 | + // BitPacking only works for non-negative values. |
| 66 | + if stats.erased().min_is_negative() { |
| 67 | + return CompressionEstimate::Verdict(EstimateVerdict::Skip); |
| 68 | + } |
| 69 | + |
| 70 | + CompressionEstimate::Deferred(DeferredEstimate::Sample) |
| 71 | + } |
| 72 | + |
| 73 | + fn compress( |
| 74 | + &self, |
| 75 | + compressor: &CascadingCompressor, |
| 76 | + data: &ArrayAndStats, |
| 77 | + compress_ctx: CompressorContext, |
| 78 | + exec_ctx: &mut ExecutionCtx, |
| 79 | + ) -> VortexResult<ArrayRef> { |
| 80 | + let primitive_array = data.array_as_primitive().into_owned(); |
| 81 | + let packed = bitpack_to_best_chunk_widths(&primitive_array, exec_ctx)?; |
| 82 | + |
| 83 | + // If every chunk needs the full bit-width, return the original array. |
| 84 | + if packed.chunk_widths().uniform_width().map(usize::from) |
| 85 | + == Some(primitive_array.ptype().bit_width()) |
| 86 | + { |
| 87 | + return Ok(primitive_array.into_array()); |
| 88 | + } |
| 89 | + |
| 90 | + let Some(patches) = packed.patches() else { |
| 91 | + return Ok( |
| 92 | + compress_width_table(compressor, packed, &compress_ctx, exec_ctx)?.into_array(), |
| 93 | + ); |
| 94 | + }; |
| 95 | + |
| 96 | + // Mostly patches means the values were never really packed: the array is sparse in all |
| 97 | + // but name, and beats raw storage only by the bytes of the few values that did fit. |
| 98 | + if patches.num_patches() * 2 >= primitive_array.len() { |
| 99 | + return Ok(primitive_array.into_array()); |
| 100 | + } |
| 101 | + |
| 102 | + // Compress the patches and place them back into the array. |
| 103 | + let packed_stats = packed.statistics().to_owned(); |
| 104 | + let ptype = packed.dtype().as_ptype(); |
| 105 | + let mut parts = BitPacked::into_parts(packed); |
| 106 | + parts.patches = Some(compress_patches(patches, exec_ctx)?); |
| 107 | + let packed = BitPacked::try_new( |
| 108 | + parts.packed, |
| 109 | + ptype, |
| 110 | + parts.validity, |
| 111 | + parts.patches, |
| 112 | + parts.widths, |
| 113 | + parts.len, |
| 114 | + parts.offset, |
| 115 | + )? |
| 116 | + .with_stats_set(packed_stats); |
| 117 | + Ok(compress_width_table(compressor, packed, &compress_ctx, exec_ctx)?.into_array()) |
| 118 | + } |
| 119 | +} |
| 120 | + |
| 121 | +/// Re-encode the width table child through the compressor. Arrays whose chunks share one width |
| 122 | +/// have no table. |
| 123 | +fn compress_width_table( |
| 124 | + compressor: &CascadingCompressor, |
| 125 | + packed: BitPackedArray, |
| 126 | + compress_ctx: &CompressorContext, |
| 127 | + exec_ctx: &mut ExecutionCtx, |
| 128 | +) -> VortexResult<BitPackedArray> { |
| 129 | + let Some(table) = packed.width_table().cloned() else { |
| 130 | + return Ok(packed); |
| 131 | + }; |
| 132 | + let compressed = compressor.compress_child( |
| 133 | + &table, |
| 134 | + compress_ctx, |
| 135 | + BitPackingV2Scheme.id(), |
| 136 | + BitPackedSlots::WIDTH_TABLE, |
| 137 | + exec_ctx, |
| 138 | + )?; |
| 139 | + BitPacked::with_width_table(packed, compressed) |
| 140 | +} |
0 commit comments