From d21cb9dadc385658087ea6bc5bb2c38b6435127a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 3 Sep 2026 17:32:52 -0400 Subject: [PATCH 1/2] feat(compressor): tell schemes which serialized IDs the writer allows CascadingCompressor carries the serialized IDs its output may use. The file writer passes the same set it gives the array context, next to the existing filter on allowed encodings, so a scheme whose encoding has several wire formats can produce the newest one still allowed without the compressor knowing about editions. The CUDA preset allows only each encoding's original format, since it has no per-chunk bit-unpacking kernel yet. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 47 ++++++++++++++++++++++++- vortex-compressor/src/compressor/mod.rs | 31 ++++++++++++++++ vortex-file/src/writer.rs | 25 +++++++++---- 3 files changed, 96 insertions(+), 7 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index fe8072d5e66..0cee563820a 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/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>, } 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) -> 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(); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..4b26439fd5a 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,9 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,6 +49,10 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// The serialized IDs the output may use, or `None` for no restriction. See + /// [`allows_serialized_id`](Self::allows_serialized_id). + allowed_serialized_ids: Option>, } impl CascadingCompressor { @@ -63,9 +70,33 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + allowed_serialized_ids: None, } } + /// Restricts the output to the serialized IDs in `allowed`, intersecting with any earlier + /// restriction. + /// + /// The file writer passes the serialized IDs its editions permit, so a scheme whose encoding + /// has several wire formats produces the newest one still allowed. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(&allowed).copied().collect(), + None => allowed, + }); + self + } + + /// Returns whether the output may use the serialized ID `id`. + /// + /// Schemes whose encoding has several wire formats consult this to pick the newest allowed + /// one. Without a restriction every ID is allowed, so that is always the newest. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns whether the compressor was configured with `scheme`. pub fn has_scheme(&self, scheme: SchemeId) -> bool { self.schemes diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index ec45653f5c1..220599733fb 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,7 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings) = + let (array_ctx, allowed_array_encodings, allowed_serialized_ids) = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); @@ -253,7 +253,8 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings), + .retain_allowed_encodings(&allowed_array_encodings) + .allow_serialized_ids(&allowed_serialized_ids), ) .build(), }; @@ -384,10 +385,12 @@ impl VortexWriteOptions { } } +/// Returns the array context, the in-memory encodings the compressor may produce, and the +/// serialized IDs its output may use. fn new_array_context( session: &VortexSession, enforce_editions: bool, -) -> (ArrayContext, HashSet) { +) -> (ArrayContext, HashSet, HashSet) { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done @@ -406,6 +409,9 @@ fn new_array_context( .filter_map(|serialized_id| arrays.registry().get(serialized_id)) .map(|plugin| plugin.id()) .collect(); + // The compressor sees the same set, so an encoding with several wire formats produces the + // newest one the editions permit. + let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { // Only permit serialized IDs in the enabled editions. @@ -413,7 +419,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings) + (array_ctx, allowed_array_encodings, allowed_serialized_ids) } /// The ids of `kind` the enabled editions permit. @@ -787,10 +793,12 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings) = new_array_context(&session, true); + let (ctx, allowed_array_encodings, allowed_serialized_ids) = + new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); + assert_eq!(allowed_serialized_ids, HashSet::from([Primitive.id()])); Ok(()) } @@ -807,9 +815,14 @@ mod tests { ) }); - let (ctx, allowed_array_encodings) = new_array_context(&session, false); + let (ctx, allowed_array_encodings, allowed_serialized_ids) = + new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); assert_eq!(allowed_array_encodings, registered_encodings); + assert_eq!( + allowed_serialized_ids, + registered_ids.iter().copied().collect::>() + ); assert!(ctx.intern(&Bool.id()).is_some()); } From 114a0f0454d0d0380b3e93f18e9884b3da4d3c64 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 3 Sep 2026 17:32:52 -0400 Subject: [PATCH 2/2] feat(btrblocks): pick per-chunk bit widths when fastlanes.bitpacked_v2 is allowed BitPackingScheme chooses a width per 1024-element chunk when the compressor may use the v2 format and one global width otherwise, so an old edition keeps writing the original format and a newer one gets per-chunk widths from the same scheme. The width table child is re-encoded through the cascade. The scheme returns the original array when half or more of the values would be patches. A few wide values otherwise pack at width 0 with nearly every value patched, which beats raw storage by a couple of buffer bytes and loses them back in footer metadata. No shipped edition permits the format yet; the writer test declares its own. Signed-off-by: Matt Katz --- Cargo.lock | 1 + .../src/schemes/integer/bitpacking.rs | 124 +++++++++----- .../schemes/integer/scheme_selection_tests.rs | 17 ++ .../golden__default__int_monotone_jitter.snap | 8 +- .../golden__default__list_of_int_runs.snap | 14 +- ...lden__default__string_fsst_structured.snap | 14 +- ...n__default__temporal_timestamp_micros.snap | 10 +- ...den__unstable__string_fsst_structured.snap | 28 +-- vortex/Cargo.toml | 1 + vortex/src/editions/tests.rs | 160 ++++++++++++++++++ 10 files changed, 311 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3bf63e6512b..cccd094f6cb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10470,6 +10470,7 @@ dependencies = [ "codspeed-divan-compat", "fastlanes 0.7.0", "mimalloc", + "parking_lot", "parquet 59.2.0", "rand 0.10.2", "serde_json", diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 78e63fb6dee..985d7431dd4 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/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 { 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 { + 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) +} diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index 1a530c61e72..4d8c46be380 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/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 = 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::(), + "got {}", + compressed.encoding_id() + ); + Ok(()) +} diff --git a/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap index 4800aae289d..82cebefe54c 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__int_monotone_jitter.snap +++ b/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 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap index c9554add05a..b05d7d17db4 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__list_of_int_runs.snap +++ b/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 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap index 327f050b0d7..28bb186bb8d 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__default__string_fsst_structured.snap @@ -3,7 +3,7 @@ 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 @@ -11,5 +11,13 @@ root: vortex.fsst(utf8, len=16384) nbytes=151382 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 diff --git a/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap index 6f38e8e6f5d..a954d2308bf 100644 --- a/vortex-btrblocks/tests/snapshots/golden__default__temporal_timestamp_micros.snap +++ b/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 diff --git a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap index c0408862141..1507a32a95a 100644 --- a/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__unstable__string_fsst_structured.snap @@ -3,31 +3,35 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=653785 -root: vortex.fsst(utf8, len=16384) nbytes=121766 +root: vortex.fsst(utf8, len=16384) nbytes=121333 metadata: len: 16384, nsymbols: 223 - uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=1802 + uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=1736 metadata: fill_value: 24u8 - patch_indices: fastlanes.delta(u16, len=1575) nbytes=1798 + patch_indices: fastlanes.delta(u16, len=1575) nbytes=1732 metadata: offset: 0 bases: vortex.primitive(u16, len=128) nbytes=256 metadata: ptype: u16 - deltas: fastlanes.bitpacked(u16, len=2048) nbytes=1542 - metadata: bit_width: 6, offset: 0 - patch_indices: vortex.primitive(u16, len=1) nbytes=2 + deltas: fastlanes.bitpacked(u16, len=2048) nbytes=1476 + metadata: bit_widths: 2 chunks, max 6, offset: 0 + patch_indices: vortex.primitive(u16, len=16) nbytes=32 + metadata: ptype: u16 + patch_values: vortex.primitive(u16, len=16) nbytes=32 metadata: ptype: u16 - patch_values: vortex.constant(u16, len=1) nbytes=2 - metadata: scalar: 67u16 patch_chunk_offsets: vortex.primitive(u8, len=2) nbytes=2 metadata: ptype: u8 + width_table: vortex.primitive(u8, len=2) nbytes=2 + metadata: ptype: u8 patch_values: vortex.constant(u8, len=1575) nbytes=2 metadata: scalar: 23u8 - codes_offsets: fastlanes.delta(u32, len=16385) nbytes=8728 + codes_offsets: fastlanes.delta(u32, len=16385) nbytes=8361 metadata: offset: 0 bases: vortex.primitive(u32, len=544) nbytes=2176 metadata: ptype: u32 - deltas: vortex.dict(u32, len=17408) nbytes=6552 + deltas: vortex.dict(u32, len=17408) nbytes=6185 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=17408) nbytes=6528 - metadata: bit_width: 3, offset: 0 + codes: fastlanes.bitpacked(u8, len=17408) nbytes=6161 + metadata: bit_widths: 17 chunks, max 3, offset: 0 + width_table: vortex.primitive(u8, len=17) nbytes=17 + metadata: ptype: u8 values: vortex.primitive(u32, len=6) nbytes=24 metadata: ptype: u32 diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index e90f6f6e8e7..2030c302460 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -57,6 +57,7 @@ vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } [dev-dependencies] +parking_lot = { workspace = true } anyhow = { workspace = true } arrow-array = { workspace = true } divan = { workspace = true } diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index f3cbedc9ffb..e57f30f609f 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -1,7 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + +use parking_lot::Mutex; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::ArrayVTable; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::array_session; @@ -15,6 +23,7 @@ use vortex_array::dtype::PType; use vortex_array::extension::datetime::Date; use vortex_array::extension::datetime::TimeUnit; use vortex_array::session::ArraySessionExt; +use vortex_array::stream::ArrayStreamExt; use vortex_buffer::ByteBufferMut; use vortex_edition::ComponentKind; use vortex_edition::Edition; @@ -28,6 +37,9 @@ use vortex_edition::EditionSessionExt; use vortex_edition::test_harness::validate_edition; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedPlugin; +use vortex_fastlanes::bitpacked_v2_id; use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; @@ -663,3 +675,151 @@ async fn serialization_context_accepts_supported_compressor_output() -> VortexRe Ok(()) } + +/// Eight FastLanes chunks whose values need 1, 4, 7, ... 22 bits: a single global bit width wastes +/// most of the array, so the compressor picks per-chunk widths whenever it may produce them. +fn drifting_integers() -> PrimitiveArray { + PrimitiveArray::from_iter((0..8 * 1024u32).map(|i| { + let width = 1 + 3 * (i / 1024); + i.wrapping_mul(2_654_435_761) >> (32 - width) + })) +} + +/// Stands in for `BitPackedPlugin` and records the serialized format ID of every bit-packed +/// array it is asked to read, so a test can see which wire format a file actually carries. +#[derive(Debug, Clone, Default)] +struct RecordingBitPacked(Arc>>); + +impl ArrayPlugin for RecordingBitPacked { + fn id(&self) -> ArrayId { + BitPackedPlugin.id() + } + + fn serialized_ids(&self) -> Vec { + BitPackedPlugin.serialized_ids() + } + + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult> { + BitPackedPlugin.serialize(array, session) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + self.0.lock().push(parts.serialized_id); + BitPackedPlugin.deserialize(parts, session) + } +} + +impl RecordingBitPacked { + /// Register a recorder on `session` in place of the stock `BitPackedPlugin`. + fn install(session: &VortexSession) -> Self { + let recorder = Self::default(); + session.arrays().register(recorder.clone()); + recorder + } + + /// Read the whole file and return the distinct serialized IDs of the bit-packed arrays in it. + async fn ids_read( + &self, + session: &VortexSession, + buffer: ByteBufferMut, + ) -> VortexResult> { + self.0.lock().clear(); + session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + let mut ids = self.0.lock().clone(); + ids.sort_unstable(); + ids.dedup(); + Ok(ids) + } +} + +/// A session with the default encodings and editions registered, enabling only the default core +/// edition regardless of the `unstable_encodings` feature. +fn core_only_session() -> VortexResult { + let session = array_session() + .with::() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + super::register_default_editions(&session); + session + .enable_edition(DEFAULT_CORE_EDITION) + .map_err(|error| vortex_err!("{error}"))?; + Ok(session) +} + +/// Per-chunk bit widths are not in any core edition, so a core-only writer never emits them even +/// when the compressor knows how to produce them. +#[tokio::test] +async fn core_writer_never_emits_bitpacked_v2() -> VortexResult<()> { + let session = core_only_session()?; + let recorder = RecordingBitPacked::install(&session); + let buffer = write_with(&session, drifting_integers().into_array()).await?; + let ids = recorder.ids_read(&session, buffer).await?; + assert_eq!( + ids, + [ArrayVTable::id(&BitPacked)], + "core-only write must bit-pack with the original format only" + ); + Ok(()) +} + +const PER_CHUNK_TEST_EDITION: EditionId = EditionId::new("bitpacked-v2-test", 2026, 9, 0); + +/// A test-only edition permitting the per-chunk bit width format, standing in for whichever +/// edition eventually ships it. +static PER_CHUNK_TEST_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: PER_CHUNK_TEST_EDITION, + min_library_version: None, + }, + added: &[EditionMember::array(&"fastlanes.bitpacked_v2")], +}; + +/// Once an enabled edition permits `fastlanes.bitpacked_v2`, the same compressor writes a +/// drifting column in that format and it reads back unchanged. +#[tokio::test] +async fn permitting_writer_emits_bitpacked_v2() -> VortexResult<()> { + use vortex_array::VortexSessionExecute; + + let session = core_only_session()?; + session + .register_edition(&PER_CHUNK_TEST_DECLARATION) + .map_err(|error| vortex_err!("{error}"))?; + session + .enable_edition(PER_CHUNK_TEST_EDITION) + .map_err(|error| vortex_err!("{error}"))?; + let recorder = RecordingBitPacked::install(&session); + let values = drifting_integers(); + let buffer = write_with(&session, values.clone().into_array()).await?; + + let read = session + .open_options() + .open_buffer(buffer.clone())? + .scan()? + .into_array_stream()? + .read_all() + .await? + .execute::(&mut session.create_execution_ctx())?; + assert_eq!(read.as_slice::(), values.as_slice::()); + + let ids = recorder.ids_read(&session, buffer).await?; + assert!( + ids.contains(&bitpacked_v2_id()), + "permitting write produced {ids:?}" + ); + Ok(()) +}