diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 341569f408c..2f2c2ec11a9 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -113,7 +113,8 @@ 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 a newer version of a scheme alongside the version it + /// [replaces](Scheme::replaces). /// /// # Panics /// @@ -203,7 +204,8 @@ impl BtrBlocksCompressorBuilder { /// Retains only schemes whose produced serialized IDs all belong to `allowed`. /// /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. + /// permit. When a newer scheme version is dropped here, the version it + /// [replaces](Scheme::replaces) is no longer replaced and compresses as before. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { self.schemes .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); diff --git a/vortex-btrblocks/tests/scheme_replacement.rs b/vortex-btrblocks/tests/scheme_replacement.rs new file mode 100644 index 00000000000..c5c8da3b942 --- /dev/null +++ b/vortex-btrblocks/tests/scheme_replacement.rs @@ -0,0 +1,125 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[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::BtrBlocksCompressorBuilder; + 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_compressor::scheme::SchemeId; + 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 replaces(&self) -> Vec { + vec![DELTA_V1.id()] + } + + 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_the_active_delta_version( + #[case] allowed_delta: ArrayId, + #[case] expect_delta: bool, + ) -> VortexResult<()> { + let session = array_session(); + vortex_fastlanes::initialize(&session); + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&IntRLEScheme) + .with_new_scheme(&DELTA_V1) + .with_new_scheme(&DeltaV2) + .retain_allowed_encodings(&[RLE.id(), allowed_delta].into_iter().collect()) + .build(); + assert_eq!(compressor.has_scheme(DELTA_V1.id()), expect_delta); + assert_eq!(compressor.has_scheme(DeltaV2.id()), !expect_delta); + 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(()) + } + + /// Excluding a replacement removes only that scheme; the version it replaces stays active. + #[test] + fn excluding_the_replacement_keeps_the_replaced_scheme() { + let allowed = [RLE.id(), Delta.id(), *DELTA_V2_ID].into_iter().collect(); + let with_both = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&IntRLEScheme) + .with_new_scheme(&DELTA_V1) + .with_new_scheme(&DeltaV2) + .retain_allowed_encodings(&allowed); + assert!(with_both.clone().build().has_scheme(DeltaV2.id())); + assert!(!with_both.clone().build().has_scheme(DELTA_V1.id())); + + let excluded = with_both.exclude_schemes([DeltaV2.id()]).build(); + assert!(!excluded.has_scheme(DeltaV2.id())); + assert!(excluded.has_scheme(DELTA_V1.id())); + } +} diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..bc260e4670f 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,8 @@ mod sample; mod select; mod structural; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -40,7 +42,8 @@ pub(crate) const ROOT_SCHEME_ID: SchemeId = SchemeId { /// embedding a custom fixed scheme list or testing scheme interactions. #[derive(Debug, Clone)] pub struct CascadingCompressor { - /// The enabled compression schemes. + /// The active compression schemes: those given, minus any that another given scheme + /// [replaces](Scheme::replaces). schemes: Vec<&'static dyn Scheme>, /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from @@ -51,8 +54,22 @@ pub struct CascadingCompressor { impl CascadingCompressor { /// Creates a new compressor with the given schemes. /// + /// A scheme that another given scheme [replaces](Scheme::replaces) is dropped, so two + /// versions of one scheme never compete. Restrict the list to the schemes whose serialized + /// IDs the writer permits before calling this, so a newer version that is not permitted is + /// gone before replacement and the version it replaces stays. + /// /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + let replaced: HashSet = schemes + .iter() + .flat_map(|scheme| scheme.replaces()) + .collect(); + let schemes = schemes + .into_iter() + .filter(|scheme| !replaced.contains(&scheme.id())) + .collect(); + // Root exclusion: exclude IntDict from list/listview offsets (monotonically // increasing data where dictionary encoding is wasteful). let root_exclusions = vec![DescendantExclusion { @@ -66,7 +83,15 @@ impl CascadingCompressor { } } - /// Returns whether the compressor was configured with `scheme`. + /// The schemes active for compression, in registration order. + pub fn schemes(&self) -> &[&'static dyn Scheme] { + &self.schemes + } + + /// Returns whether `scheme` is active for compression. + /// + /// A scheme given to [`new`](Self::new) is inactive when another given scheme + /// [replaces](Scheme::replaces) it. pub fn has_scheme(&self, scheme: SchemeId) -> bool { self.schemes .iter() @@ -78,3 +103,6 @@ impl CascadingCompressor { #[cfg(test)] mod tests; + +#[cfg(test)] +mod replacement_tests; diff --git a/vortex-compressor/src/compressor/replacement_tests.rs b/vortex-compressor/src/compressor/replacement_tests.rs new file mode 100644 index 00000000000..5150e2e2dd6 --- /dev/null +++ b/vortex-compressor/src/compressor/replacement_tests.rs @@ -0,0 +1,154 @@ +// 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_error::VortexResult; +use vortex_session::registry::CachedId; + +use super::*; +use crate::scheme::CompressionEstimate; +use crate::scheme::CompressorContext; +use crate::scheme::EstimateVerdict; +use crate::stats::ArrayAndStats; + +static V1_ID: CachedId = CachedId::new("test.format_v1"); +static V2_ID: CachedId = CachedId::new("test.format_v2"); +static V3_ID: CachedId = CachedId::new("test.format_v3"); +static OTHER_ID: CachedId = CachedId::new("test.other"); + +/// A scheme version. `produced` are the wire IDs it writes and `replaces` the versions it +/// supersedes. +struct TestScheme { + name: &'static str, + produced: &'static [&'static CachedId], + replaces: &'static [&'static TestScheme], +} + +impl std::fmt::Debug for TestScheme { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.name) + } +} + +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 { + self.produced.iter().map(|id| ***id).collect() + } + + fn replaces(&self) -> Vec { + self.replaces.iter().map(|scheme| scheme.id()).collect() + } + + fn num_children(&self) -> usize { + 1 + } + + 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()) + } +} + +/// The frozen format. +static V1: TestScheme = TestScheme { + name: "test.scheme_v1", + produced: &[&V1_ID], + replaces: &[], +}; +/// Writes the frozen format for some values and a new one for others, like wide decimals. +static V2: TestScheme = TestScheme { + name: "test.scheme_v2", + produced: &[&V1_ID, &V2_ID], + replaces: &[&V1], +}; +/// A format that stands alone. It names only the version before it. +static V3: TestScheme = TestScheme { + name: "test.scheme_v3", + produced: &[&V3_ID], + replaces: &[&V2], +}; +static OTHER: TestScheme = TestScheme { + name: "test.other", + produced: &[&OTHER_ID], + replaces: &[], +}; + +fn active(compressor: &CascadingCompressor) -> Vec { + compressor + .schemes() + .iter() + .map(|scheme| scheme.id()) + .collect() +} + +#[test] +fn a_replacement_drops_the_schemes_it_names() { + assert_eq!( + active(&CascadingCompressor::new(vec![&V2, &V1])), + vec![V2.id()] + ); + assert_eq!(active(&CascadingCompressor::new(vec![&V1])), vec![V1.id()]); +} + +/// V3 names only V2, but V2 names V1, and the lists of dropped schemes still apply. +#[test] +fn replacement_is_transitive() { + assert_eq!( + active(&CascadingCompressor::new(vec![&V3, &V2, &V1])), + vec![V3.id()] + ); +} + +/// Without V2 in the list, nothing names V1, so V3 and V1 are both active. +#[test] +fn replacement_only_follows_given_schemes() { + assert_eq!( + active(&CascadingCompressor::new(vec![&V3, &V1])), + vec![V3.id(), V1.id()] + ); +} + +#[test] +fn registration_order_is_preserved() { + assert_eq!( + active(&CascadingCompressor::new(vec![&OTHER, &V3, &V2, &V1])), + vec![OTHER.id(), V3.id()] + ); +} + +#[test] +fn replacing_an_unregistered_scheme_is_a_no_op() { + assert_eq!(active(&CascadingCompressor::new(vec![&V2])), vec![V2.id()]); +} + +#[test] +fn has_scheme_reports_the_active_version() { + let compressor = CascadingCompressor::new(vec![&V2, &V1]); + assert!(compressor.has_scheme(V2.id())); + assert!(!compressor.has_scheme(V1.id())); +} diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..676b3f3d8bf 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -133,8 +133,26 @@ pub trait Scheme: Debug + Send + Sync { /// /// For most encodings this is the in-memory encoding ID. An encoding with several wire /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. + /// A newer scheme version that writes additional IDs declares all of them and names the + /// version it [`replaces`](Scheme::replaces). fn produced_encodings(&self) -> Vec; + /// Schemes this one supersedes. + /// + /// Register a newer version of a scheme alongside the version it replaces, and the two never + /// compete: the compressor drops every listed scheme it was given together with this one. + /// The writer restricts the scheme list to permitted serialized IDs first, so a newer version + /// whose [`produced_encodings`](Scheme::produced_encodings) are not all permitted is gone + /// before replacement and the listed schemes stay and compress as before. + /// + /// Listing a scheme that is not registered has no effect. Exclusion rules and [`has_scheme`] + /// name one version, so a replacement declares its own rules rather than inheriting them. + /// + /// [`has_scheme`]: crate::compressor::CascadingCompressor::has_scheme + fn replaces(&self) -> Vec { + vec![] + } + /// 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 640de874d2b..f740aacb318 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();