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/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-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-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()); } 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(()) +}