diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder/mod.rs similarity index 68% rename from vortex-btrblocks/src/builder.rs rename to vortex-btrblocks/src/builder/mod.rs index fe8072d5e66..49cd98213af 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -18,7 +18,7 @@ use crate::schemes::integer; use crate::schemes::string; use crate::schemes::temporal; -/// All available compression schemes. +/// The newest versions of all available compression schemes. /// /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. @@ -91,12 +91,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + allowed_serialized_ids: Option>, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + allowed_serialized_ids: None, } } } @@ -108,13 +110,14 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_serialized_ids: None, } } /// 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. + /// with the compressor. Register only the newest version of a scheme. /// /// # Panics /// @@ -198,122 +201,61 @@ impl BtrBlocksCompressorBuilder { } /// Removes the specified compression schemes by their [`SchemeId`]. + /// + /// An ID anywhere in a registered predecessor chain removes the entire chain. + /// + /// # Panics + /// + /// Panics if a traversed predecessor chain contains a cycle. pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { let ids: HashSet<_> = ids.into_iter().collect(); - self.schemes.retain(|s| !ids.contains(&s.id())); + self.schemes.retain(|scheme| { + let mut seen = HashSet::new(); + let mut candidate = Some(*scheme); + while let Some(version) = candidate { + assert!( + seen.insert(version.id()), + "cycle in scheme predecessor chain" + ); + if ids.contains(&version.id()) { + return false; + } + candidate = version.predecessor(); + } + true + }); self } - /// Retains only schemes whose produced encodings all belong to `allowed`. + /// Restricts compression to the serialized IDs in `allowed`, intersecting with any earlier + /// call. /// - /// The file writer uses this to restrict compression to the encodings of its configured - /// editions. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + /// At build time, each scheme is replaced by the newest version in its predecessor chain + /// whose [`produced_encodings`](Scheme::produced_encodings) are all permitted. + /// Schemes with no eligible version are removed. This also applies to schemes added after + /// this call. The file writer passes the serialized IDs its enabled editions permit. + pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { + let allowed: HashSet = match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(allowed).copied().collect(), + None => allowed.clone(), + }; + self.allowed_serialized_ids = Some(allowed); self } /// Builds the configured [`BtrBlocksCompressor`]. + /// + /// # Panics + /// + /// Panics if predecessor chains contain a cycle or share a scheme ID. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + let compressor = CascadingCompressor::new(self.schemes); + BtrBlocksCompressor(match self.allowed_serialized_ids { + Some(allowed) => compressor.with_allowed_serialized_ids(allowed), + None => compressor, + }) } } #[cfg(test)] -mod tests { - use vortex_array::VTable; - use vortex_fastlanes::FoR; - - use super::*; - - #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); - assert!(builder.schemes.is_empty()); - } - - #[test] - fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[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 none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) - ); - } - - /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. - #[test] - fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - for excluded in [ - integer::SparseScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - ] { - assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), - "{excluded} should be excluded" - ); - } - } - - #[test] - fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); - #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); - } - - #[test] - #[cfg(feature = "pco")] - fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); - for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); - } - } -} +mod tests; diff --git a/vortex-btrblocks/src/builder/tests.rs b/vortex-btrblocks/src/builder/tests.rs new file mode 100644 index 00000000000..4310b9416c8 --- /dev/null +++ b/vortex-btrblocks/src/builder/tests.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::VarBin; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; +use vortex_fastlanes::FoR; +use vortex_fsst::FSST; +use vortex_session::registry::CachedId; + +use super::*; +use crate::ArrayAndStats; +use crate::CompressorContext; + +#[test] +fn empty_starts_with_no_schemes() { + assert!(BtrBlocksCompressorBuilder::empty().schemes.is_empty()); +} + +#[test] +fn default_includes_all_schemes() { + assert_eq!( + BtrBlocksCompressorBuilder::default().schemes.len(), + ALL_SCHEMES.len() + ); +} + +#[test] +fn allowed_serialized_ids_filter_schemes_at_build() { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .build(); + for scheme in ALL_SCHEMES { + assert_eq!( + compressor.has_scheme(scheme.id()), + scheme.id() == integer::FoRScheme.id() + ); + } +} + +#[test] +fn allowing_all_declared_outputs_keeps_every_scheme() { + let allowed = ALL_SCHEMES + .iter() + .flat_map(|s| s.produced_encodings()) + .collect(); + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed) + .build(); + for scheme in ALL_SCHEMES { + assert!(compressor.has_scheme(scheme.id())); + } +} + +#[rstest] +#[case::neither(vec![], false)] +#[case::fsst_only(vec![FSST.id()], false)] +#[case::varbin_only(vec![VarBin.id()], false)] +#[case::both(vec![FSST.id(), VarBin.id()], true)] +fn all_required_outputs_must_be_allowed(#[case] allowed: Vec, #[case] expected: bool) { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed.into_iter().collect()) + .build(); + assert_eq!(compressor.has_scheme(string::FSSTScheme.id()), expected); +} + +#[rstest] +#[case::forbidden(HashSet::new(), false)] +#[case::permitted(HashSet::from([FoR.id()]), true)] +fn restriction_applies_to_schemes_added_later( + #[case] allowed: HashSet, + #[case] expected: bool, +) { + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_serialized_ids(&allowed) + .with_new_scheme(&integer::FoRScheme) + .build(); + assert_eq!(compressor.has_scheme(integer::FoRScheme.id()), expected); +} + +#[test] +fn repeated_restrictions_intersect() { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id(), FSST.id()])) + .allow_serialized_ids(&HashSet::from([FSST.id(), VarBin.id()])) + .build(); + assert!(!compressor.has_scheme(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(string::FSSTScheme.id())); +} + +#[test] +fn cuda_compatible_excludes_alprd() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + !builder + .schemes + .iter() + .any(|s| s.id() == float::ALPRDScheme.id()) + ); +} + +/// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. +#[test] +fn cuda_compatible_excludes_every_sparse_scheme() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + for excluded in [ + integer::SparseScheme.id(), + float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), + ] { + assert!( + !builder.schemes.iter().any(|s| s.id() == excluded), + "{excluded} should be excluded" + ); + } +} + +#[test] +fn cuda_compatible_uses_fsst_for_strings() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::FSSTScheme.id()) + ); + #[cfg(feature = "zstd")] + assert!( + !builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::ZstdScheme.id()) + ); +} + +#[test] +#[cfg(feature = "pco")] +fn cuda_compatible_excludes_pco() { + let builder = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&integer::PcoScheme) + .with_new_scheme(&float::PcoScheme) + .only_cuda_compatible(); + for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { + assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); + } +} + +static FOR_V2_ID: CachedId = CachedId::new("test.for_v2"); + +#[derive(Debug)] +struct NewFoRScheme; + +impl Scheme for NewFoRScheme { + fn scheme_name(&self) -> &'static str { + "test.for_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + integer::FoRScheme.matches(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![*FOR_V2_ID] + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + Some(&integer::FoRScheme) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +#[test] +fn restrictions_select_predecessors_of_schemes_added_later() { + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .with_new_scheme(&NewFoRScheme) + .build(); + assert!(compressor.has_scheme(integer::FoRScheme.id())); + assert!(compressor.has_scheme_family(NewFoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); +} + +#[rstest] +#[case::old(integer::FoRScheme.id())] +#[case::new(NewFoRScheme.id())] +fn excluding_any_version_removes_the_chain(#[case] excluded: SchemeId) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&NewFoRScheme) + .exclude_schemes([excluded]) + .build(); + assert!(!compressor.has_scheme_family(integer::FoRScheme.id())); + assert!(!compressor.has_scheme_family(NewFoRScheme.id())); +} diff --git a/vortex-btrblocks/tests/scheme_versions.rs b/vortex-btrblocks/tests/scheme_versions.rs new file mode 100644 index 00000000000..ff4b67961c4 --- /dev/null +++ b/vortex-btrblocks/tests/scheme_versions.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg(feature = "unstable_encodings")] + +#[cfg(test)] +mod tests { + use rstest::rstest; + 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_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_btrblocks::ArrayAndStats; + use vortex_btrblocks::CascadingCompressor; + use vortex_btrblocks::CompressorContext; + use vortex_btrblocks::Scheme; + use vortex_btrblocks::SchemeExt; + use vortex_btrblocks::schemes::integer::DeltaScheme; + use vortex_btrblocks::schemes::integer::IntRLEScheme; + use vortex_compressor::scheme::CompressionEstimate; + use vortex_compressor::scheme::EstimateVerdict; + use vortex_error::VortexResult; + use vortex_fastlanes::Delta; + use vortex_fastlanes::RLE; + use vortex_session::registry::CachedId; + + static DELTA_V2_ID: CachedId = CachedId::new("test.delta_v2"); + static DELTA_V1: DeltaScheme = DeltaScheme::new(1.25); + + #[derive(Debug)] + struct DeltaV2; + + impl Scheme for DeltaV2 { + fn scheme_name(&self) -> &'static str { + "test.delta_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + DELTA_V1.matches(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![*DELTA_V2_ID] + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + Some(&DELTA_V1) + } + + fn num_children(&self) -> usize { + 2 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } + } + + #[rstest] + #[case::predecessor(Delta.id(), true)] + #[case::replacement(*DELTA_V2_ID, false)] + fn rle_respects_selected_delta_version( + #[case] allowed_delta: ArrayId, + #[case] expect_delta: bool, + ) -> VortexResult<()> { + let session = array_session(); + vortex_fastlanes::initialize(&session); + let compressor = CascadingCompressor::new(vec![&IntRLEScheme, &DeltaV2]) + .with_allowed_serialized_ids([RLE.id(), allowed_delta].into_iter().collect()); + assert!(compressor.has_scheme_family(DELTA_V1.id())); + let array = PrimitiveArray::from_iter((0..65_536u32).map(|i| (i / 64) % 100)).into_array(); + let mut ctx = session.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + assert_eq!(compressed.encoding_id(), RLE.id()); + let has_delta = compressed + .depth_first_traversal() + .any(|array| array.encoding_id() == Delta.id()); + assert_eq!(has_delta, expect_delta); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) + } +} diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..ecfd3c2c542 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -59,7 +59,7 @@ impl CascadingCompressor { let canonical = array.clone().execute::(exec_ctx)?.0; let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + let compressed = self.compress_canonical(compact, self.root_context(), exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); @@ -93,7 +93,7 @@ impl CascadingCompressor { let child_ctx = parent_ctx .clone() - .descend_with_scheme(parent_id, child_index); + .descend_with_scheme(self.resolve_scheme_id(parent_id), child_index); self.compress_canonical(compact, child_ctx, exec_ctx) } diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..3f88debbc24 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,8 +9,13 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; +use crate::scheme::CompressorContext; use crate::scheme::DescendantExclusion; use crate::scheme::Scheme; use crate::scheme::SchemeExt; @@ -46,13 +51,39 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// Maps every registered version to the version selected for compression. + scheme_aliases: HashMap, + + /// Configuration only: retained so repeated restrictions intersect exactly. + allowed_serialized_ids: Option>, } impl CascadingCompressor { /// Creates a new compressor with the given schemes. /// + /// Register only the newest version of each scheme. Predecessor IDs are aliases for the + /// selected version in exclusions and [`has_scheme_family`](Self::has_scheme_family) checks. /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. + /// + /// # Panics + /// + /// Panics if predecessor chains contain a cycle or share a scheme ID, including when multiple + /// versions of the same scheme are registered separately. pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + let mut scheme_aliases = HashMap::new(); + for &scheme in &schemes { + let mut candidate = Some(scheme); + while let Some(version) = candidate { + assert!( + scheme_aliases.insert(version.id(), scheme.id()).is_none(), + "scheme {} appears more than once in the registered predecessor chains", + version.id(), + ); + candidate = version.predecessor(); + } + } + // Root exclusion: exclude IntDict from list/listview offsets (monotonically // increasing data where dictionary encoding is wasteful). let root_exclusions = vec![DescendantExclusion { @@ -63,14 +94,76 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + scheme_aliases, + allowed_serialized_ids: None, } } - /// Returns whether the compressor was configured with `scheme`. + /// Selects the newest eligible version of each scheme, intersecting with any earlier call. + /// + /// A version is eligible only when all of its [`Scheme::produced_encodings`] are allowed. + /// Otherwise its predecessors are tried in order; the scheme is removed if none is eligible. + /// Selection preserves registration order and happens before any compression or estimation. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + let allowed = match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(&allowed).copied().collect(), + None => allowed, + }; + let mut replacements = HashMap::new(); + self.schemes = self + .schemes + .into_iter() + .filter_map(|scheme| { + let mut candidate = Some(scheme); + while let Some(version) = candidate { + if version + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + { + replacements.insert(scheme.id(), version.id()); + return Some(version); + } + candidate = version.predecessor(); + } + None + }) + .collect(); + self.scheme_aliases.retain(|_, selected| { + if let Some(replacement) = replacements.get(selected) { + *selected = *replacement; + true + } else { + false + } + }); + self.allowed_serialized_ids = Some(allowed); + self + } + + /// The context a compress call starts from. + pub(crate) fn root_context(&self) -> CompressorContext { + CompressorContext::new() + } + + /// Returns whether any version in the scheme's family is enabled. + /// + /// A family is a registered scheme and its predecessor chain. `scheme` may name any version + /// in that chain. Use [`Self::has_scheme`] to check the exact selected version. + pub fn has_scheme_family(&self, scheme: SchemeId) -> bool { + self.scheme_aliases.contains_key(&scheme) + } + + /// Returns whether this exact scheme version is selected for compression. + /// + /// Use this before invoking a specific implementation directly. pub fn has_scheme(&self, scheme: SchemeId) -> bool { - self.schemes - .iter() - .any(|candidate| candidate.id() == scheme) + self.scheme_aliases.get(&scheme) == Some(&scheme) + } + + /// Resolves a registered version to the selected version, leaving unknown IDs unchanged. + fn resolve_scheme_id(&self, scheme: SchemeId) -> SchemeId { + self.scheme_aliases.get(&scheme).copied().unwrap_or(scheme) } } @@ -78,3 +171,6 @@ impl CascadingCompressor { #[cfg(test)] mod tests; + +#[cfg(test)] +mod version_tests; diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..c492729f77c 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -152,10 +152,9 @@ impl CascadingCompressor { // The root entry is always first in the history (if present). Check if the root has // excluded us. if let Some((_, child_idx)) = iter.next_if(|&(sid, _)| sid == ROOT_SCHEME_ID) - && self - .root_exclusions - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + && self.root_exclusions.iter().any(|rule| { + self.resolve_scheme_id(rule.excluded) == id && rule.children.contains(child_idx) + }) { return true; } @@ -163,10 +162,9 @@ impl CascadingCompressor { // Push rules: Check if any of our ancestors have excluded us. for (ancestor_id, child_idx) in iter { if let Some(ancestor) = self.schemes.iter().find(|s| s.id() == ancestor_id) - && ancestor - .descendant_exclusions() - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + && ancestor.descendant_exclusions().iter().any(|rule| { + self.resolve_scheme_id(rule.excluded) == id && rule.children.contains(child_idx) + }) { return true; } @@ -174,10 +172,9 @@ impl CascadingCompressor { // Pull rules: Check if we have excluded ourselves because of our ancestors. for rule in candidate.ancestor_exclusions() { - if history - .iter() - .any(|(sid, cidx)| *sid == rule.ancestor && rule.children.contains(*cidx)) - { + if history.iter().any(|(sid, cidx)| { + *sid == self.resolve_scheme_id(rule.ancestor) && rule.children.contains(*cidx) + }) { return true; } } diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs new file mode 100644 index 00000000000..1fc6dfbac99 --- /dev/null +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -0,0 +1,272 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use super::*; +use crate::scheme::AncestorExclusion; +use crate::scheme::CompressionEstimate; +use crate::scheme::EstimateVerdict; +use crate::stats::ArrayAndStats; +use crate::stats::GenerateStatsOptions; + +static V1_ID: CachedId = CachedId::new("test.version_1"); +static V2_ID: CachedId = CachedId::new("test.version_2"); +static V3_ID: CachedId = CachedId::new("test.version_3"); +static AUX_ID: CachedId = CachedId::new("test.auxiliary"); + +#[derive(Debug)] +struct TestScheme { + name: &'static str, + version: u8, + predecessor: Option<&'static dyn Scheme>, + push: Option<&'static dyn Scheme>, + pull: Option<&'static dyn Scheme>, +} + +impl TestScheme { + const fn new( + name: &'static str, + version: u8, + predecessor: Option<&'static dyn Scheme>, + ) -> Self { + Self { + name, + version, + predecessor, + push: None, + pull: None, + } + } +} + +impl Scheme for TestScheme { + fn scheme_name(&self) -> &'static str { + self.name + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_int() + } + + fn produced_encodings(&self) -> Vec { + match self.version { + 1 => vec![*V1_ID], + 2 => vec![*V2_ID, *AUX_ID], + 3 => vec![*V3_ID], + _ => vec![], + } + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + self.predecessor + } + + fn num_children(&self) -> usize { + 2 + } + + fn descendant_exclusions(&self) -> Vec { + self.push + .map(|scheme| DescendantExclusion { + excluded: scheme.id(), + children: ChildSelection::One(1), + }) + .into_iter() + .collect() + } + + fn ancestor_exclusions(&self) -> Vec { + self.pull + .map(|scheme| AncestorExclusion { + ancestor: scheme.id(), + children: ChildSelection::One(1), + }) + .into_iter() + .collect() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // Older versions would beat newer versions if they reached estimation together. + CompressionEstimate::Verdict(EstimateVerdict::Ratio(5.0 - f64::from(self.version))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +static V1: TestScheme = TestScheme::new("test.scheme_v1", 1, None); +static V2: TestScheme = TestScheme::new("test.scheme_v2", 2, Some(&V1)); +static V3: TestScheme = TestScheme::new("test.scheme_v3", 3, Some(&V2)); +static OTHER: TestScheme = TestScheme::new("test.other", 0, None); + +#[test] +fn newest_eligible_version_is_selected_before_estimation() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut exec_ctx = session.create_execution_ctx(); + let data = ArrayAndStats::new( + PrimitiveArray::from_iter(0..128i32).into_array(), + GenerateStatsOptions::default(), + ); + for (allowed, expected) in [ + (None, V3.id()), + ( + Some(HashSet::from([*V1_ID, *V2_ID, *AUX_ID, *V3_ID])), + V3.id(), + ), + (Some(HashSet::from([*V1_ID, *V2_ID, *AUX_ID])), V2.id()), + (Some(HashSet::from([*V2_ID, *AUX_ID])), V2.id()), + (Some(HashSet::from([*V1_ID, *V2_ID])), V1.id()), + (Some(HashSet::from([*V1_ID])), V1.id()), + ] { + let mut compressor = CascadingCompressor::new(vec![&V3]); + if let Some(allowed) = allowed { + compressor = compressor.with_allowed_serialized_ids(allowed); + } + assert_eq!(compressor.schemes.len(), 1); + let winner = compressor.choose_best_scheme( + &compressor.schemes, + &data, + compressor.root_context(), + &mut exec_ctx, + )?; + assert_eq!(winner.map(|(scheme, _)| scheme.id()), Some(expected)); + for version in [&V1, &V2, &V3] { + assert!(compressor.has_scheme_family(version.id())); + assert_eq!( + compressor.has_scheme(version.id()), + version.id() == expected + ); + } + assert!(!compressor.has_scheme(OTHER.id())); + } + Ok(()) +} + +#[test] +fn no_eligible_version_removes_the_entire_chain() { + for allowed in [HashSet::new(), HashSet::from([*V2_ID])] { + let compressor = CascadingCompressor::new(vec![&V3]).with_allowed_serialized_ids(allowed); + assert!(compressor.schemes.is_empty()); + for version in [&V1, &V2, &V3] { + assert!(!compressor.has_scheme_family(version.id())); + assert!(!compressor.has_scheme(version.id())); + } + } +} + +#[test] +fn fallback_preserves_registration_order() { + let compressor = CascadingCompressor::new(vec![&V3, &OTHER]) + .with_allowed_serialized_ids(HashSet::from([*V1_ID])); + assert_eq!( + compressor + .schemes + .iter() + .map(|s| s.id()) + .collect::>(), + vec![V1.id(), OTHER.id()] + ); +} + +#[test] +fn successive_restrictions_keep_aliases_and_intersect_wire_ids() { + let compressor = CascadingCompressor::new(vec![&V3]) + .with_allowed_serialized_ids(HashSet::from([*V1_ID, *V2_ID, *AUX_ID])) + .with_allowed_serialized_ids(HashSet::from([*V1_ID])); + assert_eq!(compressor.schemes[0].id(), V1.id()); + assert_eq!(compressor.resolve_scheme_id(V3.id()), V1.id()); + + let compressor = compressor.with_allowed_serialized_ids(HashSet::from([*V2_ID, *AUX_ID])); + assert!(compressor.schemes.is_empty()); + assert!(!compressor.has_scheme_family(V3.id())); +} + +static PUSH_OLD: TestScheme = TestScheme { + push: Some(&V1), + ..TestScheme::new("test.push_old", 0, None) +}; +static PUSH_NEW: TestScheme = TestScheme { + push: Some(&V3), + ..TestScheme::new("test.push_new", 0, None) +}; +static PULL_OLD: TestScheme = TestScheme { + pull: Some(&V1), + ..TestScheme::new("test.pull_old", 0, None) +}; +static PULL_NEW: TestScheme = TestScheme { + pull: Some(&V3), + ..TestScheme::new("test.pull_new", 0, None) +}; + +#[test] +fn exclusions_follow_upgrades_and_fallbacks() { + for allowed in [HashSet::from([*V1_ID]), HashSet::from([*V3_ID])] { + let compressor = + CascadingCompressor::new(vec![&V3, &PUSH_OLD, &PUSH_NEW, &PULL_OLD, &PULL_NEW]) + .with_allowed_serialized_ids(allowed); + let selected = compressor.schemes[0]; + for child in [0, 1] { + for pusher in [&PUSH_OLD, &PUSH_NEW] { + let ctx = compressor + .root_context() + .descend_with_scheme(pusher.id(), child); + assert_eq!(compressor.is_excluded(selected, &ctx), child == 1); + } + let ctx = compressor + .root_context() + .descend_with_scheme(selected.id(), child); + for puller in [&PULL_OLD, &PULL_NEW] { + assert_eq!(compressor.is_excluded(puller, &ctx), child == 1); + } + assert!(compressor.is_excluded(selected, &ctx)); + } + } +} + +#[test] +fn root_exclusions_follow_new_versions() { + static DICT_V2: TestScheme = TestScheme::new("test.dict_v2", 3, Some(&IntDictScheme)); + let compressor = CascadingCompressor::new(vec![&DICT_V2]); + let ctx = compressor + .root_context() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::OFFSETS); + assert!(compressor.is_excluded(&DICT_V2, &ctx)); + let ctx = compressor + .root_context() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::SIZES); + assert!(!compressor.is_excluded(&DICT_V2, &ctx)); +} + +#[test] +#[should_panic(expected = "appears more than once")] +fn predecessor_cycles_are_rejected() { + static CYCLE: TestScheme = TestScheme::new("test.cycle", 1, Some(&CYCLE)); + CascadingCompressor::new(vec![&CYCLE]); +} + +#[test] +#[should_panic(expected = "appears more than once")] +fn registering_multiple_versions_is_rejected() { + CascadingCompressor::new(vec![&V3, &V1]); +} diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..0b9d8e3d4b8 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -41,7 +41,7 @@ pub struct CompressorContext { } impl CompressorContext { - /// Creates a new `CompressorContext`. + /// Creates a new root `CompressorContext`. /// /// This should **only** be created by the compressor. pub(crate) fn new() -> Self { diff --git a/vortex-compressor/src/scheme/exclusion.rs b/vortex-compressor/src/scheme/exclusion.rs index 2dba6b85046..46ca12d7735 100644 --- a/vortex-compressor/src/scheme/exclusion.rs +++ b/vortex-compressor/src/scheme/exclusion.rs @@ -34,7 +34,8 @@ impl ChildSelection { /// `ZigZag` excludes `Dict` from all its children. #[derive(Debug, Clone, Copy)] pub struct DescendantExclusion { - /// The scheme to exclude from descendants. + /// The scheme to exclude from descendants. Any version in its registered predecessor chain + /// refers to the selected version. pub excluded: SchemeId, /// Which children of the declaring scheme this rule applies to. pub children: ChildSelection, @@ -47,7 +48,8 @@ pub struct DescendantExclusion { /// `Sequence` excludes itself when `IntDict` is an ancestor on its codes child. #[derive(Debug, Clone, Copy)] pub struct AncestorExclusion { - /// The ancestor scheme that makes the declaring scheme ineligible. + /// The ancestor scheme that makes the declaring scheme ineligible. Any version in its + /// registered predecessor chain refers to the selected version. pub ancestor: SchemeId, /// Which children of the ancestor this rule applies to. pub children: ChildSelection, diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index de9e67690d4..f00cde3abd9 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,13 +124,31 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The array encodings this scheme itself may introduce into its compressed output. + /// The serialized IDs this scheme may write its output under. Every ID must be permitted before + /// this scheme can be selected. /// - /// Cascaded children are compressed by other schemes, which declare their own encodings, - /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here. - /// Canonical arrays the scheme merely rearranges do not need to be declared. + /// 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. + /// + /// Alternative versions belong in the [`predecessor`](Scheme::predecessor) chain, rather than + /// in this list. Once selected, a scheme must produce output compatible with these IDs without + /// consulting the writer's configuration. fn produced_encodings(&self) -> Vec; + /// The preceding version of this scheme, used when this version's serialized IDs are unavailable. + /// + /// Register only the newest version. The compressor selects the first eligible version in + /// this chain during configuration, before matching, generating statistics, or estimating. + /// A predecessor is a compatibility fallback, not an alternative compression candidate. + /// + /// Versions must have distinct scheme IDs and form an acyclic chain. They must support the + /// same input types and preserve child indices, because exclusions and scheme dependencies + /// referring to any version in the registered chain apply to the selected version. + fn predecessor(&self) -> Option<&'static 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 /// every scheme. diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 13dac7ce7b5..4d3b7bf1a56 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1729,6 +1729,41 @@ async fn test_encoding_registered_after_write_options() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::sparse(PrimitiveArray::from_iter( + (0..4096i32).map(|i| if i % 100 == 0 { i + 1 } else { 0 }), +).into_array())] +#[case::fsst(VarBinViewArray::from_iter( + (0..4096).map(|i| Some(format!("this_is_a_common_prefix_with_some_variation_{i}_and_a_common_suffix_pattern"))), + DType::Utf8(Nullability::NonNullable), +).into_array())] +#[tokio::test] +async fn test_writer_excludes_schemes_with_unavailable_outputs( + #[case] array: ArrayRef, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + // Permit Constant and VarBin, but not the subsequently registered Sparse and FSST. + crate::enable_all_registered_array_encodings(&session); + crate::register_default_encodings(&session); + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, array.clone().to_array_stream()) + .await?; + let read = session + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(read, array, &mut session.create_execution_ctx()); + Ok(()) +} + #[tokio::test] async fn test_writer_empty_chunks() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index ec45653f5c1..f5fe944add9 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,7 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings) = + let (array_ctx, allowed_serialized_ids) = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); @@ -253,7 +253,7 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings), + .allow_serialized_ids(&allowed_serialized_ids), ) .build(), }; @@ -384,6 +384,7 @@ impl VortexWriteOptions { } } +/// Returns the array context and the serialized IDs the compressor may write its output under. fn new_array_context( session: &VortexSession, enforce_editions: bool, @@ -401,11 +402,14 @@ fn new_array_context( .registry() .read(|registry| registry.keys().copied().collect()) }; - let allowed_array_encodings = serialized_ids - .iter() - .filter_map(|serialized_id| arrays.registry().get(serialized_id)) - .map(|plugin| plugin.id()) - .collect(); + // Editions grant permission to use a wire format, but it also must be registered in the session. + let allowed_serialized_ids = arrays.registry().read(|registry| { + serialized_ids + .iter() + .copied() + .filter(|id| registry.contains_key(id)) + .collect() + }); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { // Only permit serialized IDs in the enabled editions. @@ -413,7 +417,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings) + (array_ctx, allowed_serialized_ids) } /// The ids of `kind` the enabled editions permit. @@ -787,29 +791,27 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings) = new_array_context(&session, true); + let (ctx, allowed_serialized_ids) = new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); - assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); + assert_eq!(allowed_serialized_ids, HashSet::from([Primitive.id()])); Ok(()) } #[test] fn disabling_editions_allows_all_registered_array_ids() { let session = array_session(); - let (registered_ids, registered_encodings) = session.arrays().registry().read(|registry| { - ( - registry.keys().copied().sorted().collect::>(), - registry - .values() - .map(|plugin| plugin.id()) - .collect::>(), - ) - }); + let registered_ids = session + .arrays() + .registry() + .read(|registry| registry.keys().copied().sorted().collect::>()); - let (ctx, allowed_array_encodings) = new_array_context(&session, false); + let (ctx, allowed_serialized_ids) = new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); - assert_eq!(allowed_array_encodings, registered_encodings); + assert_eq!( + allowed_serialized_ids, + registered_ids.iter().copied().collect::>() + ); assert!(ctx.intern(&Bool.id()).is_some()); }