Handle array evolution in compressor schemes with replacement semantics

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<ArrayId>) -> Self {
         self.schemes
             .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id)));

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<ArrayId> {
+            vec![*DELTA_V2_ID]
+        }
+
+        fn replaces(&self) -> Vec<SchemeId> {
+            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<ArrayRef> {
+            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()));
+    }
+}

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<SchemeId> = 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;

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<ArrayId> {
+        self.produced.iter().map(|id| ***id).collect()
+    }
+
+    fn replaces(&self) -> Vec<SchemeId> {
+        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<ArrayRef> {
+        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<SchemeId> {
+    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()));
+}

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<ArrayId>;
 
+    /// 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<SchemeId> {
+        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.