From 2add9d701b6ee309153075cf46293e33165041d7 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 14:55:07 -0400 Subject: [PATCH 1/6] Configure compressor schemes from serialized ID permissions Resolve allowed and forbidden serialized IDs when building the compressor, select compatible Decimal modes, and compress wide decimal parts independently. Default to Decimal v2 while keeping the CUDA preset restricted to v1. Signed-off-by: "Matt Katz" --- vortex-btrblocks/src/builder.rs | 136 ++++++++++-- vortex-btrblocks/src/schemes/decimal.rs | 110 +++++++-- vortex-btrblocks/tests/decimal_config.rs | 270 +++++++++++++++++++++++ 3 files changed, 477 insertions(+), 39 deletions(-) create mode 100644 vortex-btrblocks/tests/decimal_config.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 3bcda909227..03a933b4d76 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; @@ -18,7 +19,7 @@ use crate::schemes::integer; use crate::schemes::string; use crate::schemes::temporal; -/// All available compression schemes. +/// All default compression schemes, including Decimal v2. /// /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. @@ -60,7 +61,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ &binary::BinaryDictScheme, &binary::VarBinScheme, // Decimal schemes. - &decimal::DecimalScheme, + &decimal::DecimalScheme::v2(), // Temporal schemes. &temporal::TemporalScheme, ]; @@ -79,6 +80,10 @@ 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. /// +/// [`Self::retain_allowed_encodings`] restricts serialized IDs. During [`Self::build`], these +/// restrictions select the compatible Decimal mode and filter all registered schemes. Without +/// an allowlist, registered modes are preserved, including Decimal v2 in the defaults. +/// /// # Examples /// /// ```rust @@ -96,12 +101,16 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + allowed_serialized_ids: Option>, + forbidden_serialized_ids: HashSet, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + allowed_serialized_ids: None, + forbidden_serialized_ids: HashSet::new(), } } } @@ -113,6 +122,8 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_serialized_ids: None, + forbidden_serialized_ids: HashSet::new(), } } @@ -165,6 +176,8 @@ impl BtrBlocksCompressorBuilder { /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, /// but belongs to the opt-in `zstd` edition, so callers filter the two through /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). + /// Decimal schemes are restricted to v1 during build, regardless of when the allowlist or + /// Decimal scheme is supplied, because CUDA does not support lower decimal parts. /// /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. @@ -189,7 +202,10 @@ 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); + let mut builder = self.exclude_schemes(excluded); + builder + .forbidden_serialized_ids + .insert(decimal_byte_parts_v2_id()); #[cfg(feature = "zstd")] let builder = builder @@ -206,24 +222,57 @@ impl BtrBlocksCompressorBuilder { self } - /// Retains only schemes whose produced serialized IDs all belong to `allowed`. + /// Restricts schemes to those whose produced serialized IDs all belong to `allowed`. /// /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. + /// permit. Repeated calls intersect the allowed sets; an empty set permits no serialized IDs. + /// + /// Configuration and filtering are deferred until [`Self::build`], including for schemes + /// registered after this call. Registered Decimal schemes use v2 when both decimal IDs are + /// allowed, or v1 when only v1 is allowed. If v1 is not allowed, Decimal is removed, since + /// single-part arrays always serialize as v1. The CUDA preset restricts Decimal to v1. 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_serialized_ids { + Some(current) => current.retain(|id| allowed.contains(id)), + None => self.allowed_serialized_ids = Some(allowed.clone()), + } self } /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + BtrBlocksCompressor(CascadingCompressor::new(self.configured_schemes())) + } + + fn configured_schemes(mut self) -> Vec<&'static dyn Scheme> { + let allowed = self.allowed_serialized_ids.as_ref(); + let forbidden = &self.forbidden_serialized_ids; + let is_allowed = |id: &ArrayId| { + !forbidden.contains(id) && allowed.is_none_or(|allowed| allowed.contains(id)) + }; + if allowed.is_some() || forbidden.contains(&decimal_byte_parts_v2_id()) { + let use_v2 = is_allowed(&decimal_byte_parts_v2_id()); + for scheme in &mut self.schemes { + if scheme.id() == decimal::DecimalScheme::v2().id() { + *scheme = if use_v2 { + &decimal::DecimalScheme::v2() + } else { + &decimal::DecimalScheme::v1() + }; + } + } + } + if allowed.is_some() || !forbidden.is_empty() { + self.schemes + .retain(|scheme| scheme.produced_encodings().iter().all(&is_allowed)); + } + self.schemes } } #[cfg(test)] mod tests { + use rstest::rstest; use vortex_array::VTable; use vortex_fastlanes::FoR; @@ -244,12 +293,16 @@ mod tests { #[test] fn retain_allowed_encodings_filters_schemes() { 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 schemes = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed) + .configured_schemes(); + assert_eq!(schemes.len(), 1); + assert_eq!(schemes[0].id(), integer::FoRScheme.id()); - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); + let none = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&HashSet::new()) + .configured_schemes(); + assert!(none.is_empty()); } #[test] @@ -258,8 +311,61 @@ mod tests { .iter() .flat_map(|scheme| scheme.produced_encodings()) .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); + let schemes = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed) + .configured_schemes(); + assert_eq!(schemes, ALL_SCHEMES); + } + + #[rstest] + fn restrictions_apply_to_later_registrations(#[values(false, true)] register_later: bool) { + let mut builder = BtrBlocksCompressorBuilder::empty(); + if !register_later { + builder = builder.with_new_scheme(&integer::FoRScheme); + } + builder = builder.retain_allowed_encodings(&HashSet::new()); + if register_later { + builder = builder.with_new_scheme(&integer::FoRScheme); + } + let allowed = integer::FoRScheme.produced_encodings().into_iter().collect(); + assert!( + builder + .retain_allowed_encodings(&allowed) + .configured_schemes() + .is_empty() + ); + } + + #[rstest] + fn forbidden_ids_filter_schemes( + #[values(false, true)] with_allowlist: bool, + #[values(false, true)] register_later: bool, + ) { + let mut builder = BtrBlocksCompressorBuilder::empty(); + if !register_later { + builder = builder.with_new_scheme(&integer::FoRScheme); + } + builder.forbidden_serialized_ids.insert(FoR.id()); + if with_allowlist { + builder = builder.retain_allowed_encodings(&HashSet::from([FoR.id()])); + } + if register_later { + builder = builder.with_new_scheme(&integer::FoRScheme); + } + assert!(builder.configured_schemes().is_empty()); + } + + #[test] + fn unrelated_forbidden_ids_preserve_explicit_decimal_mode() { + let mut builder = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&decimal::DecimalScheme::v1()); + builder.forbidden_serialized_ids.insert(FoR.id()); + let schemes = builder.configured_schemes(); + assert_eq!(schemes.len(), 1); + assert_eq!( + schemes[0].produced_encodings(), + decimal::DecimalScheme::v1().produced_encodings() + ); } #[test] diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index f77a77d8c50..ac0526332ba 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -9,13 +9,15 @@ 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::DecimalBytePartsSlots; 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; @@ -24,12 +26,49 @@ use crate::CompressorContext; use crate::Scheme; use crate::SchemeExt; +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum DecimalSchemeMode { + V1, + V2, +} + /// 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 compresses its byte parts independently. +/// The v1 mode leaves values wider than `i64` canonical; v2 splits them into a signed most +/// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1 +/// in either mode, while arrays with lower parts serialize as v2. +/// +/// The default uses v2. A builder with allowed serialized IDs selects the latest permitted mode, +/// including for explicitly registered Decimal schemes. The CUDA preset restricts the mode to v1. #[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct DecimalScheme; +pub struct DecimalScheme { + mode: DecimalSchemeMode, +} + +impl DecimalScheme { + /// Creates a decimal scheme configured for v1, disallowing splitting of wide decimals. + /// + /// Values that remain wider than `i64` after narrowing stay canonical. + pub const fn v1() -> Self { + Self { + mode: DecimalSchemeMode::V1, + } + } + + /// Creates a decimal scheme configured for v2, allowing splitting of wide decimals. + pub const fn v2() -> Self { + Self { + mode: DecimalSchemeMode::V2, + } + } +} + +impl Default for DecimalScheme { + fn default() -> Self { + Self::v2() + } +} impl Scheme for DecimalScheme { fn scheme_name(&self) -> &'static str { @@ -41,14 +80,21 @@ 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. - vec![decimal_byte_parts_v1_id()] + let mut ids = vec![decimal_byte_parts_v1_id()]; + + if matches!(self.mode, DecimalSchemeMode::V2) { + ids.push(decimal_byte_parts_v2_id()); + } + + ids } - /// Children: primitive=0. + /// Children: msp=0, then up to three lower parts in v2 mode. fn num_children(&self) -> usize { - 1 + match self.mode { + DecimalSchemeMode::V1 => 1, + DecimalSchemeMode::V2 => 4, + } } fn expected_compression_ratio( @@ -68,22 +114,38 @@ 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)?; - - DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) + if self.mode == DecimalSchemeMode::V1 + && matches!(decimal.values_type(), DecimalType::I128 | DecimalType::I256) + { + return Ok(decimal.into_array()); + } + + let parts = split_decimal(&decimal, exec_ctx)?; + let msp = compressor.compress_child( + &parts.msp, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::MSP, + exec_ctx, + )?; + let lower_parts = parts + .lower_parts + .iter() + .enumerate() + .map(|(idx, part)| { + compressor.compress_child( + part, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + 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_config.rs b/vortex-btrblocks/tests/decimal_config.rs new file mode 100644 index 00000000000..2bcf0ab6138 --- /dev/null +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -0,0 +1,270 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal mode selection, serialized permissions, and compression of wide decimal parts. + +#![cfg(test)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::serde::SerializeOptions; +use vortex_array::serde::SerializedArray; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::Scheme; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::decimal::DecimalScheme; +use vortex_btrblocks::schemes::integer::BitPackingScheme; +use vortex_btrblocks::schemes::integer::FoRScheme; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +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_session::registry::ReadContext; +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 +}); + +fn decimal_array(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 assert_decimal_output( + builder: BtrBlocksCompressorBuilder, + wide: bool, + expected_id: Option, +) -> VortexResult<()> { + let array = decimal_array(wide); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = builder.build().compress(&array, &mut ctx)?; + if let Some(expected_id) = expected_id { + assert!(compressed.is::()); + let serialized = SESSION + .array_serialize(&compressed)? + .ok_or_else(|| vortex_err!("expected serializable decimal byte parts"))?; + assert_eq!(serialized.serialized_id, expected_id); + } else { + assert!(compressed.is::()); + } + assert_arrays_eq!(array, compressed, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::unrestricted(None, Some(true))] +#[case::neither(Some(vec![]), None)] +#[case::v1(Some(vec![decimal_byte_parts_v1_id()]), Some(false))] +#[case::v2_only(Some(vec![decimal_byte_parts_v2_id()]), None)] +#[case::both(Some(vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]), Some(true))] +fn decimal_mode_follows_permissions( + #[case] ids: Option>, + #[case] v2: Option, + #[values(false, true)] wide: bool, +) -> VortexResult<()> { + let mut builder = BtrBlocksCompressorBuilder::default(); + if let Some(ids) = ids { + builder = builder.retain_allowed_encodings(&ids.into_iter().collect()); + } + let expected = match v2 { + Some(true) if wide => Some(decimal_byte_parts_v2_id()), + Some(_) if !wide => Some(decimal_byte_parts_v1_id()), + _ => None, + }; + assert_decimal_output(builder, wide, expected) +} + +#[rstest] +#[case::unrestricted(None)] +#[case::v1(Some(false))] +#[case::both(Some(true))] +fn explicit_decimal_modes( + #[case] allowed_v2: Option, + #[values(false, true)] initial_v2: bool, + #[values(false, true)] register_later: bool, +) -> VortexResult<()> { + let scheme: &'static dyn Scheme = if initial_v2 { + &DecimalScheme::v2() + } else { + &DecimalScheme::v1() + }; + let mut builder = BtrBlocksCompressorBuilder::empty(); + if !register_later { + builder = builder.with_new_scheme(scheme); + } + if let Some(v2) = allowed_v2 { + let mut allowed = HashSet::from([decimal_byte_parts_v1_id()]); + if v2 { + allowed.insert(decimal_byte_parts_v2_id()); + } + builder = builder.retain_allowed_encodings(&allowed); + } + if register_later { + builder = builder.with_new_scheme(scheme); + } + assert_decimal_output( + builder, + true, + allowed_v2 + .unwrap_or(initial_v2) + .then(decimal_byte_parts_v2_id), + ) +} + +#[rstest] +fn decimal_permissions_intersect(#[values(false, true)] restrictive_first: bool) -> VortexResult<()> { + let v1 = HashSet::from([decimal_byte_parts_v1_id()]); + let both = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + let (first, second) = if restrictive_first { + (&v1, &both) + } else { + (&both, &v1) + }; + let builder = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(first) + .retain_allowed_encodings(second); + assert_decimal_output(builder.clone(), false, Some(decimal_byte_parts_v1_id()))?; + assert_decimal_output(builder, true, None) +} + +#[test] +fn permissions_do_not_restore_excluded_decimal() -> VortexResult<()> { + let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + let builder = BtrBlocksCompressorBuilder::default() + .exclude_schemes([DecimalScheme::default().id()]) + .retain_allowed_encodings(&allowed); + assert_decimal_output(builder, false, None) +} + +#[rstest] +#[case::unrestricted(None)] +#[case::permissions_first(Some(true))] +#[case::permissions_last(Some(false))] +fn cuda_keeps_decimal_v1( + #[case] permissions_first: Option, + #[values(false, true)] register_later: bool, + #[values(false, true)] wide: bool, +) -> VortexResult<()> { + let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + let mut builder = BtrBlocksCompressorBuilder::default(); + if register_later { + builder = builder.exclude_schemes([DecimalScheme::default().id()]); + } + if permissions_first == Some(true) { + builder = builder.retain_allowed_encodings(&allowed); + } + builder = builder.only_cuda_compatible(); + if register_later { + builder = builder.with_new_scheme(&DecimalScheme::v2()); + } + if permissions_first.is_some() { + builder = builder.retain_allowed_encodings(&allowed); + } + assert_decimal_output(builder, wide, (!wide).then(decimal_byte_parts_v1_id)) +} + +#[rstest] +#[case::i128(false, 1)] +#[case::i256(true, 3)] +fn wide_decimal_parts_roundtrip( + #[case] use_i256: bool, + #[case] lower_part_count: usize, + #[values(false, true)] negative: bool, + #[values(false, true)] nullable: bool, + #[values(false, true)] compress_children: bool, +) -> VortexResult<()> { + let validity = if nullable { + Validity::from_iter((0..2048).map(|i| i % 7 != 0)) + } else { + Validity::NonNullable + }; + let array = if use_i256 { + let values = (1..=2048u32) + .map(|i| { + let msp = if negative { + -i128::from(i) + } else { + i128::from(i) + }; + i256::from_parts( + (u128::from(i) << 64) | u128::from(i * 131 + 17), + (msp << 64) | i128::from(i * 3 + 1), + ) + }) + .collect::>(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) + } else { + let values = (1..=2048i128) + .map(|i| { + let msp = if negative { -i } else { i }; + (msp << 70) + i * 131 + 17 + }) + .collect::>(); + DecimalArray::new(values, DecimalDType::new(38, 2), validity) + } + .into_array(); + let mut allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); + if compress_children { + allowed.extend(FoRScheme.produced_encodings()); + allowed.extend(BitPackingScheme.produced_encodings()); + } + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&DecimalScheme::v2()) + .with_new_scheme(&FoRScheme) + .with_new_scheme(&BitPackingScheme) + .retain_allowed_encodings(&allowed) + .build(); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + // Every part varies, so without child compression splitting alone cannot save space. + assert_eq!(compressed.is::(), compress_children); + if compress_children { + let parts = compressed + .as_opt::() + .ok_or_else(|| vortex_err!("expected decimal byte parts"))?; + assert_eq!(parts.lower_parts().len(), lower_part_count); + assert!(!parts.msp().is_canonical()); + assert!(parts.lower_parts().iter().all(|part| !part.is_canonical())); + } + + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + for buffer in compressed.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + array.dtype(), + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) +} From ee65f48da9c2101763d7e6165fd0e02db08616d9 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 16:38:18 -0400 Subject: [PATCH 2/6] Configure compression schemes through serialized ID permissions Replace produced_encodings with Scheme::configure and store owned AllowedSerializedIds in the BtrBlocks builder. Configure schemes in one pass, select decimal v2 whenever permitted, and preserve CUDA restrictions. Signed-off-by: "Matt Katz" --- vortex-btrblocks/src/builder.rs | 76 ++++++----------- vortex-btrblocks/src/lib.rs | 1 + vortex-btrblocks/src/schemes/binary/varbin.rs | 11 ++- vortex-btrblocks/src/schemes/binary/zstd.rs | 11 ++- .../src/schemes/binary/zstd_buffers.rs | 11 ++- vortex-btrblocks/src/schemes/decimal.rs | 18 ++-- vortex-btrblocks/src/schemes/float/alp.rs | 16 ++-- vortex-btrblocks/src/schemes/float/alprd.rs | 11 ++- vortex-btrblocks/src/schemes/float/pco.rs | 11 ++- vortex-btrblocks/src/schemes/float/rle.rs | 11 ++- vortex-btrblocks/src/schemes/float/sparse.rs | 11 ++- .../src/schemes/integer/bitpacking.rs | 16 ++-- vortex-btrblocks/src/schemes/integer/delta.rs | 11 ++- vortex-btrblocks/src/schemes/integer/for_.rs | 11 ++- vortex-btrblocks/src/schemes/integer/pco.rs | 11 ++- vortex-btrblocks/src/schemes/integer/rle.rs | 11 ++- .../src/schemes/integer/runend.rs | 11 ++- .../src/schemes/integer/sequence.rs | 11 ++- .../src/schemes/integer/sparse.rs | 11 ++- .../src/schemes/integer/zigzag.rs | 11 ++- vortex-btrblocks/src/schemes/string/fsst.rs | 11 ++- vortex-btrblocks/src/schemes/string/onpair.rs | 11 ++- vortex-btrblocks/src/schemes/string/sparse.rs | 11 ++- vortex-btrblocks/src/schemes/string/zstd.rs | 11 ++- .../src/schemes/string/zstd_buffers.rs | 11 ++- vortex-btrblocks/src/schemes/temporal.rs | 11 ++- vortex-btrblocks/tests/decimal_config.rs | 13 +-- vortex-btrblocks/tests/scheme_config.rs | 85 +++++++++++++++++++ vortex-compressor/src/builtins/dict/binary.rs | 11 ++- vortex-compressor/src/builtins/dict/float.rs | 11 ++- .../src/builtins/dict/integer.rs | 11 ++- vortex-compressor/src/builtins/dict/string.rs | 11 ++- vortex-compressor/src/compressor/tests.rs | 65 +++++++++----- .../src/scheme/allowed_serialized_ids.rs | 59 +++++++++++++ vortex-compressor/src/scheme/mod.rs | 18 ++-- 35 files changed, 458 insertions(+), 184 deletions(-) create mode 100644 vortex-btrblocks/tests/scheme_config.rs create mode 100644 vortex-compressor/src/scheme/allowed_serialized_ids.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 03a933b4d76..cc748fdcb95 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -7,6 +7,7 @@ use vortex_array::ArrayId; use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; use vortex_utils::aliases::hash_set::HashSet; +use crate::AllowedSerializedIds; use crate::BtrBlocksCompressor; use crate::CascadingCompressor; use crate::Scheme; @@ -81,8 +82,8 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); /// `zstd` feature is enabled. /// /// [`Self::retain_allowed_encodings`] restricts serialized IDs. During [`Self::build`], these -/// restrictions select the compatible Decimal mode and filter all registered schemes. Without -/// an allowlist, registered modes are preserved, including Decimal v2 in the defaults. +/// restrictions select the latest compatible Decimal mode and filter all registered schemes. +/// Decimal uses v2 whenever both serialized IDs are permitted, including without an allowlist. /// /// # Examples /// @@ -101,16 +102,14 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, - allowed_serialized_ids: Option>, - forbidden_serialized_ids: HashSet, + allowed_serialized_ids: AllowedSerializedIds, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), - allowed_serialized_ids: None, - forbidden_serialized_ids: HashSet::new(), + allowed_serialized_ids: AllowedSerializedIds::default(), } } } @@ -122,8 +121,7 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), - allowed_serialized_ids: None, - forbidden_serialized_ids: HashSet::new(), + allowed_serialized_ids: AllowedSerializedIds::default(), } } @@ -204,8 +202,8 @@ impl BtrBlocksCompressorBuilder { excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]); let mut builder = self.exclude_schemes(excluded); builder - .forbidden_serialized_ids - .insert(decimal_byte_parts_v2_id()); + .allowed_serialized_ids + .exclude(decimal_byte_parts_v2_id()); #[cfg(feature = "zstd")] let builder = builder @@ -232,10 +230,7 @@ impl BtrBlocksCompressorBuilder { /// allowed, or v1 when only v1 is allowed. If v1 is not allowed, Decimal is removed, since /// single-part arrays always serialize as v1. The CUDA preset restricts Decimal to v1. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - match &mut self.allowed_serialized_ids { - Some(current) => current.retain(|id| allowed.contains(id)), - None => self.allowed_serialized_ids = Some(allowed.clone()), - } + self.allowed_serialized_ids.restrict_to(allowed); self } @@ -244,29 +239,14 @@ impl BtrBlocksCompressorBuilder { BtrBlocksCompressor(CascadingCompressor::new(self.configured_schemes())) } - fn configured_schemes(mut self) -> Vec<&'static dyn Scheme> { - let allowed = self.allowed_serialized_ids.as_ref(); - let forbidden = &self.forbidden_serialized_ids; - let is_allowed = |id: &ArrayId| { - !forbidden.contains(id) && allowed.is_none_or(|allowed| allowed.contains(id)) - }; - if allowed.is_some() || forbidden.contains(&decimal_byte_parts_v2_id()) { - let use_v2 = is_allowed(&decimal_byte_parts_v2_id()); - for scheme in &mut self.schemes { - if scheme.id() == decimal::DecimalScheme::v2().id() { - *scheme = if use_v2 { - &decimal::DecimalScheme::v2() - } else { - &decimal::DecimalScheme::v1() - }; - } + fn configured_schemes(self) -> Vec<&'static dyn Scheme> { + let mut final_schemes = Vec::with_capacity(self.schemes.len()); + for scheme in self.schemes { + if let Some(scheme) = scheme.configure(&self.allowed_serialized_ids) { + final_schemes.push(scheme); } } - if allowed.is_some() || !forbidden.is_empty() { - self.schemes - .retain(|scheme| scheme.produced_encodings().iter().all(&is_allowed)); - } - self.schemes + final_schemes } } @@ -306,14 +286,8 @@ mod tests { } #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let schemes = BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed) - .configured_schemes(); + fn unrestricted_configuration_preserves_scheme_order() { + let schemes = BtrBlocksCompressorBuilder::default().configured_schemes(); assert_eq!(schemes, ALL_SCHEMES); } @@ -327,7 +301,7 @@ mod tests { if register_later { builder = builder.with_new_scheme(&integer::FoRScheme); } - let allowed = integer::FoRScheme.produced_encodings().into_iter().collect(); + let allowed = HashSet::from([FoR.id()]); assert!( builder .retain_allowed_encodings(&allowed) @@ -345,7 +319,7 @@ mod tests { if !register_later { builder = builder.with_new_scheme(&integer::FoRScheme); } - builder.forbidden_serialized_ids.insert(FoR.id()); + builder.allowed_serialized_ids.exclude(FoR.id()); if with_allowlist { builder = builder.retain_allowed_encodings(&HashSet::from([FoR.id()])); } @@ -356,15 +330,15 @@ mod tests { } #[test] - fn unrelated_forbidden_ids_preserve_explicit_decimal_mode() { - let mut builder = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&decimal::DecimalScheme::v1()); - builder.forbidden_serialized_ids.insert(FoR.id()); + fn unrelated_forbidden_ids_allow_decimal_v2() { + let mut builder = + BtrBlocksCompressorBuilder::empty().with_new_scheme(&decimal::DecimalScheme::v1()); + builder.allowed_serialized_ids.exclude(FoR.id()); let schemes = builder.configured_schemes(); assert_eq!(schemes.len(), 1); assert_eq!( - schemes[0].produced_encodings(), - decimal::DecimalScheme::v1().produced_encodings() + schemes[0].num_children(), + decimal::DecimalScheme::v2().num_children() ); } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 2e8ae484f90..680bfd57676 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/vortex-btrblocks/src/lib.rs @@ -82,6 +82,7 @@ pub use builder::DELTA_SCHEME; pub use canonical_compressor::BtrBlocksCompressor; pub use schemes::patches::compress_patches; pub use vortex_compressor::CascadingCompressor; +pub use vortex_compressor::scheme::AllowedSerializedIds; pub use vortex_compressor::scheme::CompressorContext; pub use vortex_compressor::scheme::MAX_CASCADE; pub use vortex_compressor::scheme::Scheme; diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index 849402a6493..f06deac7f68 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -9,7 +9,6 @@ //! cascading compressor can then compress with the ordinary integer schemes. For fixed-width //! values the offsets are a constant-stride sequence and collapse to nothing. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -21,6 +20,7 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::builders::VarBinBuilder; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::SchemeExt; @@ -44,8 +44,13 @@ impl Scheme for VarBinScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { - vec![VarBin.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&VarBin.id()) + .then_some(self) } fn num_children(&self) -> usize { diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index d652e344db2..5080cc26a20 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -3,12 +3,12 @@ //! Zstd compression for binary arrays. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,8 +31,13 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::Zstd.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_zstd::Zstd.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 3f06d65b061..751559b059d 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -3,12 +3,12 @@ //! Zstd buffer-level binary compression preserving array layout for GPU decompression. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,8 +31,13 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::ZstdBuffers.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_zstd::ZstdBuffers.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index ac0526332ba..d72ebf75006 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -3,7 +3,6 @@ //! Decimal compression scheme using byte-part decomposition. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; @@ -39,7 +39,7 @@ enum DecimalSchemeMode { /// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1 /// in either mode, while arrays with lower parts serialize as v2. /// -/// The default uses v2. A builder with allowed serialized IDs selects the latest permitted mode, +/// The default uses v2. The builder always selects the latest permitted mode, /// including for explicitly registered Decimal schemes. The CUDA preset restricts the mode to v1. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DecimalScheme { @@ -79,14 +79,14 @@ impl Scheme for DecimalScheme { matches!(canonical, Canonical::Decimal(_)) } - fn produced_encodings(&self) -> Vec { - let mut ids = vec![decimal_byte_parts_v1_id()]; - - if matches!(self.mode, DecimalSchemeMode::V2) { - ids.push(decimal_byte_parts_v2_id()); + fn configure(&self, allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme> { + if !allowed_serialized_ids.contains(&decimal_byte_parts_v1_id()) { + return None; } - - ids + if !allowed_serialized_ids.contains(&decimal_byte_parts_v2_id()) { + return Some(&Self::v1()); + } + Some(&Self::v2()) } /// Children: msp=0, then up to three lower parts in v2 mode. diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index f9fc7066bf4..ed624a8d0ce 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -7,7 +7,6 @@ use vortex_alp::ALP; use vortex_alp::ALPArrayExt; use vortex_alp::ALPArraySlotsExt; use vortex_alp::alp_encode; -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -17,6 +16,7 @@ use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -42,12 +42,14 @@ impl Scheme for ALPScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { - let mut encodings = vec![ALP.id()]; - if use_experimental_patches() { - encodings.push(Patched.id()); - } - encodings + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + (allowed_serialized_ids.contains(&ALP.id()) + && (!use_experimental_patches() + || allowed_serialized_ids.contains(&Patched.id()))) + .then_some(self) } /// Children: encoded_ints=0. diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 09c0ee85b0d..15b10a575c0 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -7,7 +7,6 @@ use vortex_alp::ALPRDArrayExt; use vortex_alp::ALPRDArrayOwnedExt; use vortex_alp::RDEncoder; use vortex_alp::RDEncoderExt; -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -15,6 +14,7 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -40,8 +40,13 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { - vec![vortex_alp::ALPRD.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_alp::ALPRD.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index 416668c2fd0..1100eed3958 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -3,12 +3,12 @@ //! Pco (pcodec) float compression. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,8 +31,13 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { - vec![vortex_pco::Pco.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_pco::Pco.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 71158b9dc3b..3645f93c96f 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -3,11 +3,11 @@ //! Run-length float encoding. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; @@ -38,8 +38,13 @@ impl Scheme for FloatRLEScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { - vec![RLE.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&RLE.id()) + .then_some(self) } /// Children: values=0, indices=1, offsets=2. diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 3d9c25b18e4..0cd1ea8a4d5 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -3,7 +3,6 @@ //! Sparse encoding for null-dominated float arrays. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -41,8 +41,13 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { - vec![Sparse.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Sparse.id()) + .then_some(self) } /// Children: indices=0. diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 5ac7d0e4078..95140f13820 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -3,7 +3,6 @@ //! BitPacking integer encoding. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -12,6 +11,7 @@ use vortex_array::VTable; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -40,12 +40,14 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - let mut encodings = vec![BitPacked.id()]; - if use_experimental_patches() { - encodings.push(Patched.id()); - } - encodings + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + (allowed_serialized_ids.contains(&BitPacked.id()) + && (!use_experimental_patches() + || allowed_serialized_ids.contains(&Patched.id()))) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 46b2f1e302e..752172ee21c 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -3,7 +3,6 @@ //! FastLanes Delta integer encoding. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -14,6 +13,7 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -97,8 +97,13 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![Delta.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Delta.id()) + .then_some(self) } fn num_children(&self) -> usize { diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 476a0dec282..ef776bdbcef 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -3,7 +3,6 @@ //! Frame of Reference integer encoding. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -14,6 +13,7 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -44,8 +44,13 @@ impl Scheme for FoRScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![FoR.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&FoR.id()) + .then_some(self) } /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op. diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 675a112d44f..2989b4ff159 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -3,12 +3,12 @@ //! Pco (pcodec) integer compression. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -32,8 +32,13 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![vortex_pco::Pco.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_pco::Pco.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 86a40d0b36a..9bbc3fe53ae 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -3,7 +3,6 @@ //! Run-length integer encoding and shared RLE compression helpers. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; @@ -106,8 +106,13 @@ impl Scheme for IntRLEScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![RLE.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&RLE.id()) + .then_some(self) } /// Children: values=0, indices=1, offsets=2. diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index 6a97f7ec37d..f6334008352 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -3,7 +3,6 @@ //! Run-end integer encoding. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -14,6 +13,7 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -48,8 +48,13 @@ impl Scheme for RunEndScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![RunEnd.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&RunEnd.id()) + .then_some(self) } /// Children: values=0, ends=1. diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index edcefb99fc2..3abc273bb1f 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -3,7 +3,6 @@ //! Sequence integer encoding for sequential patterns. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -12,6 +11,7 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -43,8 +43,13 @@ impl Scheme for SequenceScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![Sequence.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Sequence.id()) + .then_some(self) } /// Sequence encoding on dictionary codes just adds a layer of indirection without compressing diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 429ff5c1a31..251862007a8 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -3,7 +3,6 @@ //! Sparse integer encoding for single-value-dominated arrays. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -15,6 +14,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::scalar::Scalar; use vortex_compressor::builtins::IntDictScheme; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -46,8 +46,13 @@ impl Scheme for SparseScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![Sparse.id(), Constant.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + (allowed_serialized_ids.contains(&Sparse.id()) + && allowed_serialized_ids.contains(&Constant.id())) + .then_some(self) } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 0e4be01845a..5d14a20e620 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -3,7 +3,6 @@ //! ZigZag integer encoding for signed integers. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -14,6 +13,7 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -46,8 +46,13 @@ impl Scheme for ZigZagScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![ZigZag.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&ZigZag.id()) + .then_some(self) } /// Children: encoded=0. diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index fd3fd28696a..9da2b95e3fd 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -5,7 +5,6 @@ use std::sync::Arc; -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -16,6 +15,7 @@ use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -50,8 +50,13 @@ impl Scheme for FSSTScheme { canonical.dtype().is_utf8() || canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { - vec![FSST.id(), VarBin.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + (allowed_serialized_ids.contains(&FSST.id()) + && allowed_serialized_ids.contains(&VarBin.id())) + .then_some(self) } /// Children: lengths=0, code_offsets=1. diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 06a7ff97f89..229b80e8d58 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -3,7 +3,6 @@ //! OnPair short-string compression (dict-12). -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::SchemeId; @@ -48,8 +48,13 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { - vec![OnPair.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&OnPair.id()) + .then_some(self) } /// 4 primitive slot children flow through the cascading compressor: diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 8620c366f77..95e9dddac0f 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -3,7 +3,6 @@ //! Sparse encoding for null-dominated string arrays. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,6 +10,7 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -42,8 +42,13 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { - vec![Sparse.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Sparse.id()) + .then_some(self) } /// Children: indices=0. diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index 84e8860d626..c066470d276 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -3,12 +3,12 @@ //! Zstd string compression without dictionaries (nvCOMP compatible). -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,8 +31,13 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::Zstd.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_zstd::Zstd.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index cf691c70fcb..d1316804aa6 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -3,12 +3,12 @@ //! Zstd buffer-level string compression preserving array layout for GPU decompression. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,8 +31,13 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { - vec![vortex_zstd::ZstdBuffers.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&vortex_zstd::ZstdBuffers.id()) + .then_some(self) } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 79748b69450..ed660b8a957 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -3,7 +3,6 @@ //! Temporal compression scheme using datetime-part decomposition. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -17,6 +16,7 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::extension::Matcher; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; +use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_datetime_parts::DateTimeParts; @@ -55,8 +55,13 @@ impl Scheme for TemporalScheme { ) } - fn produced_encodings(&self) -> Vec { - vec![DateTimeParts.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&DateTimeParts.id()) + .then_some(self) } /// Children: days=0, seconds=1, subseconds=2. diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs index 2bcf0ab6138..1651175b574 100644 --- a/vortex-btrblocks/tests/decimal_config.rs +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -12,9 +12,11 @@ use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Decimal; use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Patched; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::i256; @@ -36,6 +38,8 @@ 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_fastlanes::BitPacked; +use vortex_fastlanes::FoR; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use vortex_utils::aliases::hash_set::HashSet; @@ -105,7 +109,7 @@ fn decimal_mode_follows_permissions( #[case::unrestricted(None)] #[case::v1(Some(false))] #[case::both(Some(true))] -fn explicit_decimal_modes( +fn explicit_decimal_modes_follow_permissions( #[case] allowed_v2: Option, #[values(false, true)] initial_v2: bool, #[values(false, true)] register_later: bool, @@ -132,9 +136,7 @@ fn explicit_decimal_modes( assert_decimal_output( builder, true, - allowed_v2 - .unwrap_or(initial_v2) - .then(decimal_byte_parts_v2_id), + allowed_v2.unwrap_or(true).then(decimal_byte_parts_v2_id), ) } @@ -232,8 +234,7 @@ fn wide_decimal_parts_roundtrip( .into_array(); let mut allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); if compress_children { - allowed.extend(FoRScheme.produced_encodings()); - allowed.extend(BitPackingScheme.produced_encodings()); + allowed.extend([FoR.id(), BitPacked.id(), Patched.id()]); } let compressor = BtrBlocksCompressorBuilder::empty() .with_new_scheme(&DecimalScheme::v2()) diff --git a/vortex-btrblocks/tests/scheme_config.rs b/vortex-btrblocks/tests/scheme_config.rs new file mode 100644 index 00000000000..0ab6c03e651 --- /dev/null +++ b/vortex-btrblocks/tests/scheme_config.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scheme configuration checks every required serialized ID. + +#![cfg(test)] + +use rstest::rstest; +use vortex_alp::ALP; +use vortex_array::ArrayId; +use vortex_array::VTable; +use vortex_array::arrays::Constant; +use vortex_array::arrays::Patched; +use vortex_array::arrays::VarBin; +use vortex_array::arrays::patched::use_experimental_patches; +use vortex_btrblocks::AllowedSerializedIds; +use vortex_btrblocks::Scheme; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::float::ALPScheme; +use vortex_btrblocks::schemes::integer::BitPackingScheme; +use vortex_btrblocks::schemes::integer::SparseScheme; +use vortex_btrblocks::schemes::string::FSSTScheme; +use vortex_fastlanes::BitPacked; +use vortex_fsst::FSST; +use vortex_sparse::Sparse; +use vortex_utils::aliases::hash_set::HashSet; + +#[rstest] +#[case::fsst(&FSSTScheme, vec![FSST.id(), VarBin.id()])] +#[case::sparse(&SparseScheme, vec![Sparse.id(), Constant.id()])] +fn every_required_id_must_be_permitted( + #[case] scheme: &dyn Scheme, + #[case] ids: Vec, +) { + let all: HashSet<_> = ids.iter().copied().collect(); + assert_eq!( + scheme + .configure(&AllowedSerializedIds::Only(all.clone())) + .map(|configured| configured.id()), + Some(scheme.id()) + ); + for id in ids { + let mut allowed = all.clone(); + allowed.remove(&id); + assert!( + scheme + .configure(&AllowedSerializedIds::Only(allowed)) + .is_none() + ); + let forbidden = HashSet::from([id]); + assert!( + scheme + .configure(&AllowedSerializedIds::AllExcept(forbidden)) + .is_none() + ); + } +} + +#[rstest] +#[case::bitpacking(&BitPackingScheme, BitPacked.id())] +#[case::alp(&ALPScheme, ALP.id())] +fn patched_permission_is_required_when_enabled( + #[case] scheme: &dyn Scheme, + #[case] id: ArrayId, +) { + let mut allowed = HashSet::from([id]); + assert_eq!( + scheme + .configure(&AllowedSerializedIds::Only(allowed.clone())) + .is_some(), + !use_experimental_patches() + ); + allowed.insert(Patched.id()); + assert!( + scheme + .configure(&AllowedSerializedIds::Only(allowed.clone())) + .is_some() + ); + allowed.remove(&id); + assert!( + scheme + .configure(&AllowedSerializedIds::Only(allowed)) + .is_none() + ); +} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..0f4111c0742 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -6,7 +6,6 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -24,6 +23,7 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; +use crate::scheme::AllowedSerializedIds; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -48,8 +48,13 @@ impl Scheme for BinaryDictScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { - vec![Dict.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Dict.id()) + .then_some(self) } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index f962c3ff967..1cd2eeb7cb9 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -6,7 +6,6 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for //! external compatibility. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; @@ -28,6 +27,7 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; +use crate::scheme::AllowedSerializedIds; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -54,8 +54,13 @@ impl Scheme for FloatDictScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { - vec![Dict.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Dict.id()) + .then_some(self) } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 27a17ef94ad..eea207a3f0a 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -6,7 +6,6 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; @@ -26,6 +25,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::CascadingCompressor; +use crate::scheme::AllowedSerializedIds; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; use crate::scheme::EstimateVerdict; @@ -49,8 +49,13 @@ impl Scheme for IntDictScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { - vec![Dict.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Dict.id()) + .then_some(self) } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..88e0e1209cf 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -6,7 +6,6 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -24,6 +23,7 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; +use crate::scheme::AllowedSerializedIds; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -48,8 +48,13 @@ impl Scheme for StringDictScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { - vec![Dict.id()] + fn configure( + &self, + allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + allowed_serialized_ids + .contains(&Dict.id()) + .then_some(self) } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 3a2a6281047..dfb2b301b7b 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -4,7 +4,6 @@ use std::sync::LazyLock; use parking_lot::Mutex; -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -35,6 +34,7 @@ use super::structural; use crate::builtins::FloatDictScheme; use crate::builtins::IntDictScheme; use crate::builtins::StringDictScheme; +use crate::scheme::AllowedSerializedIds; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; use crate::scheme::DeferredEstimate; @@ -72,8 +72,11 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -108,8 +111,11 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -144,8 +150,11 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -182,8 +191,11 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -220,8 +232,11 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -258,8 +273,11 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -294,8 +312,11 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -507,8 +528,11 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( @@ -548,8 +572,11 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { - Vec::new() + fn configure( + &self, + _allowed_serialized_ids: &AllowedSerializedIds, + ) -> Option<&dyn Scheme> { + Some(self) } fn expected_compression_ratio( diff --git a/vortex-compressor/src/scheme/allowed_serialized_ids.rs b/vortex-compressor/src/scheme/allowed_serialized_ids.rs new file mode 100644 index 00000000000..e36abd9c5a2 --- /dev/null +++ b/vortex-compressor/src/scheme/allowed_serialized_ids.rs @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialized ID permissions used when configuring compression schemes. + +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_set::HashSet; + +/// The serialized IDs permitted for a configured compression scheme. +/// +/// Defaults to permitting every ID. Restrictions and exclusions only narrow the permitted set. +#[derive(Debug, Clone, Default)] +pub enum AllowedSerializedIds { + /// All IDs are permitted. + #[default] + All, + /// Only these IDs are permitted. + Only(HashSet), + /// All IDs except these are permitted. + AllExcept(HashSet), +} + +impl AllowedSerializedIds { + /// Returns whether a serialized ID is permitted. + pub fn contains(&self, id: &ArrayId) -> bool { + match self { + Self::All => true, + Self::Only(allowed) => allowed.contains(id), + Self::AllExcept(forbidden) => !forbidden.contains(id), + } + } + + /// Restricts the permitted IDs to their intersection with `allowed`. + /// + /// Previously excluded IDs remain excluded, and an empty set permits no IDs. + pub fn restrict_to(&mut self, allowed: &HashSet) { + match self { + Self::All => *self = Self::Only(allowed.clone()), + Self::Only(current) => current.retain(|id| allowed.contains(id)), + Self::AllExcept(forbidden) => { + let permitted = allowed.difference(forbidden).copied().collect(); + *self = Self::Only(permitted); + } + } + } + + /// Excludes an ID, including from subsequent calls to [`Self::restrict_to`]. + pub fn exclude(&mut self, id: ArrayId) { + match self { + Self::All => *self = Self::AllExcept(HashSet::from([id])), + Self::Only(allowed) => { + allowed.remove(&id); + } + Self::AllExcept(forbidden) => { + forbidden.insert(id); + } + } + } +} diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..04375351aa9 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -4,6 +4,9 @@ //! Everything a scheme author implements or receives: the [`Scheme`] trait, exclusion rules, //! compression estimates, and the compression context. +mod allowed_serialized_ids; +pub use allowed_serialized_ids::AllowedSerializedIds; + mod ctx; pub use ctx::CompressorContext; pub use ctx::MAX_CASCADE; @@ -23,7 +26,6 @@ pub use estimate::EstimateVerdict; pub use exclusion::AncestorExclusion; pub use exclusion::ChildSelection; pub use exclusion::DescendantExclusion; -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -124,16 +126,12 @@ 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. - /// - /// 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 - /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme - /// merely rearranges do not need to be declared. + /// Returns a compatible scheme, or `None` if this scheme can produce encodings that serialize + /// with unpermitted IDs. /// - /// 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. - fn produced_encodings(&self) -> Vec; + /// A scheme may return itself or a configuration with different serialized outputs. The + /// returned scheme must keep the same [`SchemeId`]. + fn configure(&self, allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme>; /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies From ad9a42850337c70170643e2413932b58e020f883 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 17:14:43 -0400 Subject: [PATCH 3/6] Preserve declared encodings and upgrade decimal schemes explicitly Restore produced_encodings and add an optional try_upgrade hook. Default decimals to v1, upgrade only with an explicit allowlist, and keep CUDA exclusions effective. Use static decimal instances and cover default writer compatibility. Signed-off-by: "Matt Katz" --- vortex-btrblocks/src/builder.rs | 77 +++++++++++++---- vortex-btrblocks/src/schemes/binary/varbin.rs | 11 +-- vortex-btrblocks/src/schemes/binary/zstd.rs | 11 +-- .../src/schemes/binary/zstd_buffers.rs | 11 +-- vortex-btrblocks/src/schemes/decimal.rs | 34 ++++++-- vortex-btrblocks/src/schemes/float/alp.rs | 16 ++-- vortex-btrblocks/src/schemes/float/alprd.rs | 11 +-- vortex-btrblocks/src/schemes/float/pco.rs | 11 +-- vortex-btrblocks/src/schemes/float/rle.rs | 11 +-- vortex-btrblocks/src/schemes/float/sparse.rs | 11 +-- .../src/schemes/integer/bitpacking.rs | 16 ++-- vortex-btrblocks/src/schemes/integer/delta.rs | 11 +-- vortex-btrblocks/src/schemes/integer/for_.rs | 11 +-- vortex-btrblocks/src/schemes/integer/pco.rs | 11 +-- vortex-btrblocks/src/schemes/integer/rle.rs | 11 +-- .../src/schemes/integer/runend.rs | 11 +-- .../src/schemes/integer/sequence.rs | 11 +-- .../src/schemes/integer/sparse.rs | 11 +-- .../src/schemes/integer/zigzag.rs | 11 +-- vortex-btrblocks/src/schemes/string/fsst.rs | 11 +-- vortex-btrblocks/src/schemes/string/onpair.rs | 11 +-- vortex-btrblocks/src/schemes/string/sparse.rs | 11 +-- vortex-btrblocks/src/schemes/string/zstd.rs | 11 +-- .../src/schemes/string/zstd_buffers.rs | 11 +-- vortex-btrblocks/src/schemes/temporal.rs | 11 +-- vortex-btrblocks/tests/decimal_config.rs | 61 +++++++------ vortex-btrblocks/tests/scheme_config.rs | 85 ------------------- vortex-compressor/src/builtins/dict/binary.rs | 11 +-- vortex-compressor/src/builtins/dict/float.rs | 11 +-- .../src/builtins/dict/integer.rs | 11 +-- vortex-compressor/src/builtins/dict/string.rs | 11 +-- vortex-compressor/src/compressor/tests.rs | 65 +++++--------- vortex-compressor/src/scheme/mod.rs | 26 ++++-- vortex-file/src/tests.rs | 73 ++++++++++++++++ 34 files changed, 325 insertions(+), 403 deletions(-) delete mode 100644 vortex-btrblocks/tests/scheme_config.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index cc748fdcb95..00b22b98da4 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -20,7 +20,7 @@ use crate::schemes::integer; use crate::schemes::string; use crate::schemes::temporal; -/// All default compression schemes, including Decimal v2. +/// All default compression schemes, including Decimal v1. /// /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. @@ -62,7 +62,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ &binary::BinaryDictScheme, &binary::VarBinScheme, // Decimal schemes. - &decimal::DecimalScheme::v2(), + &decimal::DecimalScheme::v1(), // Temporal schemes. &temporal::TemporalScheme, ]; @@ -82,8 +82,9 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); /// `zstd` feature is enabled. /// /// [`Self::retain_allowed_encodings`] restricts serialized IDs. During [`Self::build`], these -/// restrictions select the latest compatible Decimal mode and filter all registered schemes. -/// Decimal uses v2 whenever both serialized IDs are permitted, including without an allowlist. +/// restrictions allow scheme upgrades and filter all registered schemes. Decimal defaults to v1 +/// and upgrades to v2 when an explicit allowlist permits both serialized IDs. Without an +/// allowlist, registered modes are preserved. /// /// # Examples /// @@ -174,8 +175,9 @@ impl BtrBlocksCompressorBuilder { /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, /// but belongs to the opt-in `zstd` edition, so callers filter the two through /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). - /// Decimal schemes are restricted to v1 during build, regardless of when the allowlist or - /// Decimal scheme is supplied, because CUDA does not support lower decimal parts. + /// The decimal v2 serialized ID is excluded because CUDA does not support lower decimal + /// parts. This prevents v1 schemes from upgrading and filters out explicitly registered v2 + /// schemes, regardless of when the allowlist or Decimal scheme is supplied. /// /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. @@ -225,10 +227,10 @@ impl BtrBlocksCompressorBuilder { /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions /// permit. Repeated calls intersect the allowed sets; an empty set permits no serialized IDs. /// - /// Configuration and filtering are deferred until [`Self::build`], including for schemes - /// registered after this call. Registered Decimal schemes use v2 when both decimal IDs are - /// allowed, or v1 when only v1 is allowed. If v1 is not allowed, Decimal is removed, since - /// single-part arrays always serialize as v1. The CUDA preset restricts Decimal to v1. + /// Upgrades and filtering are deferred until [`Self::build`], including for schemes registered + /// after this call. Decimal v1 upgrades to v2 when both decimal IDs are allowed. Schemes are + /// never downgraded: an explicitly registered v2 scheme is removed if either ID is forbidden. + /// Decimal always requires v1 permission, since single-part arrays serialize as v1. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { self.allowed_serialized_ids.restrict_to(allowed); self @@ -241,8 +243,19 @@ impl BtrBlocksCompressorBuilder { fn configured_schemes(self) -> Vec<&'static dyn Scheme> { let mut final_schemes = Vec::with_capacity(self.schemes.len()); + let allowed = &self.allowed_serialized_ids; + let try_upgrade = matches!(allowed, AllowedSerializedIds::Only(_)); for scheme in self.schemes { - if let Some(scheme) = scheme.configure(&self.allowed_serialized_ids) { + let scheme = if try_upgrade { + scheme.try_upgrade(allowed).unwrap_or(scheme) + } else { + scheme + }; + if scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + { final_schemes.push(scheme); } } @@ -258,6 +271,8 @@ mod tests { use super::*; + static DECIMAL_V1: decimal::DecimalScheme = decimal::DecimalScheme::v1(); + #[test] fn empty_starts_with_no_schemes() { let builder = BtrBlocksCompressorBuilder::empty(); @@ -291,6 +306,39 @@ mod tests { assert_eq!(schemes, ALL_SCHEMES); } + #[test] + fn retaining_all_declared_outputs_keeps_every_scheme() { + let allowed = ALL_SCHEMES + .iter() + .flat_map(|scheme| scheme.produced_encodings()) + .collect(); + let schemes = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed) + .configured_schemes(); + assert_eq!(schemes, ALL_SCHEMES); + } + + #[rstest] + fn every_declared_output_must_be_permitted(#[values(false, true)] with_allowlist: bool) { + for scheme in ALL_SCHEMES { + for id in scheme.produced_encodings() { + let mut builder = BtrBlocksCompressorBuilder::empty().with_new_scheme(*scheme); + if with_allowlist { + let mut allowed: HashSet<_> = scheme.produced_encodings().into_iter().collect(); + allowed.remove(&id); + builder = builder.retain_allowed_encodings(&allowed); + } else { + builder.allowed_serialized_ids.exclude(id); + } + assert!( + builder.configured_schemes().is_empty(), + "{} requires {id}", + scheme.scheme_name() + ); + } + } + } + #[rstest] fn restrictions_apply_to_later_registrations(#[values(false, true)] register_later: bool) { let mut builder = BtrBlocksCompressorBuilder::empty(); @@ -330,15 +378,14 @@ mod tests { } #[test] - fn unrelated_forbidden_ids_allow_decimal_v2() { - let mut builder = - BtrBlocksCompressorBuilder::empty().with_new_scheme(&decimal::DecimalScheme::v1()); + fn unrelated_forbidden_ids_preserve_decimal_v1() { + let mut builder = BtrBlocksCompressorBuilder::empty().with_new_scheme(&DECIMAL_V1); builder.allowed_serialized_ids.exclude(FoR.id()); let schemes = builder.configured_schemes(); assert_eq!(schemes.len(), 1); assert_eq!( schemes[0].num_children(), - decimal::DecimalScheme::v2().num_children() + decimal::DecimalScheme::v1().num_children() ); } diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index f06deac7f68..849402a6493 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -9,6 +9,7 @@ //! cascading compressor can then compress with the ordinary integer schemes. For fixed-width //! values the offsets are a constant-stride sequence and collapse to nothing. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -20,7 +21,6 @@ use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; use vortex_array::builders::VarBinBuilder; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::SchemeExt; @@ -44,13 +44,8 @@ impl Scheme for VarBinScheme { canonical.dtype().is_binary() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&VarBin.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![VarBin.id()] } fn num_children(&self) -> usize { diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index 5080cc26a20..d652e344db2 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -3,12 +3,12 @@ //! Zstd compression for binary arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,13 +31,8 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_zstd::Zstd.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::Zstd.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 751559b059d..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -3,12 +3,12 @@ //! Zstd buffer-level binary compression preserving array layout for GPU decompression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,13 +31,8 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_zstd::ZstdBuffers.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::ZstdBuffers.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index d72ebf75006..28ac381a333 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -3,6 +3,7 @@ //! Decimal compression scheme using byte-part decomposition. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -32,6 +33,8 @@ enum DecimalSchemeMode { V2, } +static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); + /// Compression scheme for decimal arrays via byte-part decomposition. /// /// Narrows the decimal to the smallest integer type and compresses its byte parts independently. @@ -39,8 +42,9 @@ enum DecimalSchemeMode { /// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1 /// in either mode, while arrays with lower parts serialize as v2. /// -/// The default uses v2. The builder always selects the latest permitted mode, -/// including for explicitly registered Decimal schemes. The CUDA preset restricts the mode to v1. +/// The default uses v1. An explicit allowlist permitting both decimal IDs lets the builder +/// upgrade v1 to v2. Without an allowlist, the registered mode is preserved. A v2 scheme is +/// filtered out if either ID is forbidden, including under the CUDA preset. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DecimalScheme { mode: DecimalSchemeMode, @@ -50,6 +54,7 @@ impl DecimalScheme { /// Creates a decimal scheme configured for v1, disallowing splitting of wide decimals. /// /// Values that remain wider than `i64` after narrowing stay canonical. + /// The builder may upgrade to v2 if an explicit allowlist permits both serialized IDs. pub const fn v1() -> Self { Self { mode: DecimalSchemeMode::V1, @@ -57,6 +62,7 @@ impl DecimalScheme { } /// Creates a decimal scheme configured for v2, allowing splitting of wide decimals. + /// The builder filters this scheme out if either decimal serialized ID is forbidden. pub const fn v2() -> Self { Self { mode: DecimalSchemeMode::V2, @@ -66,7 +72,7 @@ impl DecimalScheme { impl Default for DecimalScheme { fn default() -> Self { - Self::v2() + Self::v1() } } @@ -79,14 +85,24 @@ impl Scheme for DecimalScheme { matches!(canonical, Canonical::Decimal(_)) } - fn configure(&self, allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme> { - if !allowed_serialized_ids.contains(&decimal_byte_parts_v1_id()) { - return None; + fn produced_encodings(&self) -> Vec { + match self.mode { + DecimalSchemeMode::V1 => vec![decimal_byte_parts_v1_id()], + DecimalSchemeMode::V2 => { + vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()] + } } - if !allowed_serialized_ids.contains(&decimal_byte_parts_v2_id()) { - return Some(&Self::v1()); + } + + fn try_upgrade(&self, allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme> { + if self.mode == DecimalSchemeMode::V1 + && allowed_serialized_ids.contains(&decimal_byte_parts_v1_id()) + && allowed_serialized_ids.contains(&decimal_byte_parts_v2_id()) + { + Some(&DECIMAL_V2) + } else { + None } - Some(&Self::v2()) } /// Children: msp=0, then up to three lower parts in v2 mode. diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index ed624a8d0ce..f9fc7066bf4 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -7,6 +7,7 @@ use vortex_alp::ALP; use vortex_alp::ALPArrayExt; use vortex_alp::ALPArraySlotsExt; use vortex_alp::alp_encode; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -16,7 +17,6 @@ use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -42,14 +42,12 @@ impl Scheme for ALPScheme { canonical.dtype().is_float() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - (allowed_serialized_ids.contains(&ALP.id()) - && (!use_experimental_patches() - || allowed_serialized_ids.contains(&Patched.id()))) - .then_some(self) + fn produced_encodings(&self) -> Vec { + let mut encodings = vec![ALP.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings } /// Children: encoded_ints=0. diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 15b10a575c0..09c0ee85b0d 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -7,6 +7,7 @@ use vortex_alp::ALPRDArrayExt; use vortex_alp::ALPRDArrayOwnedExt; use vortex_alp::RDEncoder; use vortex_alp::RDEncoderExt; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -14,7 +15,6 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -40,13 +40,8 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_alp::ALPRD.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_alp::ALPRD.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index 1100eed3958..416668c2fd0 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -3,12 +3,12 @@ //! Pco (pcodec) float compression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,13 +31,8 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_pco::Pco.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 3645f93c96f..71158b9dc3b 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -3,11 +3,11 @@ //! Run-length float encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; @@ -38,13 +38,8 @@ impl Scheme for FloatRLEScheme { canonical.dtype().is_float() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&RLE.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![RLE.id()] } /// Children: values=0, indices=1, offsets=2. diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 0cd1ea8a4d5..3d9c25b18e4 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -3,6 +3,7 @@ //! Sparse encoding for null-dominated float arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -10,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -41,13 +41,8 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_float() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Sparse.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Sparse.id()] } /// Children: indices=0. diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 95140f13820..5ac7d0e4078 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -3,6 +3,7 @@ //! BitPacking integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,7 +12,6 @@ use vortex_array::VTable; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -40,14 +40,12 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - (allowed_serialized_ids.contains(&BitPacked.id()) - && (!use_experimental_patches() - || allowed_serialized_ids.contains(&Patched.id()))) - .then_some(self) + fn produced_encodings(&self) -> Vec { + let mut encodings = vec![BitPacked.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 752172ee21c..46b2f1e302e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -3,6 +3,7 @@ //! FastLanes Delta integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -13,7 +14,6 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -97,13 +97,8 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Delta.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Delta.id()] } fn num_children(&self) -> usize { diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index ef776bdbcef..476a0dec282 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -3,6 +3,7 @@ //! Frame of Reference integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -13,7 +14,6 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -44,13 +44,8 @@ impl Scheme for FoRScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&FoR.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![FoR.id()] } /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op. diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 2989b4ff159..675a112d44f 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -3,12 +3,12 @@ //! Pco (pcodec) integer compression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -32,13 +32,8 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_pco::Pco.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 9bbc3fe53ae..86a40d0b36a 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -3,6 +3,7 @@ //! Run-length integer encoding and shared RLE compression helpers. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -10,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; @@ -106,13 +106,8 @@ impl Scheme for IntRLEScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&RLE.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![RLE.id()] } /// Children: values=0, indices=1, offsets=2. diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index f6334008352..6a97f7ec37d 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -3,6 +3,7 @@ //! Run-end integer encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -13,7 +14,6 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -48,13 +48,8 @@ impl Scheme for RunEndScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&RunEnd.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![RunEnd.id()] } /// Children: values=0, ends=1. diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index 3abc273bb1f..edcefb99fc2 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -3,6 +3,7 @@ //! Sequence integer encoding for sequential patterns. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -11,7 +12,6 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -43,13 +43,8 @@ impl Scheme for SequenceScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Sequence.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Sequence.id()] } /// Sequence encoding on dictionary codes just adds a layer of indirection without compressing diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 251862007a8..429ff5c1a31 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -3,6 +3,7 @@ //! Sparse integer encoding for single-value-dominated arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -14,7 +15,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::scalar::Scalar; use vortex_compressor::builtins::IntDictScheme; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -46,13 +46,8 @@ impl Scheme for SparseScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - (allowed_serialized_ids.contains(&Sparse.id()) - && allowed_serialized_ids.contains(&Constant.id())) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Sparse.id(), Constant.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 5d14a20e620..0e4be01845a 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -3,6 +3,7 @@ //! ZigZag integer encoding for signed integers. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -13,7 +14,6 @@ use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; @@ -46,13 +46,8 @@ impl Scheme for ZigZagScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&ZigZag.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![ZigZag.id()] } /// Children: encoded=0. diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 9da2b95e3fd..fd3fd28696a 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -5,6 +5,7 @@ use std::sync::Arc; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -15,7 +16,6 @@ use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -50,13 +50,8 @@ impl Scheme for FSSTScheme { canonical.dtype().is_utf8() || canonical.dtype().is_binary() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - (allowed_serialized_ids.contains(&FSST.id()) - && allowed_serialized_ids.contains(&VarBin.id())) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![FSST.id(), VarBin.id()] } /// Children: lengths=0, code_offsets=1. diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 229b80e8d58..06a7ff97f89 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -3,6 +3,7 @@ //! OnPair short-string compression (dict-12). +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -10,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::SchemeId; @@ -48,13 +48,8 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&OnPair.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![OnPair.id()] } /// 4 primitive slot children flow through the cascading compressor: diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 95e9dddac0f..8620c366f77 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -3,6 +3,7 @@ //! Sparse encoding for null-dominated string arrays. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -10,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::ChildSelection; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; @@ -42,13 +42,8 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_utf8() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Sparse.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Sparse.id()] } /// Children: indices=0. diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index c066470d276..84e8860d626 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -3,12 +3,12 @@ //! Zstd string compression without dictionaries (nvCOMP compatible). +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,13 +31,8 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_zstd::Zstd.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::Zstd.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index d1316804aa6..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -3,12 +3,12 @@ //! Zstd buffer-level string compression preserving array layout for GPU decompression. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -31,13 +31,8 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&vortex_zstd::ZstdBuffers.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![vortex_zstd::ZstdBuffers.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index ed660b8a957..79748b69450 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -3,6 +3,7 @@ //! Temporal compression scheme using datetime-part decomposition. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -16,7 +17,6 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::extension::Matcher; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; -use vortex_compressor::scheme::AllowedSerializedIds; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_datetime_parts::DateTimeParts; @@ -55,13 +55,8 @@ impl Scheme for TemporalScheme { ) } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&DateTimeParts.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![DateTimeParts.id()] } /// Children: days=0, seconds=1, subseconds=2. diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs index 1651175b574..8eedb5db78c 100644 --- a/vortex-btrblocks/tests/decimal_config.rs +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -12,11 +12,9 @@ use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::IntoArray; -use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Decimal; use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::Patched; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::i256; @@ -38,12 +36,13 @@ 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_fastlanes::BitPacked; -use vortex_fastlanes::FoR; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use vortex_utils::aliases::hash_set::HashSet; +static DECIMAL_V1: DecimalScheme = DecimalScheme::v1(); +static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); + static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); vortex_decimal_byte_parts::initialize(&session); @@ -83,7 +82,7 @@ fn assert_decimal_output( } #[rstest] -#[case::unrestricted(None, Some(true))] +#[case::unrestricted(None, Some(false))] #[case::neither(Some(vec![]), None)] #[case::v1(Some(vec![decimal_byte_parts_v1_id()]), Some(false))] #[case::v2_only(Some(vec![decimal_byte_parts_v2_id()]), None)] @@ -109,16 +108,13 @@ fn decimal_mode_follows_permissions( #[case::unrestricted(None)] #[case::v1(Some(false))] #[case::both(Some(true))] -fn explicit_decimal_modes_follow_permissions( +fn explicit_decimal_modes_only_upgrade( #[case] allowed_v2: Option, #[values(false, true)] initial_v2: bool, #[values(false, true)] register_later: bool, + #[values(false, true)] wide: bool, ) -> VortexResult<()> { - let scheme: &'static dyn Scheme = if initial_v2 { - &DecimalScheme::v2() - } else { - &DecimalScheme::v1() - }; + let scheme: &'static dyn Scheme = if initial_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; let mut builder = BtrBlocksCompressorBuilder::empty(); if !register_later { builder = builder.with_new_scheme(scheme); @@ -133,15 +129,23 @@ fn explicit_decimal_modes_follow_permissions( if register_later { builder = builder.with_new_scheme(scheme); } - assert_decimal_output( - builder, - true, - allowed_v2.unwrap_or(true).then(decimal_byte_parts_v2_id), - ) + let mode = if initial_v2 && allowed_v2 == Some(false) { + None + } else { + Some(allowed_v2.unwrap_or(initial_v2)) + }; + let expected = match mode { + Some(true) if wide => Some(decimal_byte_parts_v2_id()), + Some(_) if !wide => Some(decimal_byte_parts_v1_id()), + _ => None, + }; + assert_decimal_output(builder, wide, expected) } #[rstest] -fn decimal_permissions_intersect(#[values(false, true)] restrictive_first: bool) -> VortexResult<()> { +fn decimal_permissions_intersect( + #[values(false, true)] restrictive_first: bool, +) -> VortexResult<()> { let v1 = HashSet::from([decimal_byte_parts_v1_id()]); let both = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); let (first, second) = if restrictive_first { @@ -169,27 +173,33 @@ fn permissions_do_not_restore_excluded_decimal() -> VortexResult<()> { #[case::unrestricted(None)] #[case::permissions_first(Some(true))] #[case::permissions_last(Some(false))] -fn cuda_keeps_decimal_v1( +fn cuda_never_uses_decimal_v2( #[case] permissions_first: Option, + #[values(false, true)] initial_v2: bool, #[values(false, true)] register_later: bool, #[values(false, true)] wide: bool, ) -> VortexResult<()> { let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); - let mut builder = BtrBlocksCompressorBuilder::default(); - if register_later { - builder = builder.exclude_schemes([DecimalScheme::default().id()]); + let scheme: &'static dyn Scheme = if initial_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; + let mut builder = BtrBlocksCompressorBuilder::empty(); + if !register_later { + builder = builder.with_new_scheme(scheme); } if permissions_first == Some(true) { builder = builder.retain_allowed_encodings(&allowed); } builder = builder.only_cuda_compatible(); if register_later { - builder = builder.with_new_scheme(&DecimalScheme::v2()); + builder = builder.with_new_scheme(scheme); } if permissions_first.is_some() { builder = builder.retain_allowed_encodings(&allowed); } - assert_decimal_output(builder, wide, (!wide).then(decimal_byte_parts_v1_id)) + assert_decimal_output( + builder, + wide, + (!initial_v2 && !wide).then(decimal_byte_parts_v1_id), + ) } #[rstest] @@ -234,10 +244,11 @@ fn wide_decimal_parts_roundtrip( .into_array(); let mut allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); if compress_children { - allowed.extend([FoR.id(), BitPacked.id(), Patched.id()]); + allowed.extend(FoRScheme.produced_encodings()); + allowed.extend(BitPackingScheme.produced_encodings()); } let compressor = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&DecimalScheme::v2()) + .with_new_scheme(&DECIMAL_V2) .with_new_scheme(&FoRScheme) .with_new_scheme(&BitPackingScheme) .retain_allowed_encodings(&allowed) diff --git a/vortex-btrblocks/tests/scheme_config.rs b/vortex-btrblocks/tests/scheme_config.rs deleted file mode 100644 index 0ab6c03e651..00000000000 --- a/vortex-btrblocks/tests/scheme_config.rs +++ /dev/null @@ -1,85 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Scheme configuration checks every required serialized ID. - -#![cfg(test)] - -use rstest::rstest; -use vortex_alp::ALP; -use vortex_array::ArrayId; -use vortex_array::VTable; -use vortex_array::arrays::Constant; -use vortex_array::arrays::Patched; -use vortex_array::arrays::VarBin; -use vortex_array::arrays::patched::use_experimental_patches; -use vortex_btrblocks::AllowedSerializedIds; -use vortex_btrblocks::Scheme; -use vortex_btrblocks::SchemeExt; -use vortex_btrblocks::schemes::float::ALPScheme; -use vortex_btrblocks::schemes::integer::BitPackingScheme; -use vortex_btrblocks::schemes::integer::SparseScheme; -use vortex_btrblocks::schemes::string::FSSTScheme; -use vortex_fastlanes::BitPacked; -use vortex_fsst::FSST; -use vortex_sparse::Sparse; -use vortex_utils::aliases::hash_set::HashSet; - -#[rstest] -#[case::fsst(&FSSTScheme, vec![FSST.id(), VarBin.id()])] -#[case::sparse(&SparseScheme, vec![Sparse.id(), Constant.id()])] -fn every_required_id_must_be_permitted( - #[case] scheme: &dyn Scheme, - #[case] ids: Vec, -) { - let all: HashSet<_> = ids.iter().copied().collect(); - assert_eq!( - scheme - .configure(&AllowedSerializedIds::Only(all.clone())) - .map(|configured| configured.id()), - Some(scheme.id()) - ); - for id in ids { - let mut allowed = all.clone(); - allowed.remove(&id); - assert!( - scheme - .configure(&AllowedSerializedIds::Only(allowed)) - .is_none() - ); - let forbidden = HashSet::from([id]); - assert!( - scheme - .configure(&AllowedSerializedIds::AllExcept(forbidden)) - .is_none() - ); - } -} - -#[rstest] -#[case::bitpacking(&BitPackingScheme, BitPacked.id())] -#[case::alp(&ALPScheme, ALP.id())] -fn patched_permission_is_required_when_enabled( - #[case] scheme: &dyn Scheme, - #[case] id: ArrayId, -) { - let mut allowed = HashSet::from([id]); - assert_eq!( - scheme - .configure(&AllowedSerializedIds::Only(allowed.clone())) - .is_some(), - !use_experimental_patches() - ); - allowed.insert(Patched.id()); - assert!( - scheme - .configure(&AllowedSerializedIds::Only(allowed.clone())) - .is_some() - ); - allowed.remove(&id); - assert!( - scheme - .configure(&AllowedSerializedIds::Only(allowed)) - .is_none() - ); -} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index 0f4111c0742..c407d0251c6 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -6,6 +6,7 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -23,7 +24,6 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; -use crate::scheme::AllowedSerializedIds; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -48,13 +48,8 @@ impl Scheme for BinaryDictScheme { canonical.dtype().is_binary() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Dict.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 1cd2eeb7cb9..f962c3ff967 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -6,6 +6,7 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for //! external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; @@ -27,7 +28,6 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; -use crate::scheme::AllowedSerializedIds; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -54,13 +54,8 @@ impl Scheme for FloatDictScheme { canonical.dtype().is_float() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Dict.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index eea207a3f0a..27a17ef94ad 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -6,6 +6,7 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; @@ -25,7 +26,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::CascadingCompressor; -use crate::scheme::AllowedSerializedIds; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; use crate::scheme::EstimateVerdict; @@ -49,13 +49,8 @@ impl Scheme for IntDictScheme { canonical.dtype().is_int() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Dict.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 88e0e1209cf..f5cbcd54d89 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -6,6 +6,7 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -23,7 +24,6 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; -use crate::scheme::AllowedSerializedIds; use crate::scheme::ChildSelection; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; @@ -48,13 +48,8 @@ impl Scheme for StringDictScheme { canonical.dtype().is_utf8() } - fn configure( - &self, - allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - allowed_serialized_ids - .contains(&Dict.id()) - .then_some(self) + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index dfb2b301b7b..3a2a6281047 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -4,6 +4,7 @@ use std::sync::LazyLock; use parking_lot::Mutex; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -34,7 +35,6 @@ use super::structural; use crate::builtins::FloatDictScheme; use crate::builtins::IntDictScheme; use crate::builtins::StringDictScheme; -use crate::scheme::AllowedSerializedIds; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; use crate::scheme::DeferredEstimate; @@ -72,11 +72,8 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -111,11 +108,8 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -150,11 +144,8 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -191,11 +182,8 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -232,11 +220,8 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -273,11 +258,8 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -312,11 +294,8 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -528,11 +507,8 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( @@ -572,11 +548,8 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } - fn configure( - &self, - _allowed_serialized_ids: &AllowedSerializedIds, - ) -> Option<&dyn Scheme> { - Some(self) + fn produced_encodings(&self) -> Vec { + Vec::new() } fn expected_compression_ratio( diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 04375351aa9..21964948572 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -26,6 +26,7 @@ pub use estimate::EstimateVerdict; pub use exclusion::AncestorExclusion; pub use exclusion::ChildSelection; pub use exclusion::DescendantExclusion; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -126,12 +127,27 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// Returns a compatible scheme, or `None` if this scheme can produce encodings that serialize - /// with unpermitted IDs. + /// The serialized IDs this scheme itself may write into its compressed output. /// - /// A scheme may return itself or a configuration with different serialized outputs. The - /// returned scheme must keep the same [`SchemeId`]. - fn configure(&self, allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme>; + /// 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 + /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme + /// 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. + fn produced_encodings(&self) -> Vec; + + /// Returns a newer configuration supported by the permitted serialized IDs, if available. + /// + /// `None` keeps the original scheme. An upgrade must preserve the [`SchemeId`] and must not + /// downgrade the registered configuration. The caller checks the selected scheme's + /// [`produced_encodings`](Self::produced_encodings) before using it. + /// + /// The builder only attempts upgrades when an explicit allowlist is supplied. + fn try_upgrade(&self, _allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme> { + None + } /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 640de874d2b..48294ac6802 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -38,6 +38,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::dtype::i256; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; @@ -72,7 +73,10 @@ use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; +use vortex_edition::EDITION_DECLARATIONS; use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::session::RuntimeSession; @@ -99,6 +103,7 @@ use crate::V1_FOOTER_FBS_SIZE; use crate::VERSION; use crate::VortexFile; use crate::WriteOptionsSessionExt; +use crate::WriteStrategyBuilder; use crate::flatbuffers::footer as fb; use crate::footer::SegmentSpec; static SESSION: LazyLock = LazyLock::new(|| { @@ -173,6 +178,74 @@ async fn test_read_simple() { assert_eq!(row_count, 8); } +#[rstest] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn default_strategy_keeps_wide_decimals_compatible( + #[values(false, true)] use_i256: bool, + #[values(false, true)] explicit_compressor: bool, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + crate::register_default_encodings(&session); + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + + let array = if use_i256 { + DecimalArray::new( + (0..1024u128) + .map(|i| i256::from_parts(i * 17, 1i128 << 70)) + .collect::>(), + DecimalDType::new(76, 2), + Validity::NonNullable, + ) + } else { + DecimalArray::new( + (0..1024i128) + .map(|i| (1i128 << 70) + i * 17) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + } + .into_array(); + let strategy = WriteStrategyBuilder::default() + .with_row_block_size(256) + .with_data_block_target_bytes(None); + let strategy = if explicit_compressor { + strategy.with_btrblocks_builder(BtrBlocksCompressorBuilder::default()) + } else { + strategy + } + .build(); + + for disable_editions in [false, true] { + let options = session.write_options().with_strategy(Arc::clone(&strategy)); + let options = if disable_editions { + options.disable_editions() + } else { + options + }; + let mut buffer = ByteBufferMut::empty(); + options + .write(&mut buffer, array.clone().to_array_stream()) + .await?; + let actual = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(array, actual, &mut session.create_execution_ctx()); + } + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_round_trip_many_types() { From 259a795342cfd30be008d16a7b154aae5c77246b Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 21:11:56 -0400 Subject: [PATCH 4/6] Initialize compressor permissions from serialized IDs Accept serialized IDs when constructing BtrBlocks builders and derive default permissions from the default core edition. Replace permission variants with a set, preserve CUDA exclusions, and migrate writer and benchmark callers to constructor-based permissions. Signed-off-by: "Matt Katz" --- benchmarks/compress-bench/src/gpu/vortex.rs | 11 +- vortex-bench/src/conversions.rs | 11 +- vortex-bench/src/lib.rs | 18 +- vortex-btrblocks/Cargo.toml | 2 +- vortex-btrblocks/src/builder.rs | 184 +++++++----------- vortex-btrblocks/src/canonical_compressor.rs | 3 +- vortex-btrblocks/src/schemes/decimal.rs | 7 +- .../schemes/string/scheme_selection_tests.rs | 20 +- vortex-btrblocks/tests/decimal_config.rs | 76 ++------ vortex-btrblocks/tests/golden.rs | 30 +-- .../src/scheme/allowed_serialized_ids.rs | 59 ------ vortex-compressor/src/scheme/mod.rs | 9 +- vortex-cuda/src/layout.rs | 6 +- vortex-edition/src/declarations/mod.rs | 4 + vortex-edition/src/lib.rs | 1 + vortex-file/src/tests.rs | 25 ++- vortex-file/src/writer.rs | 5 +- vortex-python/src/io.rs | 5 +- vortex-tui/src/convert.rs | 5 +- vortex/src/editions/mod.rs | 4 +- 20 files changed, 161 insertions(+), 324 deletions(-) delete mode 100644 vortex-compressor/src/scheme/allowed_serialized_ids.rs diff --git a/benchmarks/compress-bench/src/gpu/vortex.rs b/benchmarks/compress-bench/src/gpu/vortex.rs index bb48b461e87..8070a75ef22 100644 --- a/benchmarks/compress-bench/src/gpu/vortex.rs +++ b/benchmarks/compress-bench/src/gpu/vortex.rs @@ -22,7 +22,6 @@ use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; -use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::error::VortexResult; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; @@ -36,8 +35,8 @@ use vortex_bench::compress::Compressed; use vortex_bench::compress::CompressedData; use vortex_bench::compress::Compressor; use vortex_bench::compress::Uncompressed; +use vortex_bench::compressor_builder_for_session; use vortex_bench::conversions::parquet_to_vortex_chunks_with_batch_size; -use vortex_bench::retain_edition_encodings; use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaExecutionCtx; use vortex_cuda::CudaOpenOptionsExt; @@ -100,11 +99,9 @@ impl Compressor for GpuVortexCompressor { // partition rather than whatever the default strategy would regroup them into. let strategy = Arc::new(ChunkedLayoutStrategy::new(CompressingStrategy::new( CudaFlatLayoutStrategy::default(), - retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().only_cuda_compatible(), - ) - .build(), + compressor_builder_for_session(&SESSION) + .only_cuda_compatible() + .build(), ))); let start = Instant::now(); SESSION diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 6b4ed871f2e..cdfd8c4d171 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -37,7 +37,6 @@ use vortex::array::arrays::struct_::StructArrayExt; use vortex::array::builders::builder_with_capacity_in; use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; -use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; use vortex::dtype::FieldPath; use vortex::dtype::StructFields; @@ -66,7 +65,7 @@ use wkb::writer::write_geometry; use crate::CompactionStrategy; use crate::Format; use crate::SESSION; -use crate::retain_edition_encodings; +use crate::compressor_builder_for_session; use crate::utils::file::idempotent_async; /// Memory budget per concurrent conversion stream in GB. This is somewhat arbitary. @@ -248,10 +247,8 @@ fn write_options_for( let mut builder = WriteStrategyBuilder::default(); if matches!(compaction, CompactionStrategy::Compact) { - builder = builder.with_btrblocks_builder(retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().with_compact(), - )); + builder = builder + .with_btrblocks_builder(compressor_builder_for_session(&SESSION).with_compact()); } for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); @@ -263,7 +260,7 @@ fn write_options_for( fn no_dict_layout() -> Arc { Arc::new(CompressingStrategy::new( ChunkedLayoutStrategy::new(FlatLayoutStrategy::default()), - retain_edition_encodings(&SESSION, BtrBlocksCompressorBuilder::default()).build(), + compressor_builder_for_session(&SESSION).build(), )) } diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 569dfc74a4d..7dcfb8a01f3 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -255,10 +255,7 @@ impl CompactionStrategy { match self { CompactionStrategy::Compact => options.with_strategy( WriteStrategyBuilder::default() - .with_btrblocks_builder(retain_edition_encodings( - &SESSION, - BtrBlocksCompressorBuilder::default().with_compact(), - )) + .with_btrblocks_builder(compressor_builder_for_session(&SESSION).with_compact()) .build(), ), CompactionStrategy::Default => options, @@ -266,19 +263,16 @@ impl CompactionStrategy { } } -/// Restrict `builder` to the encodings permitted by the session's enabled editions. +/// Create a compressor builder permitting the session's enabled array encodings. /// -/// The default writer applies this filter itself. An explicit strategy bypasses it, so a -/// benchmark that builds its own compressor applies it here to stay within editions. -pub fn retain_edition_encodings( - session: &VortexSession, - builder: BtrBlocksCompressorBuilder, -) -> BtrBlocksCompressorBuilder { +/// Benchmarks supplying an explicit strategy use the session's permissions, including opt-in +/// editions, instead of the compressor builder's default core edition. +pub fn compressor_builder_for_session(session: &VortexSession) -> BtrBlocksCompressorBuilder { let allowed = session .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - builder.retain_allowed_encodings(&allowed) + BtrBlocksCompressorBuilder::new(allowed) } /// Verify that local data has already been prepared for the requested benchmark formats. diff --git a/vortex-btrblocks/Cargo.toml b/vortex-btrblocks/Cargo.toml index 24a03337768..a92c0a27fa9 100644 --- a/vortex-btrblocks/Cargo.toml +++ b/vortex-btrblocks/Cargo.toml @@ -26,6 +26,7 @@ vortex-buffer = { workspace = true } vortex-compressor = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-fsst = { workspace = true } @@ -49,7 +50,6 @@ tpchgen = { workspace = true } tpchgen-arrow = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-arrow = { workspace = true } -vortex-edition = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 00b22b98da4..5e28eaab404 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -3,8 +3,10 @@ //! Builder for configuring `BtrBlocksCompressor` instances. -use vortex_array::ArrayId; use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_edition::ComponentKind; +use vortex_edition::DEFAULT_CORE_EDITION; +use vortex_edition::EDITION_DECLARATIONS; use vortex_utils::aliases::hash_set::HashSet; use crate::AllowedSerializedIds; @@ -76,15 +78,15 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); /// Builder for creating configured [`BtrBlocksCompressor`] instances. /// -/// By default, all schemes in [`ALL_SCHEMES`] are enabled in a deterministic order. Feature-gated +/// By default, all schemes in [`ALL_SCHEMES`] are registered in a deterministic order. Feature-gated /// schemes (Pco, Zstd) are not in `ALL_SCHEMES` and must be added explicitly via /// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) or `with_compact` when the /// `zstd` feature is enabled. /// -/// [`Self::retain_allowed_encodings`] restricts serialized IDs. During [`Self::build`], these -/// restrictions allow scheme upgrades and filter all registered schemes. Decimal defaults to v1 -/// and upgrades to v2 when an explicit allowlist permits both serialized IDs. Without an -/// allowlist, registered modes are preserved. +/// [`Self::new`] takes the permitted serialized IDs once. During [`Self::build`], these +/// permissions allow scheme upgrades and filter all registered schemes. The default builder +/// permits the array IDs in [`DEFAULT_CORE_EDITION`]. Decimal defaults to v1 and upgrades to v2 +/// when both serialized IDs are permitted. /// /// # Examples /// @@ -92,7 +94,7 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); /// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt}; /// use vortex_btrblocks::schemes::integer::IntDictScheme; /// -/// // Default compressor with all schemes in ALL_SCHEMES. +/// // Default compressor restricted to the default core edition. /// let compressor = BtrBlocksCompressorBuilder::default().build(); /// /// // Remove specific schemes. @@ -108,21 +110,37 @@ pub struct BtrBlocksCompressorBuilder { impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { + let allowed_serialized_ids = EDITION_DECLARATIONS + .iter() + .filter(|declaration| declaration.edition.id.is_at_or_before(&DEFAULT_CORE_EDITION)) + .flat_map(|declaration| declaration.added) + .filter(|member| member.kind == ComponentKind::Array) + .map(|member| member.component.component_id()) + .collect(); + Self::new(allowed_serialized_ids) + } +} + +impl BtrBlocksCompressorBuilder { + /// Creates a builder with all default schemes and the supplied serialized ID permissions. + /// + /// An empty set permits no serialized IDs. Upgrades and filtering are deferred until + /// [`Self::build`], including for schemes registered later. Schemes are never downgraded. + pub fn new(allowed_serialized_ids: AllowedSerializedIds) -> Self { Self { schemes: ALL_SCHEMES.to_vec(), - allowed_serialized_ids: AllowedSerializedIds::default(), + allowed_serialized_ids, } } -} -impl BtrBlocksCompressorBuilder { /// Creates a builder with no schemes registered. /// /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor. - pub fn empty() -> Self { + /// The supplied serialized ID permissions apply to every scheme registered later. + pub fn empty(allowed_serialized_ids: AllowedSerializedIds) -> Self { Self { schemes: Vec::new(), - allowed_serialized_ids: AllowedSerializedIds::default(), + allowed_serialized_ids, } } @@ -173,11 +191,11 @@ impl BtrBlocksCompressorBuilder { /// /// 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, - /// but belongs to the opt-in `zstd` edition, so callers filter the two through - /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). + /// but belongs to the opt-in `zstd` edition. The permissions supplied to [`Self::new`] + /// determine which schemes survive. /// The decimal v2 serialized ID is excluded because CUDA does not support lower decimal /// parts. This prevents v1 schemes from upgrading and filters out explicitly registered v2 - /// schemes, regardless of when the allowlist or Decimal scheme is supplied. + /// schemes, including Decimal schemes registered after this call. /// /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. @@ -205,7 +223,7 @@ impl BtrBlocksCompressorBuilder { let mut builder = self.exclude_schemes(excluded); builder .allowed_serialized_ids - .exclude(decimal_byte_parts_v2_id()); + .remove(&decimal_byte_parts_v2_id()); #[cfg(feature = "zstd")] let builder = builder @@ -222,20 +240,6 @@ impl BtrBlocksCompressorBuilder { self } - /// Restricts schemes to those whose produced serialized IDs all belong to `allowed`. - /// - /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. Repeated calls intersect the allowed sets; an empty set permits no serialized IDs. - /// - /// Upgrades and filtering are deferred until [`Self::build`], including for schemes registered - /// after this call. Decimal v1 upgrades to v2 when both decimal IDs are allowed. Schemes are - /// never downgraded: an explicitly registered v2 scheme is removed if either ID is forbidden. - /// Decimal always requires v1 permission, since single-part arrays serialize as v1. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.allowed_serialized_ids.restrict_to(allowed); - self - } - /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { BtrBlocksCompressor(CascadingCompressor::new(self.configured_schemes())) @@ -244,13 +248,8 @@ impl BtrBlocksCompressorBuilder { fn configured_schemes(self) -> Vec<&'static dyn Scheme> { let mut final_schemes = Vec::with_capacity(self.schemes.len()); let allowed = &self.allowed_serialized_ids; - let try_upgrade = matches!(allowed, AllowedSerializedIds::Only(_)); for scheme in self.schemes { - let scheme = if try_upgrade { - scheme.try_upgrade(allowed).unwrap_or(scheme) - } else { - scheme - }; + let scheme = scheme.try_upgrade(allowed).unwrap_or(scheme); if scheme .produced_encodings() .iter() @@ -265,17 +264,17 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { - use rstest::rstest; use vortex_array::VTable; + use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; + use vortex_edition::EditionSession; + use vortex_error::VortexResult; use vortex_fastlanes::FoR; use super::*; - static DECIMAL_V1: decimal::DecimalScheme = decimal::DecimalScheme::v1(); - #[test] fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); + let builder = BtrBlocksCompressorBuilder::empty(HashSet::new()); assert!(builder.schemes.is_empty()); } @@ -286,50 +285,57 @@ mod tests { } #[test] - fn retain_allowed_encodings_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let schemes = BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed) - .configured_schemes(); + fn allowed_encodings_filter_schemes() { + let schemes = + BtrBlocksCompressorBuilder::new(HashSet::from([FoR.id()])).configured_schemes(); assert_eq!(schemes.len(), 1); assert_eq!(schemes[0].id(), integer::FoRScheme.id()); - let none = BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&HashSet::new()) - .configured_schemes(); + let none = BtrBlocksCompressorBuilder::new(HashSet::new()).configured_schemes(); assert!(none.is_empty()); } #[test] - fn unrestricted_configuration_preserves_scheme_order() { + fn default_configuration_preserves_scheme_order() { let schemes = BtrBlocksCompressorBuilder::default().configured_schemes(); assert_eq!(schemes, ALL_SCHEMES); } #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { + fn default_permissions_match_default_edition() -> VortexResult<()> { + let editions = EditionSession::empty(); + for declaration in EDITION_DECLARATIONS { + editions.declare(declaration)?; + } + let allowed: HashSet<_> = editions + .components_in(&DEFAULT_CORE_EDITION, ComponentKind::Array) + .into_iter() + .map(|inclusion| inclusion.component_id) + .collect(); + let builder = BtrBlocksCompressorBuilder::default(); + assert_eq!(builder.allowed_serialized_ids, allowed); + assert!(allowed.contains(&decimal_byte_parts_v1_id())); + assert!(!allowed.contains(&decimal_byte_parts_v2_id())); + Ok(()) + } + + #[test] + fn allowing_all_declared_outputs_keeps_every_scheme() { let allowed = ALL_SCHEMES .iter() .flat_map(|scheme| scheme.produced_encodings()) .collect(); - let schemes = BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed) - .configured_schemes(); + let schemes = BtrBlocksCompressorBuilder::new(allowed).configured_schemes(); assert_eq!(schemes, ALL_SCHEMES); } - #[rstest] - fn every_declared_output_must_be_permitted(#[values(false, true)] with_allowlist: bool) { + #[test] + fn every_declared_output_must_be_permitted() { for scheme in ALL_SCHEMES { for id in scheme.produced_encodings() { - let mut builder = BtrBlocksCompressorBuilder::empty().with_new_scheme(*scheme); - if with_allowlist { - let mut allowed: HashSet<_> = scheme.produced_encodings().into_iter().collect(); - allowed.remove(&id); - builder = builder.retain_allowed_encodings(&allowed); - } else { - builder.allowed_serialized_ids.exclude(id); - } + let mut allowed: HashSet<_> = scheme.produced_encodings().into_iter().collect(); + allowed.remove(&id); + let builder = BtrBlocksCompressorBuilder::empty(allowed).with_new_scheme(*scheme); assert!( builder.configured_schemes().is_empty(), "{} requires {id}", @@ -339,54 +345,12 @@ mod tests { } } - #[rstest] - fn restrictions_apply_to_later_registrations(#[values(false, true)] register_later: bool) { - let mut builder = BtrBlocksCompressorBuilder::empty(); - if !register_later { - builder = builder.with_new_scheme(&integer::FoRScheme); - } - builder = builder.retain_allowed_encodings(&HashSet::new()); - if register_later { - builder = builder.with_new_scheme(&integer::FoRScheme); - } - let allowed = HashSet::from([FoR.id()]); - assert!( - builder - .retain_allowed_encodings(&allowed) - .configured_schemes() - .is_empty() - ); - } - - #[rstest] - fn forbidden_ids_filter_schemes( - #[values(false, true)] with_allowlist: bool, - #[values(false, true)] register_later: bool, - ) { - let mut builder = BtrBlocksCompressorBuilder::empty(); - if !register_later { - builder = builder.with_new_scheme(&integer::FoRScheme); - } - builder.allowed_serialized_ids.exclude(FoR.id()); - if with_allowlist { - builder = builder.retain_allowed_encodings(&HashSet::from([FoR.id()])); - } - if register_later { - builder = builder.with_new_scheme(&integer::FoRScheme); - } - assert!(builder.configured_schemes().is_empty()); - } - #[test] - fn unrelated_forbidden_ids_preserve_decimal_v1() { - let mut builder = BtrBlocksCompressorBuilder::empty().with_new_scheme(&DECIMAL_V1); - builder.allowed_serialized_ids.exclude(FoR.id()); - let schemes = builder.configured_schemes(); - assert_eq!(schemes.len(), 1); - assert_eq!( - schemes[0].num_children(), - decimal::DecimalScheme::v1().num_children() - ); + fn permissions_apply_to_later_registrations() { + let builder = BtrBlocksCompressorBuilder::new(HashSet::new()) + .exclude_schemes([integer::FoRScheme.id()]) + .with_new_scheme(&integer::FoRScheme); + assert!(builder.configured_schemes().is_empty()); } #[test] diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index d93be365550..f080de89ba9 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -292,9 +292,8 @@ mod tests { // The CUDA preset carries both Zstd schemes; the edition filter decides which one // survives. - let compressor = BtrBlocksCompressorBuilder::default() + let compressor = BtrBlocksCompressorBuilder::new(HashSet::from([allowed])) .only_cuda_compatible() - .retain_allowed_encodings(&HashSet::from([allowed])) .build(); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 28ac381a333..7a3a32b08b4 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -42,9 +42,8 @@ static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); /// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1 /// in either mode, while arrays with lower parts serialize as v2. /// -/// The default uses v1. An explicit allowlist permitting both decimal IDs lets the builder -/// upgrade v1 to v2. Without an allowlist, the registered mode is preserved. A v2 scheme is -/// filtered out if either ID is forbidden, including under the CUDA preset. +/// The default uses v1. Permitting both decimal IDs lets the builder upgrade v1 to v2. +/// A v2 scheme is filtered out if either ID is forbidden, including under the CUDA preset. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DecimalScheme { mode: DecimalSchemeMode, @@ -54,7 +53,7 @@ impl DecimalScheme { /// Creates a decimal scheme configured for v1, disallowing splitting of wide decimals. /// /// Values that remain wider than `i64` after narrowing stay canonical. - /// The builder may upgrade to v2 if an explicit allowlist permits both serialized IDs. + /// The builder may upgrade to v2 if its permissions include both serialized IDs. pub const fn v1() -> Self { Self { mode: DecimalSchemeMode::V1, diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index aac0b4de4de..55ccf0159cb 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -17,6 +17,11 @@ use vortex_fsst::FSST; use vortex_session::VortexSession; use crate::BtrBlocksCompressor; +use crate::BtrBlocksCompressorBuilder; +use crate::Scheme; +use crate::SchemeExt; +use crate::schemes::string::FSSTScheme; +use crate::schemes::string::onpair::OnPairScheme; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); @@ -48,9 +53,6 @@ fn test_dict_compressed() -> VortexResult<()> { #[test] fn test_all_schemes_includes_onpair() { - use crate::SchemeExt; - use crate::schemes::string::onpair::OnPairScheme; - let ids: Vec<_> = crate::ALL_SCHEMES.iter().map(|s| s.id()).collect(); assert!( ids.contains(&OnPairScheme.id()), @@ -85,10 +87,6 @@ fn test_default_btrblocks_compressor_selects_onpair() -> VortexResult<()> { /// still produces an FSST array. #[test] fn test_fsst_in_default_scheme_list() -> VortexResult<()> { - use crate::BtrBlocksCompressorBuilder; - use crate::SchemeExt; - use crate::schemes::string::FSSTScheme; - // FSST is registered by default. assert!( crate::ALL_SCHEMES.iter().any(|s| s.id() == FSSTScheme.id()), @@ -106,9 +104,11 @@ fn test_fsst_in_default_scheme_list() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressor = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&FSSTScheme) - .build(); + let compressor = BtrBlocksCompressorBuilder::empty( + FSSTScheme.produced_encodings().into_iter().collect(), + ) + .with_new_scheme(&FSSTScheme) + .build(); let compressed = compressor.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs index 8eedb5db78c..09224d91c5d 100644 --- a/vortex-btrblocks/tests/decimal_config.rs +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -82,7 +82,7 @@ fn assert_decimal_output( } #[rstest] -#[case::unrestricted(None, Some(false))] +#[case::default_edition(None, Some(false))] #[case::neither(Some(vec![]), None)] #[case::v1(Some(vec![decimal_byte_parts_v1_id()]), Some(false))] #[case::v2_only(Some(vec![decimal_byte_parts_v2_id()]), None)] @@ -92,10 +92,9 @@ fn decimal_mode_follows_permissions( #[case] v2: Option, #[values(false, true)] wide: bool, ) -> VortexResult<()> { - let mut builder = BtrBlocksCompressorBuilder::default(); - if let Some(ids) = ids { - builder = builder.retain_allowed_encodings(&ids.into_iter().collect()); - } + let builder = ids + .map(|ids| BtrBlocksCompressorBuilder::new(ids.into_iter().collect())) + .unwrap_or_default(); let expected = match v2 { Some(true) if wide => Some(decimal_byte_parts_v2_id()), Some(_) if !wide => Some(decimal_byte_parts_v1_id()), @@ -105,34 +104,21 @@ fn decimal_mode_follows_permissions( } #[rstest] -#[case::unrestricted(None)] -#[case::v1(Some(false))] -#[case::both(Some(true))] fn explicit_decimal_modes_only_upgrade( - #[case] allowed_v2: Option, + #[values(false, true)] allowed_v2: bool, #[values(false, true)] initial_v2: bool, - #[values(false, true)] register_later: bool, #[values(false, true)] wide: bool, ) -> VortexResult<()> { let scheme: &'static dyn Scheme = if initial_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; - let mut builder = BtrBlocksCompressorBuilder::empty(); - if !register_later { - builder = builder.with_new_scheme(scheme); - } - if let Some(v2) = allowed_v2 { - let mut allowed = HashSet::from([decimal_byte_parts_v1_id()]); - if v2 { - allowed.insert(decimal_byte_parts_v2_id()); - } - builder = builder.retain_allowed_encodings(&allowed); - } - if register_later { - builder = builder.with_new_scheme(scheme); + let mut allowed = HashSet::from([decimal_byte_parts_v1_id()]); + if allowed_v2 { + allowed.insert(decimal_byte_parts_v2_id()); } - let mode = if initial_v2 && allowed_v2 == Some(false) { + let builder = BtrBlocksCompressorBuilder::empty(allowed).with_new_scheme(scheme); + let mode = if initial_v2 && !allowed_v2 { None } else { - Some(allowed_v2.unwrap_or(initial_v2)) + Some(allowed_v2) }; let expected = match mode { Some(true) if wide => Some(decimal_byte_parts_v2_id()), @@ -142,59 +128,30 @@ fn explicit_decimal_modes_only_upgrade( assert_decimal_output(builder, wide, expected) } -#[rstest] -fn decimal_permissions_intersect( - #[values(false, true)] restrictive_first: bool, -) -> VortexResult<()> { - let v1 = HashSet::from([decimal_byte_parts_v1_id()]); - let both = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); - let (first, second) = if restrictive_first { - (&v1, &both) - } else { - (&both, &v1) - }; - let builder = BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(first) - .retain_allowed_encodings(second); - assert_decimal_output(builder.clone(), false, Some(decimal_byte_parts_v1_id()))?; - assert_decimal_output(builder, true, None) -} - #[test] -fn permissions_do_not_restore_excluded_decimal() -> VortexResult<()> { +fn upgrades_do_not_restore_excluded_decimal() -> VortexResult<()> { let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); - let builder = BtrBlocksCompressorBuilder::default() - .exclude_schemes([DecimalScheme::default().id()]) - .retain_allowed_encodings(&allowed); + let builder = + BtrBlocksCompressorBuilder::new(allowed).exclude_schemes([DecimalScheme::default().id()]); assert_decimal_output(builder, false, None) } #[rstest] -#[case::unrestricted(None)] -#[case::permissions_first(Some(true))] -#[case::permissions_last(Some(false))] fn cuda_never_uses_decimal_v2( - #[case] permissions_first: Option, #[values(false, true)] initial_v2: bool, #[values(false, true)] register_later: bool, #[values(false, true)] wide: bool, ) -> VortexResult<()> { let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); let scheme: &'static dyn Scheme = if initial_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; - let mut builder = BtrBlocksCompressorBuilder::empty(); + let mut builder = BtrBlocksCompressorBuilder::empty(allowed); if !register_later { builder = builder.with_new_scheme(scheme); } - if permissions_first == Some(true) { - builder = builder.retain_allowed_encodings(&allowed); - } builder = builder.only_cuda_compatible(); if register_later { builder = builder.with_new_scheme(scheme); } - if permissions_first.is_some() { - builder = builder.retain_allowed_encodings(&allowed); - } assert_decimal_output( builder, wide, @@ -247,11 +204,10 @@ fn wide_decimal_parts_roundtrip( allowed.extend(FoRScheme.produced_encodings()); allowed.extend(BitPackingScheme.produced_encodings()); } - let compressor = BtrBlocksCompressorBuilder::empty() + let compressor = BtrBlocksCompressorBuilder::empty(allowed) .with_new_scheme(&DECIMAL_V2) .with_new_scheme(&FoRScheme) .with_new_scheme(&BitPackingScheme) - .retain_allowed_encodings(&allowed) .build(); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array, &mut ctx)?; diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index fc636252c27..9a105f32467 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -414,35 +414,18 @@ fn edition_session(editions: &[EditionId]) -> VortexResult { Ok(session) } -fn compressor_for_session( - session: &VortexSession, - builder: BtrBlocksCompressorBuilder, -) -> BtrBlocksCompressor { +fn compressor_builder_for_session(session: &VortexSession) -> BtrBlocksCompressorBuilder { let allowed = session .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - without_onpair(builder) - .retain_allowed_encodings(&allowed) - .build() -} - -/// Like [`compressor_for_session`] but keeps OnPair in the scheme pool. -fn compressor_with_onpair( - session: &VortexSession, - builder: BtrBlocksCompressorBuilder, -) -> BtrBlocksCompressor { - let allowed = session - .enabled_component_ids(ComponentKind::Array) - .into_iter() - .collect(); - builder.retain_allowed_encodings(&allowed).build() + BtrBlocksCompressorBuilder::new(allowed) } #[test] fn golden_regular() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3])?; - let compressor = compressor_for_session(&session, BtrBlocksCompressorBuilder::default()); + let compressor = without_onpair(compressor_builder_for_session(&session)).build(); golden_corpus_snapshots("regular", &compressor) } @@ -450,7 +433,7 @@ fn golden_regular() -> VortexResult<()> { #[test] fn golden_onpair() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3])?; - let compressor = compressor_with_onpair(&session, BtrBlocksCompressorBuilder::default()); + let compressor = compressor_builder_for_session(&session).build(); golden_snapshots( "onpair", &compressor, @@ -464,9 +447,6 @@ fn golden_compact() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3])?; vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; - let compressor = compressor_for_session( - &session, - BtrBlocksCompressorBuilder::default().with_compact(), - ); + let compressor = without_onpair(compressor_builder_for_session(&session).with_compact()).build(); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-compressor/src/scheme/allowed_serialized_ids.rs b/vortex-compressor/src/scheme/allowed_serialized_ids.rs deleted file mode 100644 index e36abd9c5a2..00000000000 --- a/vortex-compressor/src/scheme/allowed_serialized_ids.rs +++ /dev/null @@ -1,59 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Serialized ID permissions used when configuring compression schemes. - -use vortex_array::ArrayId; -use vortex_utils::aliases::hash_set::HashSet; - -/// The serialized IDs permitted for a configured compression scheme. -/// -/// Defaults to permitting every ID. Restrictions and exclusions only narrow the permitted set. -#[derive(Debug, Clone, Default)] -pub enum AllowedSerializedIds { - /// All IDs are permitted. - #[default] - All, - /// Only these IDs are permitted. - Only(HashSet), - /// All IDs except these are permitted. - AllExcept(HashSet), -} - -impl AllowedSerializedIds { - /// Returns whether a serialized ID is permitted. - pub fn contains(&self, id: &ArrayId) -> bool { - match self { - Self::All => true, - Self::Only(allowed) => allowed.contains(id), - Self::AllExcept(forbidden) => !forbidden.contains(id), - } - } - - /// Restricts the permitted IDs to their intersection with `allowed`. - /// - /// Previously excluded IDs remain excluded, and an empty set permits no IDs. - pub fn restrict_to(&mut self, allowed: &HashSet) { - match self { - Self::All => *self = Self::Only(allowed.clone()), - Self::Only(current) => current.retain(|id| allowed.contains(id)), - Self::AllExcept(forbidden) => { - let permitted = allowed.difference(forbidden).copied().collect(); - *self = Self::Only(permitted); - } - } - } - - /// Excludes an ID, including from subsequent calls to [`Self::restrict_to`]. - pub fn exclude(&mut self, id: ArrayId) { - match self { - Self::All => *self = Self::AllExcept(HashSet::from([id])), - Self::Only(allowed) => { - allowed.remove(&id); - } - Self::AllExcept(forbidden) => { - forbidden.insert(id); - } - } - } -} diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 21964948572..6af68761f2e 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -4,9 +4,6 @@ //! Everything a scheme author implements or receives: the [`Scheme`] trait, exclusion rules, //! compression estimates, and the compression context. -mod allowed_serialized_ids; -pub use allowed_serialized_ids::AllowedSerializedIds; - mod ctx; pub use ctx::CompressorContext; pub use ctx::MAX_CASCADE; @@ -31,11 +28,15 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_error::VortexResult; +use vortex_utils::aliases::hash_set::HashSet; use crate::CascadingCompressor; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; +/// The serialized IDs permitted for compression schemes. +pub type AllowedSerializedIds = HashSet; + /// Unique identifier for a compression scheme. /// /// The only way to obtain a [`SchemeId`] is through [`SchemeExt::id()`], which is auto-implemented @@ -143,8 +144,6 @@ pub trait Scheme: Debug + Send + Sync { /// `None` keeps the original scheme. An upgrade must preserve the [`SchemeId`] and must not /// downgrade the registered configuration. The caller checks the selected scheme's /// [`produced_encodings`](Self::produced_encodings) before using it. - /// - /// The builder only attempts upgrades when an explicit allowlist is supplied. fn try_upgrade(&self, _allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme> { None } diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 7cc7d0322da..942ea78ba0a 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -557,9 +557,7 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc Arc VortexResult<()> { let session = array_session() .with::() @@ -193,7 +199,7 @@ async fn default_strategy_keeps_wide_decimals_compatible( for declaration in EDITION_DECLARATIONS { session.register_edition(declaration)?; } - session.enable_edition(CORE_2026_08_3)?; + session.enable_edition(DEFAULT_CORE_EDITION)?; let array = if use_i256 { DecimalArray::new( @@ -224,7 +230,10 @@ async fn default_strategy_keeps_wide_decimals_compatible( .build(); for disable_editions in [false, true] { - let options = session.write_options().with_strategy(Arc::clone(&strategy)); + let mut options = session.write_options(); + if custom_strategy { + options = options.with_strategy(Arc::clone(&strategy)); + } let options = if disable_editions { options.disable_editions() } else { @@ -241,6 +250,12 @@ async fn default_strategy_keeps_wide_decimals_compatible( .into_array_stream()? .read_all() .await?; + let uses_v2 = actual.depth_first_traversal().any(|array| { + array + .as_opt::() + .is_some_and(|parts| !parts.lower_parts().is_empty()) + }); + assert_eq!(uses_v2, disable_editions && !custom_strategy); assert_arrays_eq!(array, actual, &mut session.create_execution_ctx()); } Ok(()) diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 874d08be306..5d6504dd7c3 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -252,10 +252,7 @@ impl VortexWriteOptions { let strategy = match self.strategy { Some(strategy) => strategy, None => WriteStrategyBuilder::default() - .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_serialized_ids), - ) + .with_btrblocks_builder(BtrBlocksCompressorBuilder::new(allowed_serialized_ids)) .build(), }; let dtype = stream.dtype().clone(); diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 7288fc5d1e3..97c3f4e898d 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -384,12 +384,11 @@ impl PyVortexWriteOptions { .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - let mut compressor = BtrBlocksCompressorBuilder::default(); + let mut compressor = BtrBlocksCompressorBuilder::new(allowed_encodings); if self.use_compact_encodings { compressor = compressor.with_compact(); } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + let strategy = WriteStrategyBuilder::default().with_btrblocks_builder(compressor); let strategy = strategy.build(); current_runtime().block_on(async move { match resolve_store(path, store.map(|x| x.into_inner()))? { diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index ab316982b27..10eb3bbf2be 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -102,12 +102,11 @@ pub async fn exec_convert(session: &VortexSession, flags: ConvertArgs) -> anyhow .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - let mut compressor = BtrBlocksCompressorBuilder::default(); + let mut compressor = BtrBlocksCompressorBuilder::new(allowed_encodings); if matches!(flags.strategy, Strategy::Compact) { compressor = compressor.with_compact(); } - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + let strategy = WriteStrategyBuilder::default().with_btrblocks_builder(compressor); let mut file = File::create(output_path).await?; session diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index 464cce2197b..6cd38268ed5 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -22,6 +22,7 @@ mod tests; pub use vortex_edition::ComponentKind; +pub use vortex_edition::DEFAULT_CORE_EDITION; pub use vortex_edition::EDITION_DECLARATIONS; pub use vortex_edition::EDITION_FAMILIES; pub use vortex_edition::Edition; @@ -47,9 +48,6 @@ use vortex_error::VortexExpect; use vortex_error::vortex_err; use vortex_session::VortexSession; -/// The `core` edition enabled for writing by the default Vortex session. -pub const DEFAULT_CORE_EDITION: EditionId = CORE_2026_08_3; - /// The newest `preview` edition. The default Vortex session registers it but does not enable it. pub const DEFAULT_PREVIEW_EDITION: EditionId = PREVIEW_2026_08_0; From 95e88c9cf3e64ba3f6d1ea04ac52317f56aec940 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 22:26:10 -0400 Subject: [PATCH 5/6] Centralize edition ID resolution and fix scheme configuration checks Resolve static edition array IDs in vortex-edition and document how builder permissions filter schemes registered later. Explicitly permit Delta in the scheme selection and trace tests while preserving their assertions. Fix the decimal upgrade lint, redundant qualifications, and formatting. Signed-off-by: "Matt Katz" --- vortex-bench/src/conversions.rs | 4 +- vortex-btrblocks/src/builder.rs | 40 ++++--------------- vortex-btrblocks/src/schemes/decimal.rs | 10 ++--- .../schemes/integer/scheme_selection_tests.rs | 31 +++++++------- .../schemes/string/scheme_selection_tests.rs | 9 ++--- vortex-btrblocks/src/trace_tests.rs | 12 +++++- vortex-btrblocks/tests/decimal_config.rs | 6 +-- vortex-btrblocks/tests/golden.rs | 3 +- vortex-edition/src/declarations/mod.rs | 16 ++++++++ vortex-edition/src/lib.rs | 1 + vortex-edition/src/tests.rs | 24 +++++++++++ vortex-file/src/tests.rs | 12 +++--- 12 files changed, 93 insertions(+), 75 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index cdfd8c4d171..596aca0cef8 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -247,8 +247,8 @@ fn write_options_for( let mut builder = WriteStrategyBuilder::default(); if matches!(compaction, CompactionStrategy::Compact) { - builder = builder - .with_btrblocks_builder(compressor_builder_for_session(&SESSION).with_compact()); + builder = + builder.with_btrblocks_builder(compressor_builder_for_session(&SESSION).with_compact()); } for name in binary_fields { builder = builder.with_field_writer(FieldPath::from_name(name), no_dict_layout()); diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 5e28eaab404..81eb182d565 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -4,9 +4,8 @@ //! Builder for configuring `BtrBlocksCompressor` instances. use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; -use vortex_edition::ComponentKind; use vortex_edition::DEFAULT_CORE_EDITION; -use vortex_edition::EDITION_DECLARATIONS; +use vortex_edition::array_ids_for_edition; use vortex_utils::aliases::hash_set::HashSet; use crate::AllowedSerializedIds; @@ -71,7 +70,8 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ /// Delta, kept out of [`ALL_SCHEMES`] because it is slower to decompress than the schemes that /// would otherwise win. Callers that want it opt in with -/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme). +/// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) and permit `fastlanes.delta` +/// in [`BtrBlocksCompressorBuilder::new`]. /// /// TODO(robert): Return it to [`ALL_SCHEMES`] once we have scheme filtering. pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); @@ -109,15 +109,10 @@ pub struct BtrBlocksCompressorBuilder { } impl Default for BtrBlocksCompressorBuilder { + /// Uses the default core edition's serialized array IDs. Use [`Self::new`] to permit + /// encodings outside that edition; otherwise, schemes requiring them are omitted at build. fn default() -> Self { - let allowed_serialized_ids = EDITION_DECLARATIONS - .iter() - .filter(|declaration| declaration.edition.id.is_at_or_before(&DEFAULT_CORE_EDITION)) - .flat_map(|declaration| declaration.added) - .filter(|member| member.kind == ComponentKind::Array) - .map(|member| member.component.component_id()) - .collect(); - Self::new(allowed_serialized_ids) + Self::new(array_ids_for_edition(&DEFAULT_CORE_EDITION).collect()) } } @@ -148,6 +143,8 @@ impl BtrBlocksCompressorBuilder { /// /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes /// with the compressor. + /// Schemes with unpermitted outputs are silently omitted during [`Self::build`]; use + /// [`Self::new`] with appropriate IDs to keep encodings outside the default core edition. /// /// # Panics /// @@ -265,9 +262,6 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { use vortex_array::VTable; - use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; - use vortex_edition::EditionSession; - use vortex_error::VortexResult; use vortex_fastlanes::FoR; use super::*; @@ -301,24 +295,6 @@ mod tests { assert_eq!(schemes, ALL_SCHEMES); } - #[test] - fn default_permissions_match_default_edition() -> VortexResult<()> { - let editions = EditionSession::empty(); - for declaration in EDITION_DECLARATIONS { - editions.declare(declaration)?; - } - let allowed: HashSet<_> = editions - .components_in(&DEFAULT_CORE_EDITION, ComponentKind::Array) - .into_iter() - .map(|inclusion| inclusion.component_id) - .collect(); - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.allowed_serialized_ids, allowed); - assert!(allowed.contains(&decimal_byte_parts_v1_id())); - assert!(!allowed.contains(&decimal_byte_parts_v2_id())); - Ok(()) - } - #[test] fn allowing_all_declared_outputs_keeps_every_scheme() { let allowed = ALL_SCHEMES diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 7a3a32b08b4..e50ce4ee5ca 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -94,14 +94,10 @@ impl Scheme for DecimalScheme { } fn try_upgrade(&self, allowed_serialized_ids: &AllowedSerializedIds) -> Option<&dyn Scheme> { - if self.mode == DecimalSchemeMode::V1 + (self.mode == DecimalSchemeMode::V1 && allowed_serialized_ids.contains(&decimal_byte_parts_v1_id()) - && allowed_serialized_ids.contains(&decimal_byte_parts_v2_id()) - { - Some(&DECIMAL_V2) - } else { - None - } + && allowed_serialized_ids.contains(&decimal_byte_parts_v2_id())) + .then_some(&DECIMAL_V2 as &dyn Scheme) } /// Children: msp=0, then up to three lower parts in v2 mode. diff --git a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs index b4726dab9b5..03888f6baf5 100644 --- a/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/integer/scheme_selection_tests.rs @@ -10,6 +10,7 @@ use rand::Rng; use rand::SeedableRng; use rand::rngs::StdRng; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Constant; use vortex_array::arrays::Dict; @@ -19,8 +20,11 @@ use vortex_array::expr::stats::Stat; use vortex_array::expr::stats::StatsProviderExt; use vortex_array::validity::Validity; use vortex_buffer::Buffer; +use vortex_edition::DEFAULT_CORE_EDITION; +use vortex_edition::array_ids_for_edition; use vortex_error::VortexResult; use vortex_fastlanes::BitPacked; +use vortex_fastlanes::Delta; use vortex_fastlanes::FoR; use vortex_runend::RunEnd; use vortex_sequence::Sequence; @@ -32,6 +36,16 @@ use crate::BtrBlocksCompressorBuilder; use crate::DELTA_SCHEME; static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +fn delta_compressor() -> BtrBlocksCompressor { + // Delta is outside the default core edition, so it needs an explicit permission too. + let allowed = array_ids_for_edition(&DEFAULT_CORE_EDITION) + .chain(iter::once(Delta.id())) + .collect(); + BtrBlocksCompressorBuilder::new(allowed) + .with_new_scheme(&DELTA_SCHEME) + .build() +} + #[test] fn test_constant_compressed() -> VortexResult<()> { let values: Vec = iter::repeat_n(42, 100).collect(); @@ -164,7 +178,6 @@ fn test_rle_compressed() -> VortexResult<()> { fn test_delta_compressed() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); use vortex_array::assert_arrays_eq; - use vortex_fastlanes::Delta; let mut rng = StdRng::seed_from_u64(7u64); let mut value = 500_000i32; @@ -176,9 +189,7 @@ fn test_delta_compressed() -> VortexResult<()> { .collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build(); + let btr = delta_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -204,7 +215,6 @@ fn test_delta_compressed() -> VortexResult<()> { fn test_delta_compressed_unaligned_length() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); use vortex_array::assert_arrays_eq; - use vortex_fastlanes::Delta; let mut rng = StdRng::seed_from_u64(7u64); let mut value = 500_000i32; @@ -216,9 +226,7 @@ fn test_delta_compressed_unaligned_length() -> VortexResult<()> { .collect(); let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable); - let btr = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build(); + let btr = delta_compressor(); let compressed = btr.compress( &array.clone().into_array(), &mut SESSION.create_execution_ctx(), @@ -239,15 +247,12 @@ fn test_delta_compressed_unaligned_length() -> VortexResult<()> { fn test_delta_nullable_unaligned_sum() -> VortexResult<()> { use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::assert_arrays_eq; - use vortex_fastlanes::Delta; let mut ctx = SESSION.create_execution_ctx(); let array = PrimitiveArray::from_option_iter(iter::once(None).chain((1i32..=100_000).map(Some))); - let btr = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&DELTA_SCHEME) - .build(); + let btr = delta_compressor(); let compressed = btr.compress(&array.clone().into_array(), &mut ctx)?; assert!( compressed.is::(), @@ -266,8 +271,6 @@ fn test_delta_nullable_unaligned_sum() -> VortexResult<()> { /// Returns true if any `Delta` array appears below an ancestor `Delta` in the tree. fn has_nested_delta(array: &vortex_array::ArrayRef, under_delta: bool) -> bool { - use vortex_fastlanes::Delta; - let is_delta = array.is::(); if is_delta && under_delta { return true; diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index 55ccf0159cb..29c36ef4f9a 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -104,11 +104,10 @@ fn test_fsst_in_default_scheme_list() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressor = BtrBlocksCompressorBuilder::empty( - FSSTScheme.produced_encodings().into_iter().collect(), - ) - .with_new_scheme(&FSSTScheme) - .build(); + let compressor = + BtrBlocksCompressorBuilder::empty(FSSTScheme.produced_encodings().into_iter().collect()) + .with_new_scheme(&FSSTScheme) + .build(); let compressed = compressor.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index e23e4ef0244..39058b3b769 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -11,6 +11,7 @@ //! rules and execute kernels fire for the scan operations TPC-H queries perform over those //! encodings. +use std::iter; use std::sync::LazyLock; use arrow_array::RecordBatch; @@ -22,6 +23,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DictArray; use vortex_array::arrays::FilterArray; @@ -49,7 +51,10 @@ use vortex_array::session::ArraySessionExt; use vortex_array::test_harness::trace::Traced; use vortex_array::test_harness::trace::trace_op; use vortex_arrow::ArrowSessionExt; +use vortex_edition::DEFAULT_CORE_EDITION; +use vortex_edition::array_ids_for_edition; use vortex_error::VortexResult; +use vortex_fastlanes::Delta; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -125,9 +130,12 @@ fn lineitem() -> VortexResult { .from_arrow_record_batch(batch, &schema) } -/// Delta is opt-in, and these traces cover the delta-encoded FSST offsets, so enable it here. +/// These traces cover delta-encoded FSST offsets, so permit and register the opt-in Delta scheme. fn compressed_lineitem() -> VortexResult { - BtrBlocksCompressorBuilder::default() + let allowed = array_ids_for_edition(&DEFAULT_CORE_EDITION) + .chain(iter::once(Delta.id())) + .collect(); + BtrBlocksCompressorBuilder::new(allowed) .with_new_scheme(&DELTA_SCHEME) .build() .compress(&lineitem()?, &mut execution_ctx()) diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs index 09224d91c5d..321c5641927 100644 --- a/vortex-btrblocks/tests/decimal_config.rs +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -115,11 +115,7 @@ fn explicit_decimal_modes_only_upgrade( allowed.insert(decimal_byte_parts_v2_id()); } let builder = BtrBlocksCompressorBuilder::empty(allowed).with_new_scheme(scheme); - let mode = if initial_v2 && !allowed_v2 { - None - } else { - Some(allowed_v2) - }; + let mode = (!initial_v2 || allowed_v2).then_some(allowed_v2); let expected = match mode { Some(true) if wide => Some(decimal_byte_parts_v2_id()), Some(_) if !wide => Some(decimal_byte_parts_v1_id()), diff --git a/vortex-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index 9a105f32467..c0539a7407f 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -447,6 +447,7 @@ fn golden_compact() -> VortexResult<()> { let session = edition_session(&[CORE_2026_08_3])?; vortex_zstd::initialize(&session); session.enable_edition(vortex_zstd::editions::ZSTD_2026_02)?; - let compressor = without_onpair(compressor_builder_for_session(&session).with_compact()).build(); + let compressor = + without_onpair(compressor_builder_for_session(&session).with_compact()).build(); golden_corpus_snapshots("compact", &compressor) } diff --git a/vortex-edition/src/declarations/mod.rs b/vortex-edition/src/declarations/mod.rs index 78dbe77f610..a402176ad1d 100644 --- a/vortex-edition/src/declarations/mod.rs +++ b/vortex-edition/src/declarations/mod.rs @@ -14,6 +14,9 @@ pub mod core; pub mod preview; +use vortex_session::registry::Id; + +use crate::ComponentKind; use crate::EditionDeclaration; use crate::EditionFamily; use crate::EditionId; @@ -35,3 +38,16 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &core::v2026_08_3::DECLARATION, &preview::v2026_08::DECLARATION, ]; + +/// Returns the serialized array IDs declared in `edition` or earlier editions of its family. +/// +/// Resolves the static first-party [`EDITION_DECLARATIONS`]. Components registered separately +/// on a session are not included; use [`crate::EditionSession::components_in`] for those. +pub fn array_ids_for_edition(edition: &EditionId) -> impl Iterator + '_ { + EDITION_DECLARATIONS + .iter() + .filter(move |declaration| declaration.edition.id.is_at_or_before(edition)) + .flat_map(|declaration| declaration.added) + .filter(|member| member.kind == ComponentKind::Array) + .map(|member| member.component.component_id()) +} diff --git a/vortex-edition/src/lib.rs b/vortex-edition/src/lib.rs index 69034f93a3d..e882f4263eb 100644 --- a/vortex-edition/src/lib.rs +++ b/vortex-edition/src/lib.rs @@ -45,6 +45,7 @@ use std::fmt::Formatter; pub use declarations::DEFAULT_CORE_EDITION; pub use declarations::EDITION_DECLARATIONS; pub use declarations::EDITION_FAMILIES; +pub use declarations::array_ids_for_edition; pub use session::EditionSession; pub use session::EditionSessionExt; pub use session::EnabledEditions; diff --git a/vortex-edition/src/tests.rs b/vortex-edition/src/tests.rs index 63cf003d11f..58bd30472f0 100644 --- a/vortex-edition/src/tests.rs +++ b/vortex-edition/src/tests.rs @@ -4,6 +4,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_session::VortexSession; +use vortex_session::registry::Id; use crate::ComponentKind; use crate::Edition; @@ -15,6 +16,29 @@ use crate::EditionMember; use crate::EditionSession; use crate::EditionSessionExt; use crate::EnabledEditions; +use crate::array_ids_for_edition; +use crate::declarations::core::CORE_2025_05_0; +use crate::declarations::core::CORE_2025_06_0; + +#[test] +fn static_array_ids_inherit_only_within_the_edition_family() { + let first_ids: Vec<_> = array_ids_for_edition(&CORE_2025_05_0).collect(); + let second_ids: Vec<_> = array_ids_for_edition(&CORE_2025_06_0).collect(); + let first: Vec<_> = first_ids.iter().map(Id::as_str).collect(); + let second: Vec<_> = second_ids.iter().map(Id::as_str).collect(); + + assert!(first.contains(&"vortex.primitive")); + assert!(!first.contains(&"vortex.sequence")); + assert!(first.iter().all(|id| second.contains(id))); + assert!(second.contains(&"vortex.sequence")); + assert!(!second.contains(&"vortex.flat")); + assert!(!second.contains(&"vortex.date")); + assert!( + array_ids_for_edition(&EditionId::new("other", 2025, 6, 0)) + .next() + .is_none() + ); +} static TEST_FAMILY: EditionFamily = EditionFamily { name: "test", diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 6fcdc9017d6..0175bbd0d8d 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1962,9 +1962,7 @@ async fn write_read_roundtrip_with_layout( array: ArrayRef, use_list_layout: bool, ) -> VortexResult { - let strategy = crate::strategy::WriteStrategyBuilder::default() - .with_list_layout() - .build(); + let strategy = WriteStrategyBuilder::default().with_list_layout().build(); let mut buf = ByteBufferMut::empty(); if use_list_layout { SESSION @@ -2348,7 +2346,7 @@ async fn timestamp_unit_mismatch_errors_with_constant_children() .into_array(); let temporal = TemporalArray::new_timestamp(ts_array, TimeUnit::Milliseconds, None); - let strategy = crate::strategy::WriteStrategyBuilder::default() + let strategy = WriteStrategyBuilder::default() .with_compressor(compressor) .build(); @@ -2666,7 +2664,7 @@ async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { let mut buf = ByteBufferMut::empty(); let summary = SESSION .write_options() - .with_strategy(crate::strategy::WriteStrategyBuilder::default().build()) + .with_strategy(WriteStrategyBuilder::default().build()) .write(&mut buf, strings.clone().to_array_stream()) .await?; assert!( @@ -2680,7 +2678,7 @@ async fn dict_probe_honours_configured_compressor() -> VortexResult<()> { let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() + WriteStrategyBuilder::default() .with_btrblocks_builder(no_string_dict) .build(), ) @@ -2710,7 +2708,7 @@ async fn probe_compressor_override_is_independent() -> VortexResult<()> { let summary = SESSION .write_options() .with_strategy( - crate::strategy::WriteStrategyBuilder::default() + WriteStrategyBuilder::default() .with_probe_compressor(probe_without_dict) .build(), ) From 0c415a71e5ecfaaa3790a1a97470a8a8b3ba5344 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 22 Sep 2026 23:12:56 -0400 Subject: [PATCH 6/6] Add builder methods to replace and extend permitted encodings Make empty() start without registered schemes or serialized ID permissions. Allow callers to set or extend permissions before build, while keeping the default core edition permissions in default(). Update callers and cover final permission filtering, decimal upgrades, and CUDA preset ordering. Signed-off-by: "Matt Katz" --- vortex-btrblocks/src/builder.rs | 111 +++++++++++++++--- .../schemes/string/scheme_selection_tests.rs | 8 +- vortex-btrblocks/tests/decimal_config.rs | 43 ++++++- vortex-cuda/src/layout.rs | 2 +- 4 files changed, 137 insertions(+), 27 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 81eb182d565..61e56951874 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -3,6 +3,7 @@ //! Builder for configuring `BtrBlocksCompressor` instances. +use vortex_array::ArrayId; use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; use vortex_edition::DEFAULT_CORE_EDITION; use vortex_edition::array_ids_for_edition; @@ -21,10 +22,14 @@ use crate::schemes::integer; use crate::schemes::string; use crate::schemes::temporal; -/// All default compression schemes, including Decimal v1. +/// All compression schemes. /// /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. +/// +/// If a scheme can be configured to support different editions like +/// [`DecimalScheme`](crate::schemes::decimal), put oldest version here to defer upgrading +/// to newest supported version when compressor is built. pub const ALL_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// // Integer schemes. @@ -63,6 +68,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ &binary::BinaryDictScheme, &binary::VarBinScheme, // Decimal schemes. + // Use v1 by default and let builder upgrade to v2 if permitted by edition. &decimal::DecimalScheme::v1(), // Temporal schemes. &temporal::TemporalScheme, @@ -71,7 +77,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ /// Delta, kept out of [`ALL_SCHEMES`] because it is slower to decompress than the schemes that /// would otherwise win. Callers that want it opt in with /// [`with_new_scheme`](BtrBlocksCompressorBuilder::with_new_scheme) and permit `fastlanes.delta` -/// in [`BtrBlocksCompressorBuilder::new`]. +/// with [`BtrBlocksCompressorBuilder::allow_encodings`]. /// /// TODO(robert): Return it to [`ALL_SCHEMES`] once we have scheme filtering. pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); @@ -83,10 +89,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. /// -/// [`Self::new`] takes the permitted serialized IDs once. During [`Self::build`], these -/// permissions allow scheme upgrades and filter all registered schemes. The default builder -/// permits the array IDs in [`DEFAULT_CORE_EDITION`]. Decimal defaults to v1 and upgrades to v2 -/// when both serialized IDs are permitted. +/// [`Self::new`] takes the initial permitted serialized IDs. [`Self::set_allowed_encodings`] +/// replaces these permissions, and [`Self::allow_encodings`] extends them. During [`Self::build`], +/// the final permissions allow scheme upgrades and filter all registered schemes. The default +/// builder permits the array IDs in [`DEFAULT_CORE_EDITION`]. Decimal defaults to v1 and upgrades +/// to v2 when both serialized IDs are permitted. /// /// # Examples /// @@ -94,7 +101,8 @@ pub static DELTA_SCHEME: integer::DeltaScheme = integer::DeltaScheme::new(1.25); /// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme, SchemeExt}; /// use vortex_btrblocks::schemes::integer::IntDictScheme; /// -/// // Default compressor restricted to the default core edition. +/// // Default compressor with all schemes in ALL_SCHEMES, restricted to the +/// // default core edition. /// let compressor = BtrBlocksCompressorBuilder::default().build(); /// /// // Remove specific schemes. @@ -109,8 +117,8 @@ pub struct BtrBlocksCompressorBuilder { } impl Default for BtrBlocksCompressorBuilder { - /// Uses the default core edition's serialized array IDs. Use [`Self::new`] to permit - /// encodings outside that edition; otherwise, schemes requiring them are omitted at build. + /// Uses the default core edition's serialized array IDs. Use [`Self::allow_encodings`] to + /// permit additional encodings; otherwise, schemes requiring them are omitted at build. fn default() -> Self { Self::new(array_ids_for_edition(&DEFAULT_CORE_EDITION).collect()) } @@ -128,23 +136,51 @@ impl BtrBlocksCompressorBuilder { } } - /// Creates a builder with no schemes registered. + /// Creates a builder with no registered schemes and no permitted serialized IDs. /// /// Useful when the caller wants explicit, scheme-by-scheme control over the compressor. - /// The supplied serialized ID permissions apply to every scheme registered later. - pub fn empty(allowed_serialized_ids: AllowedSerializedIds) -> Self { + /// Register schemes with [`Self::with_new_scheme`] and permit their serialized IDs with + /// [`Self::allow_encodings`]. + /// + /// ```rust + /// use vortex_btrblocks::{BtrBlocksCompressorBuilder, Scheme}; + /// use vortex_btrblocks::schemes::integer::FoRScheme; + /// + /// let compressor = BtrBlocksCompressorBuilder::empty() + /// .allow_encodings(FoRScheme.produced_encodings()) + /// .with_new_scheme(&FoRScheme) + /// .build(); + /// ``` + pub fn empty() -> Self { Self { schemes: Vec::new(), - allowed_serialized_ids, + allowed_serialized_ids: HashSet::new(), } } + /// Replaces the permitted serialized IDs, including any default edition permissions. + /// + /// An empty iterator permits no serialized IDs. The final permissions apply to every + /// registered scheme during [`Self::build`]. + pub fn set_allowed_encodings(mut self, ids: impl IntoIterator) -> Self { + self.allowed_serialized_ids = ids.into_iter().collect(); + self + } + + /// Adds permitted serialized IDs while preserving existing permissions. + /// + /// The final permissions apply to every registered scheme during [`Self::build`]. + pub fn allow_encodings(mut self, ids: impl IntoIterator) -> Self { + self.allowed_serialized_ids.extend(ids); + self + } + /// Adds an external compression scheme not in [`ALL_SCHEMES`]. /// /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes /// with the compressor. /// Schemes with unpermitted outputs are silently omitted during [`Self::build`]; use - /// [`Self::new`] with appropriate IDs to keep encodings outside the default core edition. + /// [`Self::allow_encodings`] to permit outputs outside the default core edition. /// /// # Panics /// @@ -188,11 +224,13 @@ impl BtrBlocksCompressorBuilder { /// /// 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, - /// but belongs to the opt-in `zstd` edition. The permissions supplied to [`Self::new`] - /// determine which schemes survive. + /// but belongs to the opt-in `zstd` edition. The final permissions determine which schemes + /// survive. /// The decimal v2 serialized ID is excluded because CUDA does not support lower decimal /// parts. This prevents v1 schemes from upgrading and filters out explicitly registered v2 /// schemes, including Decimal schemes registered after this call. + /// Apply this preset after changing permissions: [`Self::set_allowed_encodings`] and + /// [`Self::allow_encodings`] can re-enable decimal v2 if called afterwards. /// /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. @@ -262,14 +300,22 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { use vortex_array::VTable; + use vortex_fastlanes::Delta; use vortex_fastlanes::FoR; use super::*; #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(HashSet::new()); + fn empty_starts_with_no_schemes_or_permissions() { + let builder = BtrBlocksCompressorBuilder::empty(); assert!(builder.schemes.is_empty()); + assert!(builder.allowed_serialized_ids.is_empty()); + + let builder = builder.with_new_scheme(&integer::FoRScheme); + assert!(builder.clone().configured_schemes().is_empty()); + let schemes = builder.allow_encodings([FoR.id()]).configured_schemes(); + assert_eq!(schemes.len(), 1); + assert_eq!(schemes[0].id(), integer::FoRScheme.id()); } #[test] @@ -289,6 +335,31 @@ mod tests { assert!(none.is_empty()); } + #[test] + fn set_allowed_encodings_replaces_permissions() { + let builder = BtrBlocksCompressorBuilder::default().set_allowed_encodings([FoR.id()]); + let schemes = builder.clone().configured_schemes(); + assert_eq!(schemes.len(), 1); + assert_eq!(schemes[0].id(), integer::FoRScheme.id()); + assert!( + builder + .set_allowed_encodings([]) + .configured_schemes() + .is_empty() + ); + } + + #[test] + fn allow_encodings_keeps_defaults_and_enables_registered_schemes() { + let builder = BtrBlocksCompressorBuilder::default().with_new_scheme(&DELTA_SCHEME); + assert_eq!(builder.clone().configured_schemes(), ALL_SCHEMES); + + let schemes = builder.allow_encodings([Delta.id()]).configured_schemes(); + assert_eq!(&schemes[..ALL_SCHEMES.len()], ALL_SCHEMES); + assert_eq!(schemes.len(), ALL_SCHEMES.len() + 1); + assert_eq!(schemes[ALL_SCHEMES.len()].id(), DELTA_SCHEME.id()); + } + #[test] fn default_configuration_preserves_scheme_order() { let schemes = BtrBlocksCompressorBuilder::default().configured_schemes(); @@ -311,7 +382,9 @@ mod tests { for id in scheme.produced_encodings() { let mut allowed: HashSet<_> = scheme.produced_encodings().into_iter().collect(); allowed.remove(&id); - let builder = BtrBlocksCompressorBuilder::empty(allowed).with_new_scheme(*scheme); + let builder = BtrBlocksCompressorBuilder::empty() + .allow_encodings(allowed) + .with_new_scheme(*scheme); assert!( builder.configured_schemes().is_empty(), "{} requires {id}", diff --git a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs index 29c36ef4f9a..60a23d488a3 100644 --- a/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs +++ b/vortex-btrblocks/src/schemes/string/scheme_selection_tests.rs @@ -104,10 +104,10 @@ fn test_fsst_in_default_scheme_list() -> VortexResult<()> { let array = VarBinViewArray::from_iter(strings, DType::Utf8(Nullability::NonNullable)); let array_ref = array.into_array(); - let compressor = - BtrBlocksCompressorBuilder::empty(FSSTScheme.produced_encodings().into_iter().collect()) - .with_new_scheme(&FSSTScheme) - .build(); + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_encodings(FSSTScheme.produced_encodings()) + .with_new_scheme(&FSSTScheme) + .build(); let compressed = compressor.compress(&array_ref, &mut SESSION.create_execution_ctx())?; assert!( compressed.is::(), diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs index 321c5641927..e20b52aca65 100644 --- a/vortex-btrblocks/tests/decimal_config.rs +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -103,6 +103,40 @@ fn decimal_mode_follows_permissions( assert_decimal_output(builder, wide, expected) } +#[rstest] +fn decimal_mode_uses_final_permissions(#[values(false, true)] wide: bool) -> VortexResult<()> { + let builder = + BtrBlocksCompressorBuilder::default().allow_encodings([decimal_byte_parts_v2_id()]); + assert_decimal_output( + builder.clone(), + wide, + Some(if wide { + decimal_byte_parts_v2_id() + } else { + decimal_byte_parts_v1_id() + }), + )?; + + assert_decimal_output( + builder.set_allowed_encodings([decimal_byte_parts_v1_id()]), + wide, + (!wide).then(decimal_byte_parts_v1_id), + ) +} + +#[rstest] +fn permissions_can_reenable_decimal_v2_after_cuda( + #[values(false, true)] replace: bool, +) -> VortexResult<()> { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + let builder = if replace { + builder.set_allowed_encodings([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]) + } else { + builder.allow_encodings([decimal_byte_parts_v2_id()]) + }; + assert_decimal_output(builder, true, Some(decimal_byte_parts_v2_id())) +} + #[rstest] fn explicit_decimal_modes_only_upgrade( #[values(false, true)] allowed_v2: bool, @@ -114,7 +148,9 @@ fn explicit_decimal_modes_only_upgrade( if allowed_v2 { allowed.insert(decimal_byte_parts_v2_id()); } - let builder = BtrBlocksCompressorBuilder::empty(allowed).with_new_scheme(scheme); + let builder = BtrBlocksCompressorBuilder::empty() + .allow_encodings(allowed) + .with_new_scheme(scheme); let mode = (!initial_v2 || allowed_v2).then_some(allowed_v2); let expected = match mode { Some(true) if wide => Some(decimal_byte_parts_v2_id()), @@ -140,7 +176,7 @@ fn cuda_never_uses_decimal_v2( ) -> VortexResult<()> { let allowed = HashSet::from([decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()]); let scheme: &'static dyn Scheme = if initial_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; - let mut builder = BtrBlocksCompressorBuilder::empty(allowed); + let mut builder = BtrBlocksCompressorBuilder::empty().allow_encodings(allowed); if !register_later { builder = builder.with_new_scheme(scheme); } @@ -200,7 +236,8 @@ fn wide_decimal_parts_roundtrip( allowed.extend(FoRScheme.produced_encodings()); allowed.extend(BitPackingScheme.produced_encodings()); } - let compressor = BtrBlocksCompressorBuilder::empty(allowed) + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_encodings(allowed) .with_new_scheme(&DECIMAL_V2) .with_new_scheme(&FoRScheme) .with_new_scheme(&BitPackingScheme) diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 942ea78ba0a..732527530be 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -566,7 +566,7 @@ pub fn cuda_write_strategy(session: &VortexSession, block_rows: usize) -> Arc