feat(btrblocks): pick per-chunk bit widths when fastlanes.bitpacked_v2 is allowed

vortex-btrblocks/src/builder.rs

@@ -91,12 +91,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
 #[derive(Debug, Clone)]
 pub struct BtrBlocksCompressorBuilder {
     schemes: Vec<&'static dyn Scheme>,
+    allowed_serialized_ids: Option<HashSet<ArrayId>>,
 }
 
 impl Default for BtrBlocksCompressorBuilder {
     fn default() -> Self {
         Self {
             schemes: ALL_SCHEMES.to_vec(),
+            allowed_serialized_ids: None,
         }
     }
 }
@@ -108,6 +110,7 @@ impl BtrBlocksCompressorBuilder {
     pub fn empty() -> Self {
         Self {
             schemes: Vec::new(),
+            allowed_serialized_ids: None,
         }
     }
 
@@ -188,6 +191,15 @@ impl BtrBlocksCompressorBuilder {
         #[cfg(feature = "pco")]
         excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]);
         let builder = self.exclude_schemes(excluded);
+        // CUDA kernels decode each encoding's original wire format, whose ID is the encoding's
+        // own. In particular bit-unpacking has no per-chunk width kernel yet, so allowing only
+        // the original formats keeps one width per array.
+        let original_formats = builder
+            .schemes
+            .iter()
+            .flat_map(|scheme| scheme.produced_encodings())
+            .collect();
+        let builder = builder.allow_serialized_ids(&original_formats);
 
         #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
         let builder = builder.with_new_scheme(&binary::ZstdBuffersScheme);
@@ -214,16 +226,35 @@ impl BtrBlocksCompressorBuilder {
         self
     }
 
+    /// Restricts the output to the serialized IDs in `allowed`, intersecting with any earlier
+    /// restriction, so a scheme whose encoding has several wire formats produces the newest one
+    /// still allowed.
+    ///
+    /// The file writer passes the serialized IDs of its configured editions.
+    pub fn allow_serialized_ids(mut self, allowed: &HashSet<ArrayId>) -> Self {
+        self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() {
+            Some(existing) => existing.intersection(allowed).copied().collect(),
+            None => allowed.clone(),
+        });
+        self
+    }
+
     /// Builds the configured [`BtrBlocksCompressor`].
     pub fn build(self) -> BtrBlocksCompressor {
-        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
+        let compressor = CascadingCompressor::new(self.schemes);
+        BtrBlocksCompressor(match self.allowed_serialized_ids {
+            Some(allowed) => compressor.with_allowed_serialized_ids(allowed),
+            None => compressor,
+        })
     }
 }
 
 #[cfg(test)]
 mod tests {
     use vortex_array::VTable;
+    use vortex_fastlanes::BitPacked;
     use vortex_fastlanes::FoR;
+    use vortex_fastlanes::bitpacked_v2_id;
 
     use super::*;
 
@@ -287,6 +318,20 @@ mod tests {
         }
     }
 
+    /// Every serialized ID is allowed by default, so bit-packing picks per-chunk widths. The CUDA
+    /// preset has no per-chunk kernel and allows only the original formats.
+    #[test]
+    fn cuda_compatible_disallows_per_chunk_bitpacking() {
+        let default = BtrBlocksCompressorBuilder::default().build();
+        assert!(default.0.allows_serialized_id(bitpacked_v2_id()));
+
+        let cuda = BtrBlocksCompressorBuilder::default()
+            .only_cuda_compatible()
+            .build();
+        assert!(!cuda.0.allows_serialized_id(bitpacked_v2_id()));
+        assert!(cuda.0.allows_serialized_id(BitPacked.id()));
+    }
+
     #[test]
     fn cuda_compatible_uses_fsst_for_strings() {
         let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();

vortex-btrblocks/src/schemes/integer/bitpacking.rs

@@ -17,17 +17,28 @@ use vortex_compressor::scheme::DeferredEstimate;
 use vortex_compressor::scheme::EstimateVerdict;
 use vortex_error::VortexResult;
 use vortex_fastlanes::BitPacked;
+use vortex_fastlanes::BitPackedArray;
+use vortex_fastlanes::BitPackedArrayExt;
+use vortex_fastlanes::BitPackedArraySlotsExt;
+use vortex_fastlanes::BitPackedSlots;
 use vortex_fastlanes::bitpack_compress::bit_width_histogram;
 use vortex_fastlanes::bitpack_compress::bitpack_encode;
+use vortex_fastlanes::bitpack_compress::bitpack_to_best_chunk_widths;
 use vortex_fastlanes::bitpack_compress::find_best_bit_width;
+use vortex_fastlanes::bitpacked_v2_id;
 
 use crate::ArrayAndStats;
 use crate::CascadingCompressor;
 use crate::CompressorContext;
 use crate::Scheme;
+use crate::SchemeExt;
 use crate::compress_patches;
 
 /// BitPacking encoding for non-negative integers.
+///
+/// Every 1024-element chunk gets its own bit width when the compressor may use the
+/// `fastlanes.bitpacked_v2` format. Otherwise every chunk shares one width, which is the original
+/// `fastlanes.bitpacked` format.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub struct BitPackingScheme;
 
@@ -66,31 +77,51 @@ impl Scheme for BitPackingScheme {
 
     fn compress(
         &self,
-        _compressor: &CascadingCompressor,
+        compressor: &CascadingCompressor,
         data: &ArrayAndStats,
-        _compress_ctx: CompressorContext,
+        compress_ctx: CompressorContext,
         exec_ctx: &mut ExecutionCtx,
     ) -> VortexResult<ArrayRef> {
         let primitive_array = data.array_as_primitive();
+        let full_width = primitive_array.ptype().bit_width();
 
-        let histogram = bit_width_histogram(primitive_array, exec_ctx)?;
-        let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?;
+        // Per-chunk widths are the newer format of BitPacked. Produce them whenever the writer
+        // may serialize them; otherwise every chunk shares one width, the original format.
+        let packed = if compressor.allows_serialized_id(bitpacked_v2_id()) {
+            bitpack_to_best_chunk_widths(&data.array_as_primitive().into_owned(), exec_ctx)?
+        } else {
+            let histogram = bit_width_histogram(primitive_array, exec_ctx)?;
+            let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?;
+            if bw as usize == full_width {
+                return Ok(primitive_array.array().clone());
+            }
+            bitpack_encode(
+                &data.array_as_primitive().into_owned(),
+                bw,
+                Some(&histogram),
+                exec_ctx,
+            )?
+        };
 
-        // If best bw is determined to be the current bit-width, return the original array.
-        if bw as usize == primitive_array.ptype().bit_width() {
+        // If every chunk needs the full bit-width, return the original array.
+        if packed.chunk_widths().uniform_width().map(usize::from) == Some(full_width) {
+            return Ok(primitive_array.array().clone());
+        }
+        // Mostly patches means the values were never really packed: the array is sparse in all
+        // but name, and beats raw storage only by the bytes of the few values that did fit.
+        if packed
+            .patches()
+            .is_some_and(|p| p.num_patches() * 2 >= packed.len())
+        {
             return Ok(primitive_array.array().clone());
         }
-
-        // Otherwise we can bitpack the array.
-        let primitive_array = primitive_array.into_owned();
-        let packed = bitpack_encode(&primitive_array, bw, Some(&histogram), exec_ctx)?;
 
         let packed_stats = packed.statistics().to_owned();
         let ptype = packed.dtype().as_ptype();
         let mut parts = BitPacked::into_parts(packed);
+        let patches = parts.patches.take();
 
-        let array = if use_experimental_patches() {
-            let patches = parts.patches.take();
+        if use_experimental_patches() {
             // Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray.
             let array = BitPacked::try_new(
                 parts.packed,
@@ -100,36 +131,53 @@ impl Scheme for BitPackingScheme {
                 parts.widths,
                 parts.len,
                 parts.offset,
-            )?
-            .into_array();
-
-            match patches {
+            )?;
+            let array =
+                compress_width_table(compressor, array, &compress_ctx, exec_ctx)?.into_array();
+            return Ok(match patches {
                 None => array,
                 Some(p) => Patched::from_array_and_patches(array, &p, exec_ctx)?
                     .with_stats_set(packed_stats)
                     .into_array(),
-            }
-        } else {
-            // Compress patches and place back into BitPackedArray.
-            let patches = parts
-                .patches
-                .take()
-                .map(|p| compress_patches(p, exec_ctx))
-                .transpose()?;
-            parts.patches = patches;
-            BitPacked::try_new(
-                parts.packed,
-                ptype,
-                parts.validity,
-                parts.patches,
-                parts.widths,
-                parts.len,
-                parts.offset,
-            )?
-            .with_stats_set(packed_stats)
-            .into_array()
-        };
+            });
+        }
 
-        Ok(array)
+        // Compress patches and place back into BitPackedArray.
+        let patches = patches.map(|p| compress_patches(p, exec_ctx)).transpose()?;
+        let array = BitPacked::try_new(
+            parts.packed,
+            ptype,
+            parts.validity,
+            patches,
+            parts.widths,
+            parts.len,
+            parts.offset,
+        )?;
+        Ok(
+            compress_width_table(compressor, array, &compress_ctx, exec_ctx)?
+                .with_stats_set(packed_stats)
+                .into_array(),
+        )
     }
 }
+
+/// Re-encode the width table child through the compressor. Arrays whose chunks share one width
+/// have no table.
+fn compress_width_table(
+    compressor: &CascadingCompressor,
+    packed: BitPackedArray,
+    compress_ctx: &CompressorContext,
+    exec_ctx: &mut ExecutionCtx,
+) -> VortexResult<BitPackedArray> {
+    let Some(table) = packed.width_table().cloned() else {
+        return Ok(packed);
+    };
+    let compressed = compressor.compress_child(
+        &table,
+        compress_ctx,
+        BitPackingScheme.id(),
+        BitPackedSlots::WIDTH_TABLE,
+        exec_ctx,
+    )?;
+    BitPacked::with_width_table(packed, compressed)
+}

vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs

@@ -13,6 +13,7 @@ use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::Constant;
 use vortex_array::arrays::Dict;
+use vortex_array::arrays::Primitive;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::expr::stats::Precision;
 use vortex_array::expr::stats::Stat;
@@ -273,3 +274,19 @@ fn has_nested_delta(array: &vortex_array::ArrayRef, under_delta: bool) -> bool {
         .iter()
         .any(|child| has_nested_delta(child, under_delta || is_delta))
 }
+
+/// A few wide values pack at width 0 with nearly every value patched, which beats raw storage
+/// only by the bytes of the values that did fit. Such an array stays primitive.
+#[test]
+fn test_mostly_patched_stays_primitive() -> VortexResult<()> {
+    let values: Vec<i64> = vec![0, 1 << 40, 1 << 41, 1 << 42, 1 << 43, 1 << 44];
+    let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable);
+    let btr = BtrBlocksCompressor::default();
+    let compressed = btr.compress(&array.into_array(), &mut SESSION.create_execution_ctx())?;
+    assert!(
+        compressed.is::<Primitive>(),
+        "got {}",
+        compressed.encoding_id()
+    );
+    Ok(())
+}

vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap

@@ -3,7 +3,9 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: u64, len=16384, nbytes=131072
-root: fastlanes.for(u64, len=16384) nbytes=49152
+root: fastlanes.for(u64, len=16384) nbytes=47248
   metadata: reference: 1700000001036u64
-  encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49152
-    metadata: bit_width: 24, offset: 0
+  encoded: fastlanes.bitpacked(u64, len=16384) nbytes=47248
+    metadata: bit_widths: 16 chunks, max 24, offset: 0
+    width_table: vortex.primitive(u8, len=16) nbytes=16
+      metadata: ptype: u8

vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap

@@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: list(i32), len=4066, nbytes=81804
-root: vortex.list(list(i32), len=4066) nbytes=11146
+root: vortex.list(list(i32), len=4066) nbytes=10944
   metadata: 
   elements: vortex.runend(i32, len=16384) nbytes=3968
     metadata: offset: 0
@@ -15,11 +15,13 @@ root: vortex.list(list(i32), len=4066) nbytes=11146
       metadata: reference: -49931i32
       encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176
         metadata: bit_width: 17, offset: 0
-  offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178
-    metadata: bit_width: 14, offset: 0
-    patch_indices: vortex.primitive(u16, len=1) nbytes=2
+  offsets: fastlanes.bitpacked(u16, len=4067) nbytes=6976
+    metadata: bit_widths: 4 chunks, max 14, offset: 0
+    patch_indices: vortex.primitive(u16, len=14) nbytes=28
+      metadata: ptype: u16
+    patch_values: vortex.primitive(u16, len=14) nbytes=28
       metadata: ptype: u16
-    patch_values: vortex.constant(u16, len=1) nbytes=4
-      metadata: scalar: 16384u16
     patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4
       metadata: ptype: u8
+    width_table: vortex.primitive(u8, len=4) nbytes=4
+      metadata: ptype: u8

vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap

@@ -3,13 +3,21 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: utf8, len=16384, nbytes=653785
-root: vortex.fsst(utf8, len=16384) nbytes=151382
+root: vortex.fsst(utf8, len=16384) nbytes=147198
   metadata: len: 16384, nsymbols: 223
   uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154
     metadata: fill_value: 24u8
     patch_indices: vortex.primitive(u16, len=1575) nbytes=3150
       metadata: ptype: u16
     patch_values: vortex.constant(u8, len=1575) nbytes=2
       metadata: scalar: 23u8
-  codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992
-    metadata: bit_width: 17, offset: 0
+  codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=32808
+    metadata: bit_widths: 17 chunks, max 17, offset: 0
+    patch_indices: vortex.primitive(u16, len=1) nbytes=2
+      metadata: ptype: u16
+    patch_values: vortex.constant(u32, len=1) nbytes=4
+      metadata: scalar: 109229u32
+    patch_chunk_offsets: vortex.primitive(u8, len=17) nbytes=17
+      metadata: ptype: u8
+    width_table: vortex.primitive(u8, len=17) nbytes=17
+      metadata: ptype: u8

vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap

@@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs
 expression: rendered
 ---
 input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072
-root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67584
+root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=65680
   metadata: 
-  storage: fastlanes.for(i64, len=16384) nbytes=67584
+  storage: fastlanes.for(i64, len=16384) nbytes=65680
     metadata: reference: 1700000000891673i64
-    encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67584
-      metadata: bit_width: 33, offset: 0
+    encoded: fastlanes.bitpacked(i64, len=16384) nbytes=65680
+      metadata: bit_widths: 16 chunks, max 33, offset: 0
+      width_table: vortex.primitive(u8, len=16) nbytes=16
+        metadata: ptype: u8