From 88dba6480050abd1b381338135b5deb3098edd36 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 4 Sep 2026 11:15:20 -0400 Subject: [PATCH 1/6] Tell the compressor which serialized IDs the writer may emit CascadingCompressor carries the snapshot of serialized IDs the writer may emit, filled by the file writer from the enabled editions through BtrBlocksCompressorBuilder::allow_serialized_ids. A scheme whose encoding has more than one wire format picks its compression mode from it with allows_serialized_id, the newest permitted one; without a restriction every ID is allowed. No scheme consults the set yet. This is the mechanism docs/specs/editions.md describes under compression with replacement encodings (#9779). Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 36 +++++++++- vortex-compressor/src/compressor/mod.rs | 34 ++++++++++ vortex-compressor/src/compressor/tests.rs | 81 +++++++++++++++++++++++ vortex-file/src/writer.rs | 25 +++++-- 4 files changed, 169 insertions(+), 7 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index fe8072d5e66..9701579c5c5 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -91,12 +91,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + allowed_serialized_ids: Option>, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + allowed_serialized_ids: None, } } } @@ -108,6 +110,7 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_serialized_ids: None, } } @@ -214,15 +217,33 @@ impl BtrBlocksCompressorBuilder { self } + /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier + /// call. A scheme whose encoding has several wire formats picks its compression mode from this + /// set: the newest permitted one. + /// + /// The file writer passes the serialized IDs its enabled editions permit. + pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(allowed).copied().collect(), + None => allowed.clone(), + }); + self + } + /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + let compressor = CascadingCompressor::new(self.schemes); + BtrBlocksCompressor(match self.allowed_serialized_ids { + Some(allowed) => compressor.with_allowed_serialized_ids(allowed), + None => compressor, + }) } } #[cfg(test)] mod tests { use vortex_array::VTable; + use vortex_array::arrays::Bool; use vortex_fastlanes::FoR; use super::*; @@ -287,6 +308,19 @@ mod tests { } } + /// Every serialized ID is allowed until the writer narrows the set to its editions. + #[test] + fn allowed_serialized_ids_reach_the_compressor() { + let default = BtrBlocksCompressorBuilder::default().build(); + assert!(default.0.allows_serialized_id(Bool.id())); + + let narrowed = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .build(); + assert!(narrowed.0.allows_serialized_id(FoR.id())); + assert!(!narrowed.0.allows_serialized_id(Bool.id())); + } + #[test] fn cuda_compatible_uses_fsst_for_strings() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..774513c66a7 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,9 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,6 +49,10 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// The serialized IDs the output may use, or `None` for no restriction. See + /// [`allows_serialized_id`](Self::allows_serialized_id). + allowed_serialized_ids: Option>, } impl CascadingCompressor { @@ -63,9 +70,36 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + allowed_serialized_ids: None, } } + /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier + /// call. + /// + /// The file writer passes the serialized IDs its enabled editions permit. A scheme whose + /// encoding has several wire formats picks its compression mode from this set, the newest + /// permitted one, before estimating or compressing. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(&allowed).copied().collect(), + None => allowed, + }); + self + } + + /// Returns whether the writer may emit the serialized ID `id`. + /// + /// Schemes whose encoding has several wire formats consult this to pick their compression + /// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The + /// serializer still emits the oldest wire form the resulting array fits, and the + /// serialization context validates that ID. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns whether the compressor was configured with `scheme`. pub fn has_scheme(&self, scheme: SchemeId) -> bool { self.schemes diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index ec14383ce36..69afcd3e9ea 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,11 +9,14 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Bool; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -26,6 +29,7 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -96,6 +100,48 @@ impl Scheme for DirectRatioScheme { } } +/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`. +static SEEN_FORMAT: Mutex> = Mutex::new(None); + +/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether +/// the newer one is allowed and records the answer. +#[derive(Debug)] +struct FormatRecordingScheme; + +impl Scheme for FormatRecordingScheme { + fn scheme_name(&self) -> &'static str { + "test.format_recording" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id())); + Ok(data.array().clone()) + } +} + #[derive(Debug)] struct ImmediateAlwaysUseScheme; @@ -841,3 +887,38 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } + +#[test] +fn allowed_serialized_ids_default_to_everything_and_intersect() { + let compressor = compressor(); + assert!(compressor.allows_serialized_id(Constant.id())); + assert!(compressor.allows_serialized_id(Bool.id())); + + let restricted = + compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); + assert!(restricted.allows_serialized_id(Constant.id())); + assert!(!restricted.allows_serialized_id(Bool.id())); + + let narrowed = + restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); + assert!(narrowed.allows_serialized_id(Primitive.id())); + assert!(!narrowed.allows_serialized_id(Constant.id())); + assert!(!narrowed.allows_serialized_id(Bool.id())); +} + +/// A scheme sees the restriction through the compressor it is handed: everything is allowed until +/// the writer narrows the set to its editions. +#[test] +fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0..4096i32).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]); + unrestricted.compress(&array, &mut exec_ctx)?; + assert_eq!(*SEEN_FORMAT.lock(), Some(true)); + + let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()])); + restricted.compress(&array, &mut exec_ctx)?; + assert_eq!(*SEEN_FORMAT.lock(), Some(false)); + Ok(()) +} diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index ec45653f5c1..220599733fb 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,7 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings) = + let (array_ctx, allowed_array_encodings, allowed_serialized_ids) = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); @@ -253,7 +253,8 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings), + .retain_allowed_encodings(&allowed_array_encodings) + .allow_serialized_ids(&allowed_serialized_ids), ) .build(), }; @@ -384,10 +385,12 @@ impl VortexWriteOptions { } } +/// Returns the array context, the in-memory encodings the compressor may produce, and the +/// serialized IDs its output may use. fn new_array_context( session: &VortexSession, enforce_editions: bool, -) -> (ArrayContext, HashSet) { +) -> (ArrayContext, HashSet, HashSet) { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done @@ -406,6 +409,9 @@ fn new_array_context( .filter_map(|serialized_id| arrays.registry().get(serialized_id)) .map(|plugin| plugin.id()) .collect(); + // The compressor sees the same set, so an encoding with several wire formats produces the + // newest one the editions permit. + let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { // Only permit serialized IDs in the enabled editions. @@ -413,7 +419,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings) + (array_ctx, allowed_array_encodings, allowed_serialized_ids) } /// The ids of `kind` the enabled editions permit. @@ -787,10 +793,12 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings) = new_array_context(&session, true); + let (ctx, allowed_array_encodings, allowed_serialized_ids) = + new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); + assert_eq!(allowed_serialized_ids, HashSet::from([Primitive.id()])); Ok(()) } @@ -807,9 +815,14 @@ mod tests { ) }); - let (ctx, allowed_array_encodings) = new_array_context(&session, false); + let (ctx, allowed_array_encodings, allowed_serialized_ids) = + new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); assert_eq!(allowed_array_encodings, registered_encodings); + assert_eq!( + allowed_serialized_ids, + registered_ids.iter().copied().collect::>() + ); assert!(ctx.intern(&Bool.id()).is_some()); } From 98adaee1c554830c7c5c0825a78ed7e269b552d1 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 10:47:58 -0400 Subject: [PATCH 2/6] Enable a scheme when any of its serialized IDs is permitted Scheme::produced_encodings now names the serialized IDs a scheme may write its output under, oldest first. BtrBlocksCompressorBuilder::allow_serialized_ids replaces retain_allowed_encodings: it keeps a scheme when at least one of those IDs is permitted and hands the set to the compressor, so the writer makes one call from the serialized IDs its editions permit instead of mapping them back to in-memory encodings, which could not tell two wire formats of one encoding apart. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 97 +++++++++++++++++++++++------ vortex-compressor/src/scheme/mod.rs | 12 ++-- vortex-file/src/writer.rs | 40 ++++-------- 3 files changed, 98 insertions(+), 51 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 9701579c5c5..2c9a54456b5 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -207,26 +207,22 @@ impl BtrBlocksCompressorBuilder { 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))); - self - } - - /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier - /// call. A scheme whose encoding has several wire formats picks its compression mode from this - /// set: the newest permitted one. + /// A scheme stays when at least one of its [produced IDs](Scheme::produced_encodings) is + /// permitted, and the compressor is handed the set so a scheme whose encoding has several + /// wire formats picks its compression mode from it: the newest permitted one. /// /// The file writer passes the serialized IDs its enabled editions permit. pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { - self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + let allowed: HashSet = match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(allowed).copied().collect(), None => allowed.clone(), - }); + }; + self.schemes + .retain(|s| s.produced_encodings().iter().any(|id| allowed.contains(id))); + self.allowed_serialized_ids = Some(allowed); self } @@ -242,11 +238,20 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; use vortex_array::VTable; use vortex_array::arrays::Bool; + use vortex_array::arrays::Primitive; + use vortex_compressor::scheme::CompressionEstimate; + use vortex_compressor::scheme::EstimateVerdict; + use vortex_error::VortexResult; use vortex_fastlanes::FoR; use super::*; + use crate::ArrayAndStats; + use crate::CompressorContext; #[test] fn empty_starts_with_no_schemes() { @@ -261,26 +266,80 @@ mod tests { } #[test] - fn retain_allowed_encodings_filters_schemes() { + fn allow_serialized_ids_filters_schemes() { let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&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()); + let none = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&HashSet::new()); assert!(none.schemes.is_empty()); } #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { + fn allowing_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); + let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + /// Stands in for a scheme whose encoding has two wire formats. + #[derive(Debug)] + struct TwoFormatScheme; + + impl Scheme for TwoFormatScheme { + fn scheme_name(&self) -> &'static str { + "test.two_formats" + } + + fn matches(&self, _canonical: &Canonical) -> bool { + false + } + + fn produced_encodings(&self) -> Vec { + vec![FoR.id(), Bool.id()] + } + + 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 { + unreachable!("test helper never matches") + } + } + + /// A scheme with several wire formats stays while any of them is permitted; which one it + /// produces is decided when compressing. + #[test] + fn any_permitted_format_keeps_the_scheme() { + static TWO_FORMATS: TwoFormatScheme = TwoFormatScheme; + + let newer_only = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&TWO_FORMATS) + .allow_serialized_ids(&HashSet::from([Bool.id()])); + assert_eq!(newer_only.schemes.len(), 1); + + let neither = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&TWO_FORMATS) + .allow_serialized_ids(&HashSet::from([Primitive.id()])); + assert!(neither.schemes.is_empty()); + } + #[test] fn cuda_compatible_excludes_alprd() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index de9e67690d4..fa42231d422 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,11 +124,15 @@ 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. /// - /// 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. + /// + /// An encoding with several wire formats lists every one of them, oldest first. The writer + /// keeps the scheme while any of them is permitted, and the scheme picks the newest + /// permitted one as its compression mode. fn produced_encodings(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 220599733fb..79d485f9b1f 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, allowed_serialized_ids) = + 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,6 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings) .allow_serialized_ids(&allowed_serialized_ids), ) .build(), @@ -385,12 +384,11 @@ impl VortexWriteOptions { } } -/// Returns the array context, the in-memory encodings the compressor may produce, and the -/// serialized IDs its output may use. +/// Returns the array context and the serialized IDs the compressor may write its output under. fn new_array_context( session: &VortexSession, enforce_editions: bool, -) -> (ArrayContext, HashSet, HashSet) { +) -> (ArrayContext, HashSet) { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done @@ -404,13 +402,8 @@ 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(); - // The compressor sees the same set, so an encoding with several wire formats produces the - // newest one the editions permit. + // The compressor sees the same set: it keeps the schemes that can write one of these IDs, and + // an encoding with several wire formats produces the newest one permitted. let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { @@ -419,7 +412,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings, allowed_serialized_ids) + (array_ctx, allowed_serialized_ids) } /// The ids of `kind` the enabled editions permit. @@ -793,11 +786,9 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings, allowed_serialized_ids) = - 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(()) } @@ -805,20 +796,13 @@ mod tests { #[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, allowed_serialized_ids) = - 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::>() From dafbfa47813ce9fd39098a40efe0d713be1631f4 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 10:49:53 -0400 Subject: [PATCH 3/6] Read the permitted serialized IDs from the CompressorContext The compressor seeds each root CompressorContext with its permitted serialized IDs and every descent inherits them, so a scheme asks compress_ctx.allows_serialized_id both while estimating and while compressing and picks the same mode in both. The per-compressor accessor goes; allowed_serialized_ids remains for inspection. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 8 +- vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 38 +++---- vortex-compressor/src/compressor/tests.rs | 110 ++++++++++++++------ vortex-compressor/src/scheme/ctx.rs | 27 ++++- 5 files changed, 127 insertions(+), 58 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 2c9a54456b5..ad25969a028 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -371,13 +371,15 @@ mod tests { #[test] fn allowed_serialized_ids_reach_the_compressor() { let default = BtrBlocksCompressorBuilder::default().build(); - assert!(default.0.allows_serialized_id(Bool.id())); + assert!(default.0.allowed_serialized_ids().is_none()); let narrowed = BtrBlocksCompressorBuilder::default() .allow_serialized_ids(&HashSet::from([FoR.id()])) .build(); - assert!(narrowed.0.allows_serialized_id(FoR.id())); - assert!(!narrowed.0.allows_serialized_id(Bool.id())); + assert_eq!( + narrowed.0.allowed_serialized_ids(), + Some(&HashSet::from([FoR.id()])) + ); } #[test] diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..dd98f4ea3c6 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()); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 774513c66a7..159842f2595 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,11 +9,14 @@ mod sample; mod select; mod structural; +use std::sync::Arc; + use vortex_array::ArrayId; use vortex_utils::aliases::hash_set::HashSet; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; +use crate::scheme::CompressorContext; use crate::scheme::DescendantExclusion; use crate::scheme::Scheme; use crate::scheme::SchemeExt; @@ -50,9 +53,9 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The serialized IDs the output may use, or `None` for no restriction. See - /// [`allows_serialized_id`](Self::allows_serialized_id). - allowed_serialized_ids: Option>, + /// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root + /// [`CompressorContext`], where schemes read it. + allowed_serialized_ids: Option>>, } impl CascadingCompressor { @@ -77,27 +80,26 @@ impl CascadingCompressor { /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier /// call. /// - /// The file writer passes the serialized IDs its enabled editions permit. A scheme whose - /// encoding has several wire formats picks its compression mode from this set, the newest - /// permitted one, before estimating or compressing. + /// The file writer passes the serialized IDs its enabled editions permit. Schemes read the + /// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has + /// several wire formats picks the newest permitted one as its mode, while estimating and + /// while compressing alike. pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { - self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(&allowed).copied().collect(), None => allowed, - }); + })); self } - /// Returns whether the writer may emit the serialized ID `id`. - /// - /// Schemes whose encoding has several wire formats consult this to pick their compression - /// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The - /// serializer still emits the oldest wire form the resulting array fits, and the - /// serialization context validates that ID. - pub fn allows_serialized_id(&self, id: ArrayId) -> bool { - self.allowed_serialized_ids - .as_ref() - .is_none_or(|allowed| allowed.contains(&id)) + /// The serialized IDs the writer may emit, or `None` when unrestricted. + pub fn allowed_serialized_ids(&self) -> Option<&HashSet> { + self.allowed_serialized_ids.as_deref() + } + + /// The context a compress call starts from. + pub(crate) fn root_context(&self) -> CompressorContext { + CompressorContext::new(self.allowed_serialized_ids.clone()) } /// Returns whether the compressor was configured with `scheme`. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 69afcd3e9ea..42a98245fef 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -132,12 +132,12 @@ impl Scheme for FormatRecordingScheme { fn compress( &self, - compressor: &CascadingCompressor, + _compressor: &CascadingCompressor, data: &ArrayAndStats, - _compress_ctx: CompressorContext, + compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id())); + *SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id())); Ok(data.array().clone()) } } @@ -420,8 +420,12 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -438,8 +442,12 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -456,8 +464,12 @@ fn callback_skip_is_ignored() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -474,8 +486,12 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -492,8 +508,12 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -510,8 +530,12 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -528,8 +552,12 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(winner.is_none()); Ok(()) @@ -630,8 +658,12 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -651,7 +683,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -672,7 +704,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner // `None`) because the zero-byte sample is never stored as the best. @@ -691,7 +723,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert_eq!(observed, Some(None)); @@ -711,7 +743,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -732,8 +764,12 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -780,7 +816,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // A context with default stats_options (count_distinct_values = false) and // marked as a sample so the function skips the sampling step and compresses // the array directly. - let ctx = CompressorContext::new().with_sampling(); + let ctx = CompressorContext::new(None).with_sampling(); // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" @@ -891,19 +927,27 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { #[test] fn allowed_serialized_ids_default_to_everything_and_intersect() { let compressor = compressor(); - assert!(compressor.allows_serialized_id(Constant.id())); - assert!(compressor.allows_serialized_id(Bool.id())); + let root = compressor.root_context(); + assert!(root.allows_serialized_id(Constant.id())); + assert!(root.allows_serialized_id(Bool.id())); let restricted = compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); - assert!(restricted.allows_serialized_id(Constant.id())); - assert!(!restricted.allows_serialized_id(Bool.id())); + let root = restricted.root_context(); + assert!(root.allows_serialized_id(Constant.id())); + assert!(!root.allows_serialized_id(Bool.id())); let narrowed = restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); - assert!(narrowed.allows_serialized_id(Primitive.id())); - assert!(!narrowed.allows_serialized_id(Constant.id())); - assert!(!narrowed.allows_serialized_id(Bool.id())); + let root = narrowed.root_context(); + assert!(root.allows_serialized_id(Primitive.id())); + assert!(!root.allows_serialized_id(Constant.id())); + assert!(!root.allows_serialized_id(Bool.id())); + + // Descending keeps the set. + let child = root.descend_with_scheme(IntDictScheme.id(), 0); + assert!(child.allows_serialized_id(Primitive.id())); + assert!(!child.allows_serialized_id(Constant.id())); } /// A scheme sees the restriction through the compressor it is handed: everything is allowed until diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..83685031c33 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,8 +4,11 @@ //! Compression context for recursive compression. use std::fmt; +use std::sync::Arc; +use vortex_array::ArrayId; use vortex_error::VortexExpect; +use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -38,18 +41,24 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, + + /// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every + /// context of one compress call, so cloning at each descent is a pointer bump. + allowed_serialized_ids: Option>>, } impl CompressorContext { - /// Creates a new `CompressorContext`. + /// Creates a new root `CompressorContext` for a compressor that may emit the given serialized + /// IDs, or any ID when `None`. /// /// This should **only** be created by the compressor. - pub(crate) fn new() -> Self { + pub(crate) fn new(allowed_serialized_ids: Option>>) -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), + allowed_serialized_ids, } } } @@ -57,7 +66,7 @@ impl CompressorContext { #[cfg(test)] impl Default for CompressorContext { fn default() -> Self { - Self::new() + Self::new(None) } } @@ -67,6 +76,18 @@ impl CompressorContext { self.is_sample } + /// Returns whether the writer may emit the serialized ID `id`. + /// + /// A scheme whose encoding has several wire formats picks its compression mode from this, + /// the newest permitted one, and the same answer is available while estimating and while + /// compressing. Without a restriction every ID is allowed. The serializer still emits the + /// oldest wire form the resulting array fits, and the serialization context validates it. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options From 1752d46cedd6171c37052f6548ada7944ea96f93 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 14:25:55 -0400 Subject: [PATCH 4/6] add predecessor schemes and make produced_encodings refer to serialized ids Signed-off-by: Matt Katz --- .../src/{builder.rs => builder/mod.rs} | 217 +++----------- vortex-btrblocks/src/builder/tests.rs | 216 ++++++++++++++ .../src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- .../src/schemes/string/zstd_buffers.rs | 2 +- vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 98 +++++-- vortex-compressor/src/compressor/select.rs | 21 +- vortex-compressor/src/compressor/tests.rs | 171 ++--------- .../src/compressor/version_tests.rs | 266 ++++++++++++++++++ vortex-compressor/src/scheme/ctx.rs | 27 +- vortex-compressor/src/scheme/exclusion.rs | 6 +- vortex-compressor/src/scheme/mod.rs | 22 +- vortex-file/src/tests.rs | 35 +++ vortex-file/src/writer.rs | 2 - 16 files changed, 686 insertions(+), 405 deletions(-) rename vortex-btrblocks/src/{builder.rs => builder/mod.rs} (58%) create mode 100644 vortex-btrblocks/src/builder/tests.rs create mode 100644 vortex-compressor/src/compressor/version_tests.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder/mod.rs similarity index 58% rename from vortex-btrblocks/src/builder.rs rename to vortex-btrblocks/src/builder/mod.rs index ad25969a028..ea0ef849dd2 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. @@ -117,7 +117,7 @@ impl BtrBlocksCompressorBuilder { /// 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 /// @@ -201,32 +201,53 @@ 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 } /// Restricts compression to the serialized IDs in `allowed`, intersecting with any earlier /// call. /// - /// A scheme stays when at least one of its [produced IDs](Scheme::produced_encodings) is - /// permitted, and the compressor is handed the set so a scheme whose encoding has several - /// wire formats picks its compression mode from it: the newest permitted one. - /// - /// The file writer passes the serialized IDs its enabled editions permit. + /// At build time, each scheme is replaced by the newest version in its predecessor chain + /// whose [`required_serialized_ids`](Scheme::required_serialized_ids) 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.schemes - .retain(|s| s.produced_encodings().iter().any(|id| allowed.contains(id))); 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 { let compressor = CascadingCompressor::new(self.schemes); BtrBlocksCompressor(match self.allowed_serialized_ids { @@ -237,178 +258,4 @@ impl BtrBlocksCompressorBuilder { } #[cfg(test)] -mod tests { - use vortex_array::ArrayRef; - use vortex_array::Canonical; - use vortex_array::ExecutionCtx; - use vortex_array::VTable; - use vortex_array::arrays::Bool; - use vortex_array::arrays::Primitive; - use vortex_compressor::scheme::CompressionEstimate; - use vortex_compressor::scheme::EstimateVerdict; - use vortex_error::VortexResult; - use vortex_fastlanes::FoR; - - use super::*; - use crate::ArrayAndStats; - use crate::CompressorContext; - - #[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 allow_serialized_ids_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn allowing_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - /// Stands in for a scheme whose encoding has two wire formats. - #[derive(Debug)] - struct TwoFormatScheme; - - impl Scheme for TwoFormatScheme { - fn scheme_name(&self) -> &'static str { - "test.two_formats" - } - - fn matches(&self, _canonical: &Canonical) -> bool { - false - } - - fn produced_encodings(&self) -> Vec { - vec![FoR.id(), Bool.id()] - } - - 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 { - unreachable!("test helper never matches") - } - } - - /// A scheme with several wire formats stays while any of them is permitted; which one it - /// produces is decided when compressing. - #[test] - fn any_permitted_format_keeps_the_scheme() { - static TWO_FORMATS: TwoFormatScheme = TwoFormatScheme; - - let newer_only = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&TWO_FORMATS) - .allow_serialized_ids(&HashSet::from([Bool.id()])); - assert_eq!(newer_only.schemes.len(), 1); - - let neither = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&TWO_FORMATS) - .allow_serialized_ids(&HashSet::from([Primitive.id()])); - assert!(neither.schemes.is_empty()); - } - - #[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" - ); - } - } - - /// Every serialized ID is allowed until the writer narrows the set to its editions. - #[test] - fn allowed_serialized_ids_reach_the_compressor() { - let default = BtrBlocksCompressorBuilder::default().build(); - assert!(default.0.allowed_serialized_ids().is_none()); - - let narrowed = BtrBlocksCompressorBuilder::default() - .allow_serialized_ids(&HashSet::from([FoR.id()])) - .build(); - assert_eq!( - narrowed.0.allowed_serialized_ids(), - Some(&HashSet::from([FoR.id()])) - ); - } - - #[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..1b736ba895b --- /dev/null +++ b/vortex-btrblocks/src/builder/tests.rs @@ -0,0 +1,216 @@ +// 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(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(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); +} diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 3f06d65b061..5204e4e2478 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 46b2f1e302e..86b69bba47e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index a1bc8643775..dddaa349de1 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index cf691c70fcb..98d5feee2d9 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index dd98f4ea3c6..ecfd3c2c542 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -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 159842f2595..03733aca5e1 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,9 +9,8 @@ mod sample; mod select; mod structural; -use std::sync::Arc; - use vortex_array::ArrayId; +use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; use crate::builtins::IntDictScheme; @@ -53,16 +52,38 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root - /// [`CompressorContext`], where schemes read it. - allowed_serialized_ids: Option>>, + /// 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`](Self::has_scheme) 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 { @@ -73,40 +94,68 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + scheme_aliases, allowed_serialized_ids: None, } } - /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier - /// call. + /// Selects the newest eligible version of each scheme, intersecting with any earlier call. /// - /// The file writer passes the serialized IDs its enabled editions permit. Schemes read the - /// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has - /// several wire formats picks the newest permitted one as its mode, while estimating and - /// while compressing alike. + /// A version is eligible only when all of its [`Scheme::required_serialized_ids`] 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 { - self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() { + 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 serialized IDs the writer may emit, or `None` when unrestricted. - pub fn allowed_serialized_ids(&self) -> Option<&HashSet> { - self.allowed_serialized_ids.as_deref() - } - /// The context a compress call starts from. pub(crate) fn root_context(&self) -> CompressorContext { - CompressorContext::new(self.allowed_serialized_ids.clone()) + CompressorContext::new() } - /// Returns whether the compressor was configured with `scheme`. + /// Returns whether a version of `scheme` is enabled. + /// + /// Any ID in a registered predecessor chain refers to the selected version, including when + /// the selected version is older or newer than the specified ID. pub fn has_scheme(&self, scheme: SchemeId) -> bool { - self.schemes - .iter() - .any(|candidate| candidate.id() == scheme) + self.scheme_aliases.contains_key(&scheme) + } + + fn resolve_scheme_id(&self, scheme: SchemeId) -> SchemeId { + self.scheme_aliases.get(&scheme).copied().unwrap_or(scheme) } } @@ -114,3 +163,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/tests.rs b/vortex-compressor/src/compressor/tests.rs index 42a98245fef..ec14383ce36 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,14 +9,11 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::VTable; use vortex_array::VortexSessionExecute; -use vortex_array::arrays::Bool; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; -use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -29,7 +26,6 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -100,48 +96,6 @@ impl Scheme for DirectRatioScheme { } } -/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`. -static SEEN_FORMAT: Mutex> = Mutex::new(None); - -/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether -/// the newer one is allowed and records the answer. -#[derive(Debug)] -struct FormatRecordingScheme; - -impl Scheme for FormatRecordingScheme { - fn scheme_name(&self) -> &'static str { - "test.format_recording" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn produced_encodings(&self) -> Vec { - Vec::new() - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - *SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id())); - Ok(data.array().clone()) - } -} - #[derive(Debug)] struct ImmediateAlwaysUseScheme; @@ -420,12 +374,8 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -442,12 +392,8 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -464,12 +410,8 @@ fn callback_skip_is_ignored() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -486,12 +428,8 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -508,12 +446,8 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -530,12 +464,8 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -552,12 +482,8 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(winner.is_none()); Ok(()) @@ -658,12 +584,8 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -683,7 +605,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -704,7 +626,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner // `None`) because the zero-byte sample is never stored as the best. @@ -723,7 +645,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert_eq!(observed, Some(None)); @@ -743,7 +665,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -764,12 +686,8 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -816,7 +734,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // A context with default stats_options (count_distinct_values = false) and // marked as a sample so the function skips the sampling step and compresses // the array directly. - let ctx = CompressorContext::new(None).with_sampling(); + let ctx = CompressorContext::new().with_sampling(); // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" @@ -923,46 +841,3 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } - -#[test] -fn allowed_serialized_ids_default_to_everything_and_intersect() { - let compressor = compressor(); - let root = compressor.root_context(); - assert!(root.allows_serialized_id(Constant.id())); - assert!(root.allows_serialized_id(Bool.id())); - - let restricted = - compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); - let root = restricted.root_context(); - assert!(root.allows_serialized_id(Constant.id())); - assert!(!root.allows_serialized_id(Bool.id())); - - let narrowed = - restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); - let root = narrowed.root_context(); - assert!(root.allows_serialized_id(Primitive.id())); - assert!(!root.allows_serialized_id(Constant.id())); - assert!(!root.allows_serialized_id(Bool.id())); - - // Descending keeps the set. - let child = root.descend_with_scheme(IntDictScheme.id(), 0); - assert!(child.allows_serialized_id(Primitive.id())); - assert!(!child.allows_serialized_id(Constant.id())); -} - -/// A scheme sees the restriction through the compressor it is handed: everything is allowed until -/// the writer narrows the set to its editions. -#[test] -fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> { - let array = PrimitiveArray::from_iter(0..4096i32).into_array(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]); - unrestricted.compress(&array, &mut exec_ctx)?; - assert_eq!(*SEEN_FORMAT.lock(), Some(true)); - - let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()])); - restricted.compress(&array, &mut exec_ctx)?; - assert_eq!(*SEEN_FORMAT.lock(), Some(false)); - Ok(()) -} diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs new file mode 100644 index 00000000000..197583331cf --- /dev/null +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -0,0 +1,266 @@ +// 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(version.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(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(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 83685031c33..0b9d8e3d4b8 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,11 +4,8 @@ //! Compression context for recursive compression. use std::fmt; -use std::sync::Arc; -use vortex_array::ArrayId; use vortex_error::VortexExpect; -use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -41,24 +38,18 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, - - /// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every - /// context of one compress call, so cloning at each descent is a pointer bump. - allowed_serialized_ids: Option>>, } impl CompressorContext { - /// Creates a new root `CompressorContext` for a compressor that may emit the given serialized - /// IDs, or any ID when `None`. + /// Creates a new root `CompressorContext`. /// /// This should **only** be created by the compressor. - pub(crate) fn new(allowed_serialized_ids: Option>>) -> Self { + pub(crate) fn new() -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), - allowed_serialized_ids, } } } @@ -66,7 +57,7 @@ impl CompressorContext { #[cfg(test)] impl Default for CompressorContext { fn default() -> Self { - Self::new(None) + Self::new() } } @@ -76,18 +67,6 @@ impl CompressorContext { self.is_sample } - /// Returns whether the writer may emit the serialized ID `id`. - /// - /// A scheme whose encoding has several wire formats picks its compression mode from this, - /// the newest permitted one, and the same answer is available while estimating and while - /// compressing. Without a restriction every ID is allowed. The serializer still emits the - /// oldest wire form the resulting array fits, and the serialization context validates it. - pub fn allows_serialized_id(&self, id: ArrayId) -> bool { - self.allowed_serialized_ids - .as_ref() - .is_none_or(|allowed| allowed.contains(&id)) - } - /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options diff --git a/vortex-compressor/src/scheme/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 fa42231d422..f00cde3abd9 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,17 +124,31 @@ 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 may write its output under. + /// 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 IDs, so only /// arrays constructed directly by [`compress`](Scheme::compress) belong here. Canonical /// arrays the scheme merely rearranges do not need to be declared. /// - /// An encoding with several wire formats lists every one of them, oldest first. The writer - /// keeps the scheme while any of them is permitted, and the scheme picks the newest - /// permitted one as its compression mode. + /// 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 79d485f9b1f..725405d3302 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -402,8 +402,6 @@ fn new_array_context( .registry() .read(|registry| registry.keys().copied().collect()) }; - // The compressor sees the same set: it keeps the schemes that can write one of these IDs, and - // an encoding with several wire formats produces the newest one permitted. let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { From 7accd56be435fd2e90f91c2b55debdaf626fabbf Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 14:40:19 -0400 Subject: [PATCH 5/6] Fix optional scheme trait methods and documentation links Restore produced_encodings in feature-gated schemes and fix stale trait links. Document scheme ID resolution for Clippy. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder/mod.rs | 2 +- vortex-btrblocks/src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- vortex-btrblocks/src/schemes/string/zstd_buffers.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 3 ++- 6 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vortex-btrblocks/src/builder/mod.rs b/vortex-btrblocks/src/builder/mod.rs index ea0ef849dd2..49cd98213af 100644 --- a/vortex-btrblocks/src/builder/mod.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -231,7 +231,7 @@ impl BtrBlocksCompressorBuilder { /// call. /// /// At build time, each scheme is replaced by the newest version in its predecessor chain - /// whose [`required_serialized_ids`](Scheme::required_serialized_ids) are all permitted. + /// 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 { diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 5204e4e2478..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 86b69bba47e..46b2f1e302e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index dddaa349de1..a1bc8643775 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index 98d5feee2d9..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 03733aca5e1..464768e8b36 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -101,7 +101,7 @@ impl CascadingCompressor { /// Selects the newest eligible version of each scheme, intersecting with any earlier call. /// - /// A version is eligible only when all of its [`Scheme::required_serialized_ids`] are allowed. + /// 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 { @@ -154,6 +154,7 @@ impl CascadingCompressor { self.scheme_aliases.contains_key(&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) } From ece1fe2478d54e1d5a2c42856fbe01f3ae819db0 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 17:52:01 -0400 Subject: [PATCH 6/6] fix Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder/tests.rs | 7 +- vortex-btrblocks/tests/scheme_versions.rs | 103 ++++++++++++++++++ vortex-compressor/src/compressor/mod.rs | 17 ++- .../src/compressor/version_tests.rs | 10 +- vortex-file/src/writer.rs | 9 +- 5 files changed, 135 insertions(+), 11 deletions(-) create mode 100644 vortex-btrblocks/tests/scheme_versions.rs diff --git a/vortex-btrblocks/src/builder/tests.rs b/vortex-btrblocks/src/builder/tests.rs index 1b736ba895b..4310b9416c8 100644 --- a/vortex-btrblocks/src/builder/tests.rs +++ b/vortex-btrblocks/src/builder/tests.rs @@ -200,7 +200,8 @@ fn restrictions_select_predecessors_of_schemes_added_later() { .with_new_scheme(&NewFoRScheme) .build(); assert!(compressor.has_scheme(integer::FoRScheme.id())); - assert!(compressor.has_scheme(NewFoRScheme.id())); + assert!(compressor.has_scheme_family(NewFoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); } #[rstest] @@ -211,6 +212,6 @@ fn excluding_any_version_removes_the_chain(#[case] excluded: SchemeId) { .with_new_scheme(&NewFoRScheme) .exclude_schemes([excluded]) .build(); - assert!(!compressor.has_scheme(integer::FoRScheme.id())); - assert!(!compressor.has_scheme(NewFoRScheme.id())); + 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/mod.rs b/vortex-compressor/src/compressor/mod.rs index 464768e8b36..3f88debbc24 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -63,7 +63,7 @@ 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`](Self::has_scheme) checks. + /// 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 @@ -146,14 +146,21 @@ impl CascadingCompressor { CompressorContext::new() } - /// Returns whether a version of `scheme` is enabled. + /// Returns whether any version in the scheme's family is enabled. /// - /// Any ID in a registered predecessor chain refers to the selected version, including when - /// the selected version is older or newer than the specified ID. - pub fn has_scheme(&self, scheme: SchemeId) -> bool { + /// 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.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) diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs index 197583331cf..1fc6dfbac99 100644 --- a/vortex-compressor/src/compressor/version_tests.rs +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -152,8 +152,13 @@ fn newest_eligible_version_is_selected_before_estimation() -> VortexResult<()> { )?; assert_eq!(winner.map(|(scheme, _)| scheme.id()), Some(expected)); for version in [&V1, &V2, &V3] { - assert!(compressor.has_scheme(version.id())); + assert!(compressor.has_scheme_family(version.id())); + assert_eq!( + compressor.has_scheme(version.id()), + version.id() == expected + ); } + assert!(!compressor.has_scheme(OTHER.id())); } Ok(()) } @@ -164,6 +169,7 @@ fn no_eligible_version_removes_the_entire_chain() { 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())); } } @@ -193,7 +199,7 @@ fn successive_restrictions_keep_aliases_and_intersect_wire_ids() { let compressor = compressor.with_allowed_serialized_ids(HashSet::from([*V2_ID, *AUX_ID])); assert!(compressor.schemes.is_empty()); - assert!(!compressor.has_scheme(V3.id())); + assert!(!compressor.has_scheme_family(V3.id())); } static PUSH_OLD: TestScheme = TestScheme { diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 725405d3302..f5fe944add9 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -402,7 +402,14 @@ fn new_array_context( .registry() .read(|registry| registry.keys().copied().collect()) }; - let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().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.