fix

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

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()));
 }

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<ArrayId> {
+            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<ArrayRef> {
+            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(())
+    }
+}

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)

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 {

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<ArrayId> = 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.