From 41b85cc1781428a6a79ce538bd28349a258d6678 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 2 Oct 2026 18:08:54 -0400 Subject: [PATCH 1/5] refine BitPackingScheme to per-block bit widths when fastlanes.bitpacked.v2 is allowed Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 7 +- .../src/schemes/integer/bitpacking.rs | 253 +++++++++++++++++- vortex-btrblocks/src/schemes/integer/for_.rs | 4 +- vortex-btrblocks/src/schemes/integer/mod.rs | 1 + vortex-btrblocks/src/session.rs | 2 +- vortex-btrblocks/tests/bitpacking_config.rs | 146 ++++++++++ vortex-btrblocks/tests/decimal_config.rs | 3 +- vortex-btrblocks/tests/for_config.rs | 2 +- 8 files changed, 406 insertions(+), 12 deletions(-) create mode 100644 vortex-btrblocks/tests/bitpacking_config.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 7bdc6661ff6..d3dea3fdb53 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -5,6 +5,7 @@ use vortex_array::ArrayId; use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_fastlanes::bitpacked_v2_id; use vortex_fastlanes::for_v2_id; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; @@ -94,9 +95,9 @@ impl CompressionMode { fn excluded_encodings(self) -> Vec { match self { Self::All | Self::Default | Self::Compact => Vec::new(), - // Multi-part DecimalByteParts arrays and FoR arrays with per-chunk references have no - // CUDA decode kernel. - Self::Cuda => vec![decimal_byte_parts_v2_id(), for_v2_id()], + // Multi-part DecimalByteParts arrays, FoR arrays with per-chunk references and + // BitPacked arrays with per-block bit widths have no CUDA decode kernel. + Self::Cuda => vec![decimal_byte_parts_v2_id(), for_v2_id(), bitpacked_v2_id()], } } } diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 4f81462ca3c..b0cfba95fbf 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -10,27 +10,84 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Patched; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::buffer::BufferHandle; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::patches::Patches; +use vortex_array::validity::Validity; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArray; +use vortex_fastlanes::BitWidths; use vortex_fastlanes::bitpack_compress::bit_width_histogram; use vortex_fastlanes::bitpack_compress::bitpack_encode; +use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_widths; use vortex_fastlanes::bitpack_compress::find_best_bit_width; use vortex_fastlanes::bitpacked_v1_id; +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; +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum BitPackingSchemeMode { + V1, + V2, +} + +/// The BitPacking scheme that only produces arrays with one bit width. +pub(crate) static BITPACKING_V1: BitPackingScheme = BitPackingScheme::v1(); + +/// The BitPacking scheme that may give each 1024-element block its own bit width. +pub(crate) static BITPACKING_V2: BitPackingScheme = BitPackingScheme::v2(); + /// BitPacking encoding for non-negative integers. +/// +/// The v1 mode packs every value at one bit width. The v2 mode chooses a width for each +/// 1024-element block. Arrays whose blocks all choose the same width keep one bit width and +/// serialize as `fastlanes.bitpacked`, while the others serialize as `fastlanes.bitpacked.v2`. +/// +/// The default uses v1. [`refine`](Scheme::refine) picks v2 when the v2 ID is allowed and v1 +/// otherwise. #[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct BitPackingScheme; +pub struct BitPackingScheme { + mode: BitPackingSchemeMode, +} + +impl BitPackingScheme { + /// Creates a BitPacking scheme configured for v1, which uses one bit width. + pub const fn v1() -> Self { + Self { + mode: BitPackingSchemeMode::V1, + } + } + + /// Creates a BitPacking scheme configured for v2, which may use one bit width per block. + pub const fn v2() -> Self { + Self { + mode: BitPackingSchemeMode::V2, + } + } +} + +impl Default for BitPackingScheme { + fn default() -> Self { + Self::v1() + } +} impl Scheme for BitPackingScheme { fn scheme_name(&self) -> &'static str { @@ -43,13 +100,32 @@ impl Scheme for BitPackingScheme { fn produced_encodings(&self) -> Vec { // Global-width arrays serialize under the frozen v1 ID. - let mut encodings = vec![bitpacked_v1_id()]; + let mut encodings = match self.mode { + BitPackingSchemeMode::V1 => vec![bitpacked_v1_id()], + BitPackingSchemeMode::V2 => vec![bitpacked_v1_id(), bitpacked_v2_id()], + }; if use_experimental_patches() { encodings.push(Patched.id()); } encodings } + fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme { + if allowed(&bitpacked_v2_id()) { + &BITPACKING_V2 + } else { + &BITPACKING_V1 + } + } + + /// Children: block offsets=0 in v2 mode. + fn num_children(&self) -> usize { + match self.mode { + BitPackingSchemeMode::V1 => 0, + BitPackingSchemeMode::V2 => 1, + } + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, @@ -68,11 +144,15 @@ impl Scheme for BitPackingScheme { fn compress( &self, - _compressor: &CascadingCompressor, + compressor: &CascadingCompressor, data: &ArrayAndStats, - _compress_ctx: CompressorContext, + compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { + if self.mode == BitPackingSchemeMode::V2 { + return self.compress_blocked(compressor, data, compress_ctx, exec_ctx); + } + let primitive_array = data.array_as_primitive(); let histogram = bit_width_histogram(primitive_array, exec_ctx)?; @@ -135,3 +215,168 @@ impl Scheme for BitPackingScheme { Ok(array) } } + +impl BitPackingScheme { + /// Bit-pack each 1024-element block at its own best width, keeping one bit width when every + /// block chooses the same one. + fn compress_blocked( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let primitive_array = data.array_as_primitive().into_owned(); + let packed = bitpack_to_best_bit_widths(&primitive_array, exec_ctx)?; + + let packed_stats = packed.statistics().to_owned(); + let ptype = packed.dtype().as_ptype(); + let mut parts = BitPacked::into_parts(packed); + let BitWidths::Blocked(block_offsets) = parts.bit_widths else { + vortex_bail!("Blocked bit-packing must produce block offsets"); + }; + + let bit_widths = match global_bit_width(&block_offsets, exec_ctx)? { + // Like v1, keep the original array when every block needs its full width. + Some(bit_width) if usize::from(bit_width) == ptype.bit_width() => { + return Ok(primitive_array.into_array()); + } + // The blocks are packed back to back, so equal widths are also one global width. + Some(bit_width) => BitWidths::Global(bit_width), + None => BitWidths::Blocked(compressor.compress_child( + &block_offsets, + &compress_ctx, + self.id(), + 0, + exec_ctx, + )?), + }; + + let array = if use_experimental_patches() { + let patches = parts.patches.take(); + // Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray. + let array = new_bitpacked( + parts.packed, + ptype, + parts.validity, + None, + bit_widths, + parts.len, + parts.offset, + )? + .into_array(); + + 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()?; + new_bitpacked( + parts.packed, + ptype, + parts.validity, + patches, + bit_widths, + parts.len, + parts.offset, + )? + .with_stats_set(packed_stats) + .into_array() + }; + + Ok(array) + } +} + +/// The bit width shared by every block between `block_offsets`, or `None` if the widths differ. +fn global_bit_width( + block_offsets: &ArrayRef, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult> { + let boundaries = block_offsets + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::(exec_ctx)?; + let mut bit_widths = boundaries + .as_slice::() + .windows(2) + .map(|pair| (pair[1] - pair[0]) / 128); + let Some(first) = bit_widths.next() else { + return Ok(Some(0)); + }; + if !bit_widths.all(|bit_width| bit_width == first) { + return Ok(None); + } + Ok(Some(u8::try_from(first)?)) +} + +/// Construct a bit-packed array at `bit_widths`. +fn new_bitpacked( + packed: BufferHandle, + ptype: PType, + validity: Validity, + patches: Option, + bit_widths: BitWidths, + len: usize, + offset: u16, +) -> VortexResult { + match bit_widths { + BitWidths::Global(bit_width) => { + BitPacked::try_new(packed, ptype, validity, patches, bit_width, len, offset) + } + BitWidths::Blocked(block_offsets) => BitPacked::try_new_with_block_offsets( + packed, + ptype, + validity, + patches, + block_offsets, + len, + offset, + ), + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayId; + use vortex_fastlanes::bitpacked_v1_id; + use vortex_fastlanes::bitpacked_v2_id; + + use super::BITPACKING_V1; + use super::BITPACKING_V2; + use crate::Scheme; + use crate::SchemeExt; + + /// Both variants refine to v2 exactly when the v2 ID is allowed. + #[rstest] + #[case::neither(false, false, false)] + #[case::v1(true, false, false)] + #[case::v2_only(false, true, true)] + #[case::both(true, true, true)] + fn refine_picks_v2_when_v2_id_is_allowed( + #[case] allow_v1: bool, + #[case] allow_v2: bool, + #[case] expect_v2: bool, + #[values(&BITPACKING_V1, &BITPACKING_V2)] scheme: &'static dyn Scheme, + ) { + let allowed = |id: &ArrayId| { + (allow_v1 && *id == bitpacked_v1_id()) || (allow_v2 && *id == bitpacked_v2_id()) + }; + let refined = scheme.refine(&allowed); + assert_eq!(refined.id(), scheme.id()); + let expected: &dyn Scheme = if expect_v2 { + &BITPACKING_V2 + } else { + &BITPACKING_V1 + }; + assert_eq!(refined.produced_encodings(), expected.produced_encodings()); + } +} diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 43c7ee56d68..34376673c11 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -25,7 +25,7 @@ use vortex_fastlanes::FoRArraySlotsExt; use vortex_fastlanes::for_v1_id; use vortex_fastlanes::for_v2_id; -use super::BitPackingScheme; +use super::BITPACKING_V1; use crate::ArrayAndStats; use crate::CascadingCompressor; use crate::CompressorContext; @@ -216,7 +216,7 @@ impl Scheme for FoRScheme { let leaf_ctx = compress_ctx.clone().as_leaf(); let biased_data = ArrayAndStats::new(biased.into_array(), compress_ctx.merged_stats_options()); - let compressed = BitPackingScheme.compress(compressor, &biased_data, leaf_ctx, exec_ctx)?; + let compressed = BITPACKING_V1.compress(compressor, &biased_data, leaf_ctx, exec_ctx)?; // TODO(connor): This should really be `new_unchecked`. let for_compressed = match for_array.constant_reference() { diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index 742b453c8ef..74edd53463f 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -15,6 +15,7 @@ mod zigzag; #[cfg(feature = "pco")] mod pco; +pub(crate) use bitpacking::BITPACKING_V1; pub use bitpacking::BitPackingScheme; pub use delta::DeltaScheme; pub(crate) use for_::FOR_V1; diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs index 4ca2a0e5de3..120345efd62 100644 --- a/vortex-btrblocks/src/session.rs +++ b/vortex-btrblocks/src/session.rs @@ -48,7 +48,7 @@ impl Default for CompressionSession { &integer::FOR_V1, // NOTE: ZigZag should precede BitPacking because we don't want negative numbers. &integer::ZigZagScheme, - &integer::BitPackingScheme, + &integer::BITPACKING_V1, &integer::SparseScheme, &integer::IntDictScheme, &integer::RunEndScheme, diff --git a/vortex-btrblocks/tests/bitpacking_config.rs b/vortex-btrblocks/tests/bitpacking_config.rs new file mode 100644 index 00000000000..25e3d67071a --- /dev/null +++ b/vortex-btrblocks/tests/bitpacking_config.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! BitPacking scheme refinement by allowed serialized IDs, and per-block bit widths. + +#![cfg(test)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::schemes::integer::BitPackingScheme; +use vortex_buffer::ByteBufferMut; +use vortex_edition::EDITION_DECLARATIONS; +use vortex_edition::EDITION_FAMILIES; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; +use vortex_error::VortexResult; +use vortex_fastlanes::bitpacked_v1_id; +use vortex_fastlanes::bitpacked_v2_id; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +static BITPACKING_V1: BitPackingScheme = BitPackingScheme::v1(); + +/// Registers the fastlanes encodings, and enables no editions. +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +/// Like [`SESSION`], with the latest core edition enabled: it allows `fastlanes.bitpacked` but not +/// `fastlanes.bitpacked.v2`. +static CORE_SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session().with::(); + for family in EDITION_FAMILIES { + session + .editions() + .declare_family(family) + .expect("first-party edition family"); + } + for declaration in EDITION_DECLARATIONS { + session + .register_edition(declaration) + .expect("first-party edition"); + } + session + .enable_edition(CORE_2026_08_3) + .expect("core edition is registered"); + vortex_fastlanes::initialize(&session); + session +}); + +/// Values whose 1024-value blocks need 1 to 8 bits. +fn drifting() -> ArrayRef { + PrimitiveArray::from_iter((0..8192u32).map(|i| i % (2 << (i / 1024)))).into_array() +} + +/// Values that need 7 bits in every block, so the blocks keep one bit width and the array +/// serializes as `fastlanes.bitpacked`. +fn uniform() -> ArrayRef { + PrimitiveArray::from_iter((0..8192u32).map(|i| i % 128)).into_array() +} + +/// Compresses `array`, round trips it through serialization, and returns the serialized IDs. +fn compress_roundtrip( + builder: BtrBlocksCompressorBuilder, + array: &ArrayRef, +) -> VortexResult> { + let mut ctx = SESSION.create_execution_ctx(); + let compressed = builder.build().compress(array, &mut ctx)?; + assert_arrays_eq!(array, compressed, &mut ctx); + + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + for buffer in compressed.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(&buffer); + } + let read = SerializedArray::try_from(bytes.freeze())?.decode( + array.dtype(), + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + assert_arrays_eq!(array, read, &mut ctx); + Ok(array_ctx.to_ids()) +} + +/// Only BitPacking, with every serialized ID allowed, so it refines to v2. +fn bitpacking_only() -> BtrBlocksCompressorBuilder { + BtrBlocksCompressorBuilder::empty().with_new_scheme(&BITPACKING_V1) +} + +#[rstest] +#[case::drifting(drifting(), true)] +#[case::uniform(uniform(), false)] +fn v2_serializes_varying_widths_as_v2( + #[case] array: ArrayRef, + #[case] expect_v2: bool, +) -> VortexResult<()> { + let ids = compress_roundtrip(bitpacking_only(), &array)?; + assert_eq!(ids.contains(&bitpacked_v2_id()), expect_v2); + if !expect_v2 { + assert!(ids.contains(&bitpacked_v1_id())); + } + Ok(()) +} + +#[test] +fn nullable_drifting_roundtrip() -> VortexResult<()> { + let array = PrimitiveArray::from_option_iter( + (0..8192u64).map(|i| (i % 7 != 0).then_some(i % (2 << (i / 1024)))), + ) + .into_array(); + let ids = compress_roundtrip(bitpacking_only(), &array)?; + assert!(ids.contains(&bitpacked_v2_id())); + Ok(()) +} + +#[test] +fn core_edition_keeps_global_width() -> VortexResult<()> { + let ids = compress_roundtrip( + BtrBlocksCompressorBuilder::from_session(&CORE_SESSION), + &drifting(), + )?; + assert!(!ids.contains(&bitpacked_v2_id())); + Ok(()) +} + +#[test] +fn cuda_preset_keeps_global_width() -> VortexResult<()> { + let ids = compress_roundtrip(bitpacking_only().only_cuda_compatible(), &drifting())?; + assert!(!ids.contains(&bitpacked_v2_id())); + Ok(()) +} diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs index e91995e0e4f..749aebc561b 100644 --- a/vortex-btrblocks/tests/decimal_config.rs +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -45,6 +45,7 @@ use vortex_session::registry::ReadContext; static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); static FOR_V1: FoRScheme = FoRScheme::v1(); +static BITPACKING_V1: BitPackingScheme = BitPackingScheme::v1(); /// Registers the decimal and fastlanes encodings, and enables no editions. static SESSION: LazyLock = LazyLock::new(|| { @@ -199,7 +200,7 @@ fn wide_decimal_parts_roundtrip( if compress_children { builder = builder .with_new_scheme(&FOR_V1) - .with_new_scheme(&BitPackingScheme); + .with_new_scheme(&BITPACKING_V1); } let mut ctx = SESSION.create_execution_ctx(); let compressed = builder.build().compress(&array, &mut ctx)?; diff --git a/vortex-btrblocks/tests/for_config.rs b/vortex-btrblocks/tests/for_config.rs index 812e462871c..cf6be6e4826 100644 --- a/vortex-btrblocks/tests/for_config.rs +++ b/vortex-btrblocks/tests/for_config.rs @@ -33,7 +33,7 @@ use vortex_session::VortexSession; use vortex_session::registry::ReadContext; static FOR_V1: FoRScheme = FoRScheme::v1(); -static BITPACKING: BitPackingScheme = BitPackingScheme; +static BITPACKING: BitPackingScheme = BitPackingScheme::v1(); /// Registers the fastlanes encodings, and enables no editions. static SESSION: LazyLock = LazyLock::new(|| { From 5ebf96c50a22f2827bdb952fde618195788c057e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 2 Oct 2026 18:21:13 -0400 Subject: [PATCH 2/5] always produce blocked arrays in bitpacking v2 Signed-off-by: Matt Katz --- .../src/schemes/integer/bitpacking.rs | 95 +++---------------- vortex-btrblocks/tests/bitpacking_config.rs | 19 ++-- 2 files changed, 19 insertions(+), 95 deletions(-) diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index b0cfba95fbf..adda3be4186 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -10,23 +10,14 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Patched; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_array::buffer::BufferHandle; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::patches::Patches; -use vortex_array::validity::Validity; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_fastlanes::BitPacked; -use vortex_fastlanes::BitPackedArray; use vortex_fastlanes::BitWidths; use vortex_fastlanes::bitpack_compress::bit_width_histogram; use vortex_fastlanes::bitpack_compress::bitpack_encode; @@ -56,9 +47,9 @@ pub(crate) static BITPACKING_V2: BitPackingScheme = BitPackingScheme::v2(); /// BitPacking encoding for non-negative integers. /// -/// The v1 mode packs every value at one bit width. The v2 mode chooses a width for each -/// 1024-element block. Arrays whose blocks all choose the same width keep one bit width and -/// serialize as `fastlanes.bitpacked`, while the others serialize as `fastlanes.bitpacked.v2`. +/// The v1 mode packs every value at one bit width, and serializes as `fastlanes.bitpacked`. The v2 +/// mode always chooses a width for each 1024-element block, and serializes as +/// `fastlanes.bitpacked.v2`. /// /// The default uses v1. [`refine`](Scheme::refine) picks v2 when the v2 ID is allowed and v1 /// otherwise. @@ -99,10 +90,10 @@ impl Scheme for BitPackingScheme { } fn produced_encodings(&self) -> Vec { - // Global-width arrays serialize under the frozen v1 ID. let mut encodings = match self.mode { + // Global-width arrays serialize under the frozen v1 ID. BitPackingSchemeMode::V1 => vec![bitpacked_v1_id()], - BitPackingSchemeMode::V2 => vec![bitpacked_v1_id(), bitpacked_v2_id()], + BitPackingSchemeMode::V2 => vec![bitpacked_v2_id()], }; if use_experimental_patches() { encodings.push(Patched.id()); @@ -217,8 +208,7 @@ impl Scheme for BitPackingScheme { } impl BitPackingScheme { - /// Bit-pack each 1024-element block at its own best width, keeping one bit width when every - /// block chooses the same one. + /// Bit-pack each 1024-element block at its own best width. fn compress_blocked( &self, compressor: &CascadingCompressor, @@ -235,32 +225,18 @@ impl BitPackingScheme { let BitWidths::Blocked(block_offsets) = parts.bit_widths else { vortex_bail!("Blocked bit-packing must produce block offsets"); }; - - let bit_widths = match global_bit_width(&block_offsets, exec_ctx)? { - // Like v1, keep the original array when every block needs its full width. - Some(bit_width) if usize::from(bit_width) == ptype.bit_width() => { - return Ok(primitive_array.into_array()); - } - // The blocks are packed back to back, so equal widths are also one global width. - Some(bit_width) => BitWidths::Global(bit_width), - None => BitWidths::Blocked(compressor.compress_child( - &block_offsets, - &compress_ctx, - self.id(), - 0, - exec_ctx, - )?), - }; + let block_offsets = + compressor.compress_child(&block_offsets, &compress_ctx, self.id(), 0, exec_ctx)?; let array = if use_experimental_patches() { let patches = parts.patches.take(); // Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray. - let array = new_bitpacked( + let array = BitPacked::try_new_with_block_offsets( parts.packed, ptype, parts.validity, None, - bit_widths, + block_offsets, parts.len, parts.offset, )? @@ -279,12 +255,12 @@ impl BitPackingScheme { .take() .map(|p| compress_patches(p, exec_ctx)) .transpose()?; - new_bitpacked( + BitPacked::try_new_with_block_offsets( parts.packed, ptype, parts.validity, patches, - bit_widths, + block_offsets, parts.len, parts.offset, )? @@ -296,53 +272,6 @@ impl BitPackingScheme { } } -/// The bit width shared by every block between `block_offsets`, or `None` if the widths differ. -fn global_bit_width( - block_offsets: &ArrayRef, - exec_ctx: &mut ExecutionCtx, -) -> VortexResult> { - let boundaries = block_offsets - .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? - .execute::(exec_ctx)?; - let mut bit_widths = boundaries - .as_slice::() - .windows(2) - .map(|pair| (pair[1] - pair[0]) / 128); - let Some(first) = bit_widths.next() else { - return Ok(Some(0)); - }; - if !bit_widths.all(|bit_width| bit_width == first) { - return Ok(None); - } - Ok(Some(u8::try_from(first)?)) -} - -/// Construct a bit-packed array at `bit_widths`. -fn new_bitpacked( - packed: BufferHandle, - ptype: PType, - validity: Validity, - patches: Option, - bit_widths: BitWidths, - len: usize, - offset: u16, -) -> VortexResult { - match bit_widths { - BitWidths::Global(bit_width) => { - BitPacked::try_new(packed, ptype, validity, patches, bit_width, len, offset) - } - BitWidths::Blocked(block_offsets) => BitPacked::try_new_with_block_offsets( - packed, - ptype, - validity, - patches, - block_offsets, - len, - offset, - ), - } -} - #[cfg(test)] mod tests { use rstest::rstest; diff --git a/vortex-btrblocks/tests/bitpacking_config.rs b/vortex-btrblocks/tests/bitpacking_config.rs index 25e3d67071a..37621fed83c 100644 --- a/vortex-btrblocks/tests/bitpacking_config.rs +++ b/vortex-btrblocks/tests/bitpacking_config.rs @@ -67,8 +67,7 @@ fn drifting() -> ArrayRef { PrimitiveArray::from_iter((0..8192u32).map(|i| i % (2 << (i / 1024)))).into_array() } -/// Values that need 7 bits in every block, so the blocks keep one bit width and the array -/// serializes as `fastlanes.bitpacked`. +/// Values that need 7 bits in every block. fn uniform() -> ArrayRef { PrimitiveArray::from_iter((0..8192u32).map(|i| i % 128)).into_array() } @@ -102,18 +101,13 @@ fn bitpacking_only() -> BtrBlocksCompressorBuilder { BtrBlocksCompressorBuilder::empty().with_new_scheme(&BITPACKING_V1) } +/// v2 produces per-block bit widths even when every block chooses the same width. #[rstest] -#[case::drifting(drifting(), true)] -#[case::uniform(uniform(), false)] -fn v2_serializes_varying_widths_as_v2( - #[case] array: ArrayRef, - #[case] expect_v2: bool, -) -> VortexResult<()> { +#[case::drifting(drifting())] +#[case::uniform(uniform())] +fn v2_always_serializes_as_v2(#[case] array: ArrayRef) -> VortexResult<()> { let ids = compress_roundtrip(bitpacking_only(), &array)?; - assert_eq!(ids.contains(&bitpacked_v2_id()), expect_v2); - if !expect_v2 { - assert!(ids.contains(&bitpacked_v1_id())); - } + assert!(ids.contains(&bitpacked_v2_id())); Ok(()) } @@ -141,6 +135,7 @@ fn core_edition_keeps_global_width() -> VortexResult<()> { #[test] fn cuda_preset_keeps_global_width() -> VortexResult<()> { let ids = compress_roundtrip(bitpacking_only().only_cuda_compatible(), &drifting())?; + assert!(ids.contains(&bitpacked_v1_id())); assert!(!ids.contains(&bitpacked_v2_id())); Ok(()) } From 6ab800eb0c7f85d0c70d8ab147bb076009e2da0a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 2 Oct 2026 18:29:23 -0400 Subject: [PATCH 3/5] keep the dictionary test out of blocked bitpacking's reach Signed-off-by: Matt Katz --- vortex-file/src/tests.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 62e6f39a4b1..f61c6e89eef 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1592,8 +1592,9 @@ async fn test_into_tokio_array_stream() -> VortexResult<()> { async fn test_array_stream_no_double_dict_encode() -> VortexResult<()> { let num_vals = 2048; let mut values = Vec::::with_capacity(num_vals); - values.extend(iter::repeat_n(0, num_vals / 2)); - values.extend(iter::repeat_n(1, num_vals / 2)); + // Far apart, so per-block bit widths can't pack them tighter than a dictionary. + values.extend(iter::repeat_n(1_000_000_000, num_vals / 2)); + values.extend(iter::repeat_n(2_000_000_000, num_vals / 2)); let array = PrimitiveArray::from_iter(values).into_array(); let mut buf = Vec::new(); From b01d1082360b92b4c0d86031b4a9f459f2388751 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 9 Oct 2026 10:51:18 -0400 Subject: [PATCH 4/5] use renamed bitpack_blocked_to_best_bit_widths in BitPackingScheme Signed-off-by: Matt Katz --- vortex-btrblocks/src/schemes/integer/bitpacking.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index adda3be4186..35d6bcc3c06 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -20,8 +20,8 @@ use vortex_error::vortex_bail; use vortex_fastlanes::BitPacked; use vortex_fastlanes::BitWidths; use vortex_fastlanes::bitpack_compress::bit_width_histogram; +use vortex_fastlanes::bitpack_compress::bitpack_blocked_to_best_bit_widths; use vortex_fastlanes::bitpack_compress::bitpack_encode; -use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_widths; use vortex_fastlanes::bitpack_compress::find_best_bit_width; use vortex_fastlanes::bitpacked_v1_id; use vortex_fastlanes::bitpacked_v2_id; @@ -217,7 +217,7 @@ impl BitPackingScheme { exec_ctx: &mut ExecutionCtx, ) -> VortexResult { let primitive_array = data.array_as_primitive().into_owned(); - let packed = bitpack_to_best_bit_widths(&primitive_array, exec_ctx)?; + let packed = bitpack_blocked_to_best_bit_widths(&primitive_array, exec_ctx)?; let packed_stats = packed.statistics().to_owned(); let ptype = packed.dtype().as_ptype(); From 3711acb08cdd21e25c45de7ae0ceebef9403dee2 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Mon, 12 Oct 2026 12:41:29 -0400 Subject: [PATCH 5/5] Fix bitpacking scheme fallback for uniform block widths Select block widths before encoding, retain primitive arrays when all widths are native, and use global bitpacking for uniform widths. Declare both serialized IDs for the V2 scheme and share patch handling. Cover the fallback and wire-format choices, register Patched in the test session, and update stock-ticker size expectations for the verified saving from blocked delta bases. Signed-off-by: Matt Katz --- .../array/bitpack_compress/blocked.rs | 24 ++- .../bitpacking/array/bitpack_compress/mod.rs | 1 + vortex-btrblocks/src/builder.rs | 18 ++ .../src/schemes/integer/bitpacking.rs | 189 ++++++++---------- vortex-btrblocks/tests/bitpacking_config.rs | 38 +++- vortex-file/src/tests.rs | 2 +- 6 files changed, 152 insertions(+), 120 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs index b2db108f15c..2a3743f5c64 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/blocked.rs @@ -50,12 +50,28 @@ pub fn bitpack_blocked_to_best_bit_widths( array: &PrimitiveArray, ctx: &mut ExecutionCtx, ) -> VortexResult { + let (bit_widths, num_exceptions) = find_best_bit_widths(array, ctx)?; + bitpack_encode_blocked(array, &bit_widths, Some(num_exceptions), ctx) +} + +/// Choose the width minimizing each 1024-value block's packed size plus the cost of its +/// exceptions, and return the widths and total number of exceptions. +/// +/// Widths may equal the array's native bit width when packing would not help. An empty array +/// produces no widths and no exceptions. +/// +/// # Errors +/// +/// Returns an error if `array` is not an integer array or contains negative values. +pub fn find_best_bit_widths( + array: &PrimitiveArray, + ctx: &mut ExecutionCtx, +) -> VortexResult<(Vec, usize)> { ensure_non_negative_integers(array, ctx)?; let validity_mask = array.validity()?.execute_mask(array.len(), ctx)?; - let (bit_widths, num_exceptions) = match_each_integer_ptype!(array.ptype(), |T| { - block_bit_widths(array.as_slice::(), &validity_mask)? - }); - bitpack_encode_blocked(array, &bit_widths, Some(num_exceptions), ctx) + match_each_integer_ptype!(array.ptype(), |T| { + block_bit_widths(array.as_slice::(), &validity_mask) + }) } /// Bit-pack each 1024-value block of `array` at its width in `bit_widths`. diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs index 7806536877c..b60bf1e6513 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress/mod.rs @@ -6,6 +6,7 @@ mod global; pub use blocked::bitpack_blocked_to_best_bit_widths; pub use blocked::bitpack_encode_blocked; +pub use blocked::find_best_bit_widths; pub use global::bit_width_histogram; pub use global::bitpack_encode; pub use global::bitpack_encode_unchecked; diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index d3dea3fdb53..93179dab638 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -250,9 +250,12 @@ mod tests { use rstest::rstest; use vortex_array::ArrayId; use vortex_array::VTable; + use vortex_array::arrays::Patched; use vortex_array::arrays::VarBin; use vortex_decimal_byte_parts::DecimalByteParts; use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; + use vortex_fastlanes::bitpacked_v1_id; + use vortex_fastlanes::bitpacked_v2_id; use vortex_fsst::FSST; use super::*; @@ -265,6 +268,21 @@ mod tests { #[case::fsst_with_codes(&string::FSSTScheme, vec![FSST.id(), VarBin.id()], true)] #[case::decimal_v1_id(&DECIMAL_V1, vec![decimal_byte_parts_v1_id()], true)] #[case::decimal_both_ids(&DECIMAL_V2, vec![decimal_byte_parts_v1_id(), DecimalByteParts.id()], true)] + #[case::bitpacking_v1_id( + &integer::BITPACKING_V1, + vec![bitpacked_v1_id(), Patched.id()], + true + )] + #[case::bitpacking_both_ids( + &integer::BITPACKING_V1, + vec![bitpacked_v1_id(), bitpacked_v2_id(), Patched.id()], + true + )] + #[case::bitpacking_v2_only( + &integer::BITPACKING_V1, + vec![bitpacked_v2_id(), Patched.id()], + false + )] fn test_allowed_schemes( #[case] scheme: &'static dyn Scheme, #[case] allowed: Vec, diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 35d6bcc3c06..941c354da7c 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -3,6 +3,7 @@ //! BitPacking integer encoding. +use itertools::Itertools; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -10,19 +11,21 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Patched; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedArray; use vortex_fastlanes::BitWidths; use vortex_fastlanes::bitpack_compress::bit_width_histogram; -use vortex_fastlanes::bitpack_compress::bitpack_blocked_to_best_bit_widths; use vortex_fastlanes::bitpack_compress::bitpack_encode; +use vortex_fastlanes::bitpack_compress::bitpack_encode_blocked; use vortex_fastlanes::bitpack_compress::find_best_bit_width; +use vortex_fastlanes::bitpack_compress::find_best_bit_widths; use vortex_fastlanes::bitpacked_v1_id; use vortex_fastlanes::bitpacked_v2_id; @@ -48,8 +51,8 @@ pub(crate) static BITPACKING_V2: BitPackingScheme = BitPackingScheme::v2(); /// BitPacking encoding for non-negative integers. /// /// The v1 mode packs every value at one bit width, and serializes as `fastlanes.bitpacked`. The v2 -/// mode always chooses a width for each 1024-element block, and serializes as -/// `fastlanes.bitpacked.v2`. +/// mode chooses a width for each 1024-element block. Uniform widths use the v1 format; varying +/// widths serialize as `fastlanes.bitpacked.v2`. /// /// The default uses v1. [`refine`](Scheme::refine) picks v2 when the v2 ID is allowed and v1 /// otherwise. @@ -93,7 +96,7 @@ impl Scheme for BitPackingScheme { let mut encodings = match self.mode { // Global-width arrays serialize under the frozen v1 ID. BitPackingSchemeMode::V1 => vec![bitpacked_v1_id()], - BitPackingSchemeMode::V2 => vec![bitpacked_v2_id()], + BitPackingSchemeMode::V2 => vec![bitpacked_v1_id(), bitpacked_v2_id()], }; if use_experimental_patches() { encodings.push(Patched.id()); @@ -140,57 +143,32 @@ impl Scheme for BitPackingScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - if self.mode == BitPackingSchemeMode::V2 { - return self.compress_blocked(compressor, data, compress_ctx, exec_ctx); - } - - let primitive_array = data.array_as_primitive(); - - let histogram = bit_width_histogram(primitive_array, exec_ctx)?; - let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?; - - // If best bw is determined to be the current bit-width, return the original array. - if bw as usize == primitive_array.ptype().bit_width() { - 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 primitive_array = data.array_as_primitive().into_owned(); + let packed = match self.mode { + BitPackingSchemeMode::V1 => compress_global(&primitive_array, exec_ctx)?, + BitPackingSchemeMode::V2 => compress_blocked(&primitive_array, exec_ctx)?, + }; + let Some(packed) = packed else { + return Ok(primitive_array.into_array()); + }; let packed_stats = packed.statistics().to_owned(); let ptype = packed.dtype().as_ptype(); let mut parts = BitPacked::into_parts(packed); - let array = if use_experimental_patches() { - let patches = parts.patches.take(); - // Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray. - let array = BitPacked::try_new( - parts.packed, - ptype, - parts.validity, - None, - bw, - parts.len, - parts.offset, - )? - .into_array(); - - match patches { - None => array, - Some(p) => Patched::from_array_and_patches(array, &p, exec_ctx)? - .with_stats_set(packed_stats) - .into_array(), - } + let patches = if use_experimental_patches() { + parts.patches.take() } else { - // Compress patches and place back into BitPackedArray. - let patches = parts + parts.patches = parts .patches .take() .map(|p| compress_patches(p, exec_ctx)) .transpose()?; - parts.patches = patches; - BitPacked::try_new( + None + }; + + let array = match parts.bit_widths { + BitWidths::Global(bw) => BitPacked::try_new( parts.packed, ptype, parts.validity, @@ -198,77 +176,68 @@ impl Scheme for BitPackingScheme { bw, parts.len, parts.offset, - )? - .with_stats_set(packed_stats) - .into_array() + )?, + BitWidths::Blocked(block_offsets) => { + let block_offsets = compressor.compress_child( + &block_offsets, + &compress_ctx, + self.id(), + 0, + exec_ctx, + )?; + BitPacked::try_new_with_block_offsets( + parts.packed, + ptype, + parts.validity, + parts.patches, + block_offsets, + parts.len, + parts.offset, + )? + } }; - Ok(array) + match patches { + None => Ok(array.with_stats_set(packed_stats).into_array()), + Some(p) => Ok( + Patched::from_array_and_patches(array.into_array(), &p, exec_ctx)? + .with_stats_set(packed_stats) + .into_array(), + ), + } } } -impl BitPackingScheme { - /// Bit-pack each 1024-element block at its own best width. - fn compress_blocked( - &self, - compressor: &CascadingCompressor, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let primitive_array = data.array_as_primitive().into_owned(); - let packed = bitpack_blocked_to_best_bit_widths(&primitive_array, exec_ctx)?; - - let packed_stats = packed.statistics().to_owned(); - let ptype = packed.dtype().as_ptype(); - let mut parts = BitPacked::into_parts(packed); - let BitWidths::Blocked(block_offsets) = parts.bit_widths else { - vortex_bail!("Blocked bit-packing must produce block offsets"); - }; - let block_offsets = - compressor.compress_child(&block_offsets, &compress_ctx, self.id(), 0, exec_ctx)?; - - let array = if use_experimental_patches() { - let patches = parts.patches.take(); - // Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray. - let array = BitPacked::try_new_with_block_offsets( - parts.packed, - ptype, - parts.validity, - None, - block_offsets, - parts.len, - parts.offset, - )? - .into_array(); +/// Returns `None` when the best width is the native width. +fn compress_global( + primitive_array: &PrimitiveArray, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult> { + let histogram = bit_width_histogram(primitive_array.as_view(), exec_ctx)?; + let bw = find_best_bit_width(primitive_array.ptype(), &histogram)?; + if usize::from(bw) == primitive_array.ptype().bit_width() { + return Ok(None); + } + bitpack_encode(primitive_array, bw, Some(&histogram), exec_ctx).map(Some) +} - 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()?; - BitPacked::try_new_with_block_offsets( - parts.packed, - ptype, - parts.validity, - patches, - block_offsets, - parts.len, - parts.offset, - )? - .with_stats_set(packed_stats) - .into_array() - }; +/// Uses global packing for uniform widths. Returns `None` for empty input or all-native widths. +fn compress_blocked( + primitive_array: &PrimitiveArray, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult> { + let (bit_widths, num_exceptions) = find_best_bit_widths(primitive_array, exec_ctx)?; + if bit_widths.is_empty() { + return Ok(None); + } - Ok(array) + match bit_widths.iter().copied().all_equal_value() { + Ok(bw) if usize::from(bw) == primitive_array.ptype().bit_width() => Ok(None), + Ok(bw) => bitpack_encode(primitive_array, bw, None, exec_ctx).map(Some), + Err(_) => { + bitpack_encode_blocked(primitive_array, &bit_widths, Some(num_exceptions), exec_ctx) + .map(Some) + } } } diff --git a/vortex-btrblocks/tests/bitpacking_config.rs b/vortex-btrblocks/tests/bitpacking_config.rs index 37621fed83c..0fe0323bfa9 100644 --- a/vortex-btrblocks/tests/bitpacking_config.rs +++ b/vortex-btrblocks/tests/bitpacking_config.rs @@ -13,10 +13,12 @@ use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Patched; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; +use vortex_array::session::ArraySessionExt; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::schemes::integer::BitPackingScheme; use vortex_buffer::ByteBufferMut; @@ -36,6 +38,7 @@ static BITPACKING_V1: BitPackingScheme = BitPackingScheme::v1(); /// Registers the fastlanes encodings, and enables no editions. static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); + session.arrays().register(Patched); vortex_fastlanes::initialize(&session); session }); @@ -101,13 +104,38 @@ fn bitpacking_only() -> BtrBlocksCompressorBuilder { BtrBlocksCompressorBuilder::empty().with_new_scheme(&BITPACKING_V1) } -/// v2 produces per-block bit widths even when every block chooses the same width. #[rstest] -#[case::drifting(drifting())] -#[case::uniform(uniform())] -fn v2_always_serializes_as_v2(#[case] array: ArrayRef) -> VortexResult<()> { +#[case::drifting(drifting(), Some(bitpacked_v2_id()))] +#[case::uniform(uniform(), Some(bitpacked_v1_id()))] +#[case::uniform_with_patches( + PrimitiveArray::from_iter( + (0..8192u32).map(|i| if i % 1024 == 0 { 1 << 20 } else { i % 128 }), + ) + .into_array(), + Some(bitpacked_v1_id()) +)] +#[case::native_u8( + PrimitiveArray::from_iter((128u8..=u8::MAX).cycle().take(8192)).into_array(), + None +)] +#[case::native_u32( + PrimitiveArray::from_iter((0..8192u32).map(|i| (1 << 31) + i)).into_array(), + None +)] +#[case::empty(PrimitiveArray::from_iter(Vec::::new()).into_array(), None)] +fn v2_selects_wire_format( + #[case] array: ArrayRef, + #[case] expected_id: Option, +) -> VortexResult<()> { let ids = compress_roundtrip(bitpacking_only(), &array)?; - assert!(ids.contains(&bitpacked_v2_id())); + assert_eq!( + ids.contains(&bitpacked_v1_id()), + expected_id == Some(bitpacked_v1_id()) + ); + assert_eq!( + ids.contains(&bitpacked_v2_id()), + expected_id == Some(bitpacked_v2_id()) + ); Ok(()) } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f61c6e89eef..0098fbeae52 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -149,7 +149,7 @@ async fn test_eof_values() { #[rstest] #[case::default( BtrBlocksCompressorBuilder::from_session(&SESSION), - if cfg!(feature = "zstd") { 70_100 } else { 70_036 } + if cfg!(feature = "zstd") { 70_068 } else { 70_004 } )] #[cfg_attr( feature = "zstd",