From d0198afbba1f40c7bb51f72e662eee66f36f4217 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 14:35:09 -0400 Subject: [PATCH] Pass allowed serialized IDs through the compressor context Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 261 ++++++++++++++++---- vortex-btrblocks/src/schemes/decimal.rs | 50 ++-- vortex-btrblocks/tests/decimal.rs | 134 ++++++++++ vortex-compressor/src/compressor/cascade.rs | 6 +- vortex-compressor/src/compressor/mod.rs | 43 ++++ vortex-compressor/src/compressor/tests.rs | 91 +++++++ vortex-compressor/src/scheme/ctx.rs | 25 ++ vortex-compressor/src/scheme/mod.rs | 7 +- vortex/src/editions/tests.rs | 114 +++++++++ 9 files changed, 660 insertions(+), 71 deletions(-) create mode 100644 vortex-btrblocks/tests/decimal.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 3bcda909227..c8400be2797 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -4,6 +4,7 @@ //! Builder for configuring `BtrBlocksCompressor` instances. use vortex_array::ArrayId; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; use vortex_utils::aliases::hash_set::HashSet; use crate::BtrBlocksCompressor; @@ -79,6 +80,11 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); /// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the /// `zstd` feature is enabled. /// +/// [`retain_allowed_encodings`](Self::retain_allowed_encodings) and +/// [`exclude_encodings`](Self::exclude_encodings) restrict the serialized IDs the compressor may +/// emit. Both apply in [`build`](Self::build), so they cover schemes registered afterwards and may +/// be called in any order. +/// /// # Examples /// /// ```rust @@ -96,12 +102,18 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + /// Serialized IDs the writer permits, or `None` when no writer restriction was supplied. + allowed_encodings: Option>, + /// Serialized IDs denied regardless of what the writer permits. + excluded_encodings: HashSet, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + allowed_encodings: None, + excluded_encodings: HashSet::new(), } } } @@ -113,6 +125,8 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_encodings: None, + excluded_encodings: HashSet::new(), } } @@ -158,8 +172,8 @@ impl BtrBlocksCompressorBuilder { builder } - /// Excludes schemes without CUDA kernel support, keeps FSST for string compression, - /// and adds Zstd for binary compression. + /// Excludes schemes and wire formats without CUDA kernel support, keeps FSST for string + /// compression, and adds Zstd for binary compression. /// /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, @@ -189,7 +203,11 @@ impl BtrBlocksCompressorBuilder { excluded.push(integer::DeltaScheme::default().id()); #[cfg(feature = "pco")] excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); - let builder = self.exclude_schemes(excluded); + // Multi-part DecimalByteParts arrays have no CUDA decode kernel, so wide decimals stay + // canonical while single-part decimals still compress under the v1 format. + let builder = self + .exclude_schemes(excluded) + .exclude_encodings([decimal_byte_parts_v2_id()]); #[cfg(feature = "zstd")] let builder = builder @@ -206,29 +224,61 @@ impl BtrBlocksCompressorBuilder { self } - /// Retains only schemes whose produced serialized IDs all belong to `allowed`. + /// Retains only schemes whose produced serialized IDs all belong to `allowed`, and permits + /// `allowed` to the schemes that remain. /// - /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. + /// `allowed` holds serialized IDs. The file writer passes the IDs its enabled editions permit. + /// Repeated calls intersect. The restriction applies in [`build`](Self::build), so it covers + /// schemes registered after this call, and the remaining schemes only emit optional wire + /// formats the set contains. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + match &mut self.allowed_encodings { + Some(current) => current.retain(|id| allowed.contains(id)), + None => self.allowed_encodings = Some(allowed.clone()), + } + self + } + + /// Denies the given serialized IDs, whatever + /// [`retain_allowed_encodings`](Self::retain_allowed_encodings) permits. + /// + /// Like `retain_allowed_encodings`, this applies in [`build`](Self::build). + pub fn exclude_encodings(mut self, ids: impl IntoIterator) -> Self { + self.excluded_encodings.extend(ids); self } /// Builds the configured [`BtrBlocksCompressor`]. + /// + /// Without [`retain_allowed_encodings`](Self::retain_allowed_encodings), the compressor may + /// emit exactly the serialized IDs its schemes declare, so optional wire formats need explicit + /// permission. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + let compressor = CascadingCompressor::new(self.schemes); + let mut allowed = self + .allowed_encodings + .unwrap_or_else(|| compressor.allowed_serialized_ids().clone()); + allowed.retain(|id| !self.excluded_encodings.contains(id)); + BtrBlocksCompressor(compressor.with_allowed_serialized_ids(allowed)) } } #[cfg(test)] mod tests { + use rstest::rstest; use vortex_array::VTable; + use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; use vortex_fastlanes::FoR; use super::*; + fn declared_ids() -> HashSet { + ALL_SCHEMES + .iter() + .flat_map(|scheme| scheme.produced_encodings()) + .collect() + } + #[test] fn empty_starts_with_no_schemes() { let builder = BtrBlocksCompressorBuilder::empty(); @@ -237,53 +287,175 @@ mod tests { #[test] fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); + let compressor = BtrBlocksCompressorBuilder::default().build(); + for scheme in ALL_SCHEMES { + assert!(compressor.has_scheme(scheme.id())); + } + } + + /// Without a writer allowlist the compressor may emit what its schemes declare and nothing + /// more, so an optional wire format such as DecimalByteParts v2 stays off. + #[test] + fn default_permits_only_declared_ids() { + let compressor = BtrBlocksCompressorBuilder::default().build(); + assert_eq!(compressor.allowed_serialized_ids(), &declared_ids()); + assert!( + !compressor + .allowed_serialized_ids() + .contains(&decimal_byte_parts_v2_id()) + ); } #[test] - fn retain_allowed_encodings_filters_schemes() { + fn allowed_encodings_filter_schemes_on_build() { let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); + let compressor = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed) + .build(); + assert!(compressor.has_scheme(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(integer::BitPackingScheme.id())); - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); + let none = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&HashSet::new()) + .build(); + for scheme in ALL_SCHEMES { + assert!(!none.has_scheme(scheme.id())); + } } #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); + fn allowing_all_declared_ids_keeps_every_scheme() { + let compressor = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&declared_ids()) + .build(); + for scheme in ALL_SCHEMES { + assert!(compressor.has_scheme(scheme.id())); + } } + #[rstest] + fn allowed_encodings_apply_regardless_of_registration_order( + #[values(false, true)] permitted: bool, + #[values(false, true)] register_later: bool, + ) { + let allowed = if permitted { + HashSet::from([FoR.id()]) + } else { + HashSet::new() + }; + let mut builder = BtrBlocksCompressorBuilder::empty(); + if !register_later { + builder = builder.with_new_scheme(&integer::FoRScheme); + } + builder = builder.retain_allowed_encodings(&allowed); + if register_later { + builder = builder.with_new_scheme(&integer::FoRScheme); + } + assert_eq!( + builder.build().has_scheme(integer::FoRScheme.id()), + permitted + ); + } + + #[rstest] + fn allowed_encodings_intersect(#[values(false, true)] restrictive_first: bool) { + let all = declared_ids(); + let restricted = HashSet::from([FoR.id()]); + let (first, second) = if restrictive_first { + (&restricted, &all) + } else { + (&all, &restricted) + }; + let compressor = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(first) + .retain_allowed_encodings(second) + .build(); + assert!(compressor.has_scheme(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(integer::BitPackingScheme.id())); + } + + /// A permitted optional format reaches the compressor's allowlist without being declared by any + /// scheme, so schemes can pick it up through the compression context. #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + fn allowed_encodings_permit_optional_ids() { + let compressor = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&HashSet::from([ + decimal_byte_parts_v1_id(), + decimal_byte_parts_v2_id(), + ])) + .build(); + assert!(compressor.has_scheme(decimal::DecimalScheme.id())); assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) + compressor + .allowed_serialized_ids() + .contains(&decimal_byte_parts_v2_id()) ); } + #[test] + fn excluded_encodings_override_allowed_encodings() { + let mut allowed = declared_ids(); + allowed.insert(decimal_byte_parts_v2_id()); + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_encodings([decimal_byte_parts_v2_id(), FoR.id()]) + .retain_allowed_encodings(&allowed) + .build(); + // Denying an optional format keeps the scheme but withholds the format. + assert!(compressor.has_scheme(decimal::DecimalScheme.id())); + assert!( + !compressor + .allowed_serialized_ids() + .contains(&decimal_byte_parts_v2_id()) + ); + // Denying a declared format drops the scheme that needs it. + assert!(!compressor.has_scheme(integer::FoRScheme.id())); + } + + #[rstest] + #[case::without_allowlist(None)] + #[case::allowlist_first(Some(true))] + #[case::allowlist_last(Some(false))] + fn cuda_compatible_denies_decimal_v2(#[case] allowlist_first: Option) { + let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + let builder = match allowlist_first { + None => BtrBlocksCompressorBuilder::default().only_cuda_compatible(), + Some(true) => BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed) + .only_cuda_compatible(), + Some(false) => BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .retain_allowed_encodings(&allowed), + }; + let compressor = builder.build(); + assert!(compressor.has_scheme(decimal::DecimalScheme.id())); + assert!( + !compressor + .allowed_serialized_ids() + .contains(&decimal_byte_parts_v2_id()) + ); + } + + #[test] + fn cuda_compatible_excludes_alprd() { + let compressor = BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .build(); + assert!(!compressor.has_scheme(float::ALPRDScheme.id())); + } + /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. #[test] fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + let compressor = BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .build(); for excluded in [ integer::SparseScheme.id(), float::NullDominatedSparseScheme.id(), string::NullDominatedSparseScheme.id(), ] { assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), + !compressor.has_scheme(excluded), "{excluded} should be excluded" ); } @@ -291,31 +463,24 @@ mod tests { #[test] fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); + let compressor = BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .build(); + assert!(compressor.has_scheme(string::FSSTScheme.id())); #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); + assert!(!compressor.has_scheme(string::ZstdScheme.id())); } #[test] #[cfg(feature = "pco")] fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() + let compressor = BtrBlocksCompressorBuilder::default() .with_new_scheme(&integer::PcoScheme) .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); + .only_cuda_compatible() + .build(); for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); + assert!(!compressor.has_scheme(scheme)); } } } diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index f77a77d8c50..c5f1eac0b61 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -9,13 +9,14 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_decimal_byte_parts::split_decimal; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -26,8 +27,11 @@ use crate::SchemeExt; /// Compression scheme for decimal arrays via byte-part decomposition. /// -/// Narrows the decimal to the smallest integer type, compresses the underlying primitive, and wraps -/// the result in a `DecimalBytePartsArray`. +/// Narrows the decimal to the smallest integer type and splits it into a signed most significant +/// part plus up to three unsigned 64-bit lower parts, each compressed as its own child. Values that +/// fit one signed part serialize under the frozen `vortex.decimal_byte_parts` format. Wider values +/// need lower parts and therefore the `vortex.decimal_byte_parts.v2` format, so they are only +/// split when the compressor may emit that format; otherwise they stay canonical. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DecimalScheme; @@ -41,14 +45,14 @@ impl Scheme for DecimalScheme { } fn produced_encodings(&self) -> Vec { - // This scheme only builds single-part arrays, which serialize under the frozen v1 ID. - // The in-memory ID is the v2 wire ID, which no edition permits yet. + // Single-part arrays serialize under the frozen v1 ID. Multi-part arrays need the v2 ID, + // which `compress` checks against the context before splitting wide values. vec![decimal_byte_parts_v1_id()] } - /// Children: primitive=0. + /// Children: msp=0, then up to three lower parts. fn num_children(&self) -> usize { - 1 + 4 } fn expected_compression_ratio( @@ -68,22 +72,28 @@ impl Scheme for DecimalScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): add support splitting i128/256 buffers into chunks of primitive values - // for compression. 2 for i128 and 4 for i256. let decimal = data.array().clone().execute::(exec_ctx)?; let decimal = narrowed_decimal(decimal); - let validity = decimal.validity()?; - let prim = match decimal.values_type() { - DecimalType::I8 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I16 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I32 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I64 => PrimitiveArray::new(decimal.buffer::(), validity), - _ => return Ok(decimal.into_array()), - }; - let compressed = - compressor.compress_child(&prim.into_array(), &compress_ctx, self.id(), 0, exec_ctx)?; + let values_type = decimal.values_type(); + let wide = matches!(values_type, DecimalType::I128 | DecimalType::I256); + if wide && !compress_ctx.allows_serialized_id(&decimal_byte_parts_v2_id()) { + // Only the v2 wire format carries lower parts. + return Ok(decimal.into_array()); + } - DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) + let parts = split_decimal(&decimal, exec_ctx)?; + let msp = compressor.compress_child(&parts.msp, &compress_ctx, self.id(), 0, exec_ctx)?; + let lower_parts = parts + .lower_parts + .iter() + .enumerate() + .map(|(idx, part)| { + compressor.compress_child(part, &compress_ctx, self.id(), 1 + idx, exec_ctx) + }) + .collect::>>()?; + + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal.decimal_dtype()) + .map(IntoArray::into_array) } } diff --git a/vortex-btrblocks/tests/decimal.rs b/vortex-btrblocks/tests/decimal.rs new file mode 100644 index 00000000000..5588127b8c5 --- /dev/null +++ b/vortex-btrblocks/tests/decimal.rs @@ -0,0 +1,134 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Wide decimals compress into multi-part byte-part arrays only when the compressor may emit the +//! `vortex.decimal_byte_parts.v2` wire format. + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::decimal::DecimalScheme; +use vortex_buffer::Buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + session +}); + +/// 128 ascending decimals. Wide ones exceed a single signed 64-bit part. +fn decimals(wide: bool) -> ArrayRef { + let base = if wide { 1i128 << 70 } else { 0 }; + DecimalArray::new( + (0..128i128).map(|i| base + i).collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array() +} + +fn serialized_id(array: &ArrayRef) -> VortexResult { + let serialized = SESSION + .array_serialize(array)? + .ok_or_else(|| vortex_err!("expected a serializable array"))?; + Ok(serialized.serialized_id) +} + +fn v1_only() -> HashSet { + HashSet::from([decimal_byte_parts_v1_id()]) +} + +fn v1_and_v2() -> HashSet { + HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]) +} + +/// The scheme always needs v1 for single-part arrays; v2 alone does not register it. +#[rstest] +#[case::neither(HashSet::new(), false)] +#[case::v1(v1_only(), true)] +#[case::v2_only(HashSet::from([decimal_byte_parts_v2_id()]), false)] +#[case::both(v1_and_v2(), true)] +fn decimal_scheme_needs_v1(#[case] allowed: HashSet, #[case] registered: bool) { + let compressor = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed) + .build(); + assert_eq!(compressor.has_scheme(DecimalScheme.id()), registered); +} + +/// Narrow decimals always compress under v1. Wide decimals split into a multi-part v2 array when +/// that format is permitted and otherwise stay canonical, including without any allowlist. +#[rstest] +fn wide_decimals_follow_permitted_formats( + #[values(None, Some(false), Some(true))] permit_v2: Option, + #[values(false, true)] wide: bool, +) -> VortexResult<()> { + let builder = BtrBlocksCompressorBuilder::default(); + let compressor = match permit_v2 { + None => builder.build(), + Some(false) => builder.retain_allowed_encodings(&v1_only()).build(), + Some(true) => builder.retain_allowed_encodings(&v1_and_v2()).build(), + }; + let array = decimals(wide); + let mut ctx = SESSION.create_execution_ctx(); + + let compressed = compressor.compress(&array, &mut ctx)?; + + let expect_byte_parts = !wide || permit_v2 == Some(true); + assert_eq!(compressed.is::(), expect_byte_parts); + if expect_byte_parts { + let expected = if wide { + decimal_byte_parts_v2_id() + } else { + decimal_byte_parts_v1_id() + }; + assert_eq!(serialized_id(&compressed)?, expected); + } + assert_arrays_eq!(array, compressed, &mut ctx); + Ok(()) +} + +/// The CUDA preset withholds v2 whether the writer allowlist is applied before or after it, while +/// single-part decimals still compress under v1. +#[rstest] +fn cuda_preset_keeps_wide_decimals_canonical( + #[values(false, true)] allowlist_first: bool, +) -> VortexResult<()> { + let builder = BtrBlocksCompressorBuilder::default(); + let builder = if allowlist_first { + builder + .retain_allowed_encodings(&v1_and_v2()) + .only_cuda_compatible() + } else { + builder + .only_cuda_compatible() + .retain_allowed_encodings(&v1_and_v2()) + }; + let compressor = builder.build(); + let mut ctx = SESSION.create_execution_ctx(); + + let wide = compressor.compress(&decimals(true), &mut ctx)?; + assert!(!wide.is::()); + + let narrow = compressor.compress(&decimals(false), &mut ctx)?; + assert_eq!(serialized_id(&narrow)?, decimal_byte_parts_v1_id()); + Ok(()) +} diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..706339d05a5 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -3,6 +3,8 @@ //! Core cascading compression flow. +use std::sync::Arc; + use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::CanonicalValidity; @@ -59,7 +61,9 @@ impl CascadingCompressor { let canonical = array.clone().execute::(exec_ctx)?.0; let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + let root_ctx = CompressorContext::new() + .with_allowed_serialized_ids(Arc::clone(&self.allowed_serialized_ids)); + let compressed = self.compress_canonical(compact, root_ctx, exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..30aa9eca65f 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,11 @@ mod sample; mod select; mod structural; +use std::sync::Arc; + +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,12 +51,18 @@ 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 compressor may emit. See [`Self::with_allowed_serialized_ids`]. + allowed_serialized_ids: Arc>, } impl CascadingCompressor { /// Creates a new compressor with the given schemes. /// /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. + /// The compressor may emit every serialized ID its schemes declare in + /// [`Scheme::produced_encodings`], and nothing else, until + /// [`with_allowed_serialized_ids`](Self::with_allowed_serialized_ids) says otherwise. pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { // Root exclusion: exclude IntDict from list/listview offsets (monotonically // increasing data where dictionary encoding is wasteful). @@ -59,12 +70,44 @@ impl CascadingCompressor { excluded: IntDictScheme.id(), children: ChildSelection::One(structural::root_list_children::OFFSETS), }]; + let allowed_serialized_ids = schemes + .iter() + .flat_map(|scheme| scheme.produced_encodings()) + .collect(); Self { schemes, root_exclusions, + allowed_serialized_ids: Arc::new(allowed_serialized_ids), } } + + /// Restricts the compressor to the serialized IDs in `allowed`. + /// + /// Schemes declaring an ID outside `allowed` are removed. The remaining schemes see `allowed` + /// through [`allows_serialized_id`](crate::scheme::CompressorContext::allows_serialized_id), + /// which lets a scheme emit an optional wire format only when the writer permits it. The file + /// writer passes the serialized IDs its enabled editions permit. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + self.schemes.retain(|scheme| { + scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + }); + self.allowed_serialized_ids = Arc::new(allowed); + self + } + + /// The serialized IDs this compressor may emit. + pub fn allowed_serialized_ids(&self) -> &HashSet { + &self.allowed_serialized_ids + } + + /// Whether a scheme with the given ID is registered. + pub fn has_scheme(&self, id: SchemeId) -> bool { + self.schemes.iter().any(|scheme| scheme.id() == id) + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 3a2a6281047..efb73b86dbd 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,9 +9,11 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::Dict; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; @@ -26,6 +28,7 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -845,3 +848,91 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } + +static OPTIONAL_WIRE_FORMAT_PERMITTED: Mutex> = Mutex::new(None); + +fn required_wire_id() -> ArrayId { + ArrayId::new_static("test.wire.required") +} + +fn optional_wire_id() -> ArrayId { + ArrayId::new_static("test.wire.optional") +} + +/// Declares one required wire format and records whether the context permits an optional one. +#[derive(Debug)] +struct OptionalWireFormatScheme; + +impl Scheme for OptionalWireFormatScheme { + fn scheme_name(&self) -> &'static str { + "test.optional_wire_format" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![required_wire_id()] + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + *OPTIONAL_WIRE_FORMAT_PERMITTED.lock() = + Some(compress_ctx.allows_serialized_id(&optional_wire_id())); + Ok(data.array().clone()) + } +} + +#[test] +fn compressor_permits_declared_ids_by_default() { + let compressor = CascadingCompressor::new(vec![&IntDictScheme, &OptionalWireFormatScheme]); + assert_eq!( + *compressor.allowed_serialized_ids(), + HashSet::from([Dict.id(), required_wire_id()]) + ); + assert!(compressor.has_scheme(IntDictScheme.id())); + assert!(compressor.has_scheme(OptionalWireFormatScheme.id())); +} + +#[test] +fn allowed_serialized_ids_drop_schemes_declaring_other_ids() { + let compressor = CascadingCompressor::new(vec![&IntDictScheme, &OptionalWireFormatScheme]) + .with_allowed_serialized_ids(HashSet::from([Dict.id()])); + assert!(compressor.has_scheme(IntDictScheme.id())); + assert!(!compressor.has_scheme(OptionalWireFormatScheme.id())); + + let none = CascadingCompressor::new(vec![&IntDictScheme]) + .with_allowed_serialized_ids(HashSet::new()); + assert!(!none.has_scheme(IntDictScheme.id())); +} + +#[test] +fn context_reports_optional_serialized_ids() -> VortexResult<()> { + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressor = CascadingCompressor::new(vec![&OptionalWireFormatScheme]); + compressor.compress(&array, &mut exec_ctx)?; + assert_eq!(*OPTIONAL_WIRE_FORMAT_PERMITTED.lock(), Some(false)); + + let compressor = compressor + .with_allowed_serialized_ids(HashSet::from([required_wire_id(), optional_wire_id()])); + compressor.compress(&array, &mut exec_ctx)?; + assert_eq!(*OPTIONAL_WIRE_FORMAT_PERMITTED.lock(), Some(true)); + Ok(()) +} diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..4e974d4f40f 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,8 +4,11 @@ //! Compression context for recursive compression. use std::fmt; +use std::sync::Arc; +use vortex_array::ArrayId; use vortex_error::VortexExpect; +use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -38,6 +41,9 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, + + /// The serialized IDs the compressor may emit, shared by every context in the cascade. + allowed_serialized_ids: Arc>, } impl CompressorContext { @@ -50,8 +56,15 @@ impl CompressorContext { allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), + allowed_serialized_ids: Arc::default(), } } + + /// Returns a context that permits the given serialized IDs. + pub(crate) fn with_allowed_serialized_ids(mut self, allowed: Arc>) -> Self { + self.allowed_serialized_ids = allowed; + self + } } #[cfg(test)] @@ -67,6 +80,18 @@ impl CompressorContext { self.is_sample } + /// Whether the compressor may emit arrays serialized under `id`. + /// + /// A scheme declares the serialized IDs it always needs in + /// [`produced_encodings`](crate::scheme::Scheme::produced_encodings) and is only registered + /// when all of them are permitted. A scheme with an optional wire format, such as a newer + /// version that only some values need, checks that format here before producing it and + /// otherwise falls back to a format it declared. The estimate and compression paths receive + /// the same context, so both can make the same decision. + pub fn allows_serialized_id(&self, id: &ArrayId) -> bool { + self.allowed_serialized_ids.contains(id) + } + /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..12c71aeabca 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,7 +124,7 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The serialized IDs this scheme itself may write into its compressed output. + /// The serialized IDs this scheme needs in order to write its compressed output. /// /// Every declared ID must be permitted for the scheme to be used. Cascaded children are /// compressed by other schemes, which declare their own IDs, so only arrays constructed @@ -132,7 +132,10 @@ pub trait Scheme: Debug + Send + Sync { /// merely rearranges do not need to be declared. /// /// For most encodings this is the in-memory encoding ID. An encoding with several wire - /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. + /// formats declares the wire IDs the scheme always writes, which may differ from its + /// in-memory ID. A wire format the scheme only needs for some inputs is not declared here; + /// the scheme checks it with [`CompressorContext::allows_serialized_id`] and falls back to a + /// declared format when it is not permitted. fn produced_encodings(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 98a4f4a7d30..c413923662e 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -6,16 +6,24 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::array_session; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::DecimalArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; 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_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; use vortex_edition::ComponentKind; use vortex_edition::Edition; use vortex_edition::EditionDeclaration; @@ -28,6 +36,7 @@ use vortex_edition::test_harness::validate_edition; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_file::OpenOptionsSessionExt; +use vortex_file::VortexWriteOptions; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; use vortex_io::session::RuntimeSession; @@ -649,3 +658,108 @@ async fn serialization_context_accepts_supported_compressor_output() -> VortexRe Ok(()) } + +/// Decimals whose values need more than one signed 64-bit part, so that byte-part compression +/// requires the `vortex.decimal_byte_parts.v2` wire format. +fn wide_decimals() -> ArrayRef { + DecimalArray::new( + (0..128i128) + .map(|i| (1i128 << 70) + i) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array() +} + +/// Write `array` with `options`, read the file back, and return the array it contains. +async fn round_trip( + session: &VortexSession, + options: VortexWriteOptions, + array: &ArrayRef, +) -> VortexResult { + let mut buffer = ByteBufferMut::empty(); + let stream = array.to_array_stream(); + options.write(&mut buffer, stream).await?; + session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await +} + +/// An explicit BtrBlocks strategy is not configured from the editions, so it may emit only the +/// wire formats its schemes declare. The opt-in v2 decimal format is not one of them: wide +/// decimals stay canonical and the write succeeds under the default editions. +#[tokio::test] +async fn explicit_default_strategy_keeps_wide_decimals_writable() -> VortexResult<()> { + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + assert!( + !session + .enabled_component_ids(ComponentKind::Array) + .contains(&decimal_byte_parts_v2_id()) + ); + let strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default()) + .build(); + let array = wide_decimals(); + + let read = round_trip( + &session, + session.write_options().with_strategy(strategy), + &array, + ) + .await?; + + assert!( + !read + .depth_first_traversal() + .any(|child| child.is::()) + ); + Ok(()) +} + +/// The writer's compressor splits wide decimals only when the v2 wire format is permitted: by an +/// enabled edition that includes it, or by disabling editions so every registered format is +/// allowed. +#[tokio::test] +async fn writer_splits_wide_decimals_when_v2_is_permitted() -> VortexResult<()> { + use crate::VortexSessionDefault; + + const EDITION: EditionId = EditionId::new("decimal-v2-test", 2026, 9, 0); + + let array = wide_decimals(); + for (enable_v2, disable_editions) in [(false, false), (true, false), (false, true)] { + let session = VortexSession::default(); + if enable_v2 { + let editions = session.editions(); + editions.declare_edition(Edition { + id: EDITION, + min_library_version: None, + })?; + editions.declare_inclusion(EditionInclusion::array( + &decimal_byte_parts_v2_id(), + EDITION, + ))?; + session.enable_edition(EDITION)?; + } + let mut options = session.write_options(); + if disable_editions { + options = options.disable_editions(); + } + + let read = round_trip(&session, options, &array).await?; + + assert_eq!( + read.depth_first_traversal() + .any(|child| child.is::()), + enable_v2 || disable_editions, + "enable_v2={enable_v2} disable_editions={disable_editions}" + ); + } + Ok(()) +}