Carry the writer's permitted serialized IDs into compression schemes

vortex-btrblocks/src/builder.rs

@@ -4,6 +4,7 @@
 //! Builder for configuring `BtrBlocksCompressor` instances.
 
 use vortex_array::ArrayId;
+use vortex_compressor::scheme::AllowedSerializedIds;
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::BtrBlocksCompressor;
@@ -90,12 +91,15 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
 #[derive(Debug, Clone)]
 pub struct BtrBlocksCompressorBuilder {
     schemes: Vec<&'static dyn Scheme>,
+    /// The serialized IDs the compressor may write under.
+    allowed_serialized_ids: AllowedSerializedIds,
 }
 
 impl Default for BtrBlocksCompressorBuilder {
     fn default() -> Self {
         Self {
             schemes: ALL_SCHEMES.to_vec(),
+            allowed_serialized_ids: AllowedSerializedIds::All,
         }
     }
 }
@@ -107,6 +111,7 @@ impl BtrBlocksCompressorBuilder {
     pub fn empty() -> Self {
         Self {
             schemes: Vec::new(),
+            allowed_serialized_ids: AllowedSerializedIds::All,
         }
     }
 
@@ -202,23 +207,29 @@ 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.
+    /// The set is also handed to the compressor, intersected with any earlier call, so a scheme
+    /// with several wire formats writes a newer one only when permitted. The file writer passes
+    /// the array IDs its enabled editions permit.
     pub fn retain_allowed_encodings(mut self, allowed: &HashSet<ArrayId>) -> Self {
         self.schemes
             .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id)));
+        self.allowed_serialized_ids.restrict(allowed);
         self
     }
 
     /// Builds the configured [`BtrBlocksCompressor`].
     pub fn build(self) -> BtrBlocksCompressor {
-        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
+        BtrBlocksCompressor(
+            CascadingCompressor::new(self.schemes)
+                .with_allowed_serialized_ids(&self.allowed_serialized_ids),
+        )
     }
 }
 
 #[cfg(test)]
 mod tests {
     use vortex_array::VTable;
+    use vortex_fastlanes::BitPacked;
     use vortex_fastlanes::FoR;
 
     use super::*;
@@ -238,12 +249,20 @@ mod tests {
     #[test]
     fn retain_allowed_encodings_filters_schemes() {
         let allowed: HashSet<ArrayId> = [FoR.id()].into_iter().collect();
-        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
-        assert_eq!(builder.schemes.len(), 1);
-        assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id());
+        let compressor = BtrBlocksCompressorBuilder::default()
+            .retain_allowed_encodings(&allowed)
+            .build();
+        assert_eq!(compressor.schemes().len(), 1);
+        assert_eq!(compressor.schemes()[0].id(), integer::FoRScheme.id());
+        assert_eq!(
+            compressor.allowed_serialized_ids(),
+            &AllowedSerializedIds::Only(allowed)
+        );
 
-        let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new());
-        assert!(none.schemes.is_empty());
+        let none = BtrBlocksCompressorBuilder::default()
+            .retain_allowed_encodings(&HashSet::new())
+            .build();
+        assert!(none.schemes().is_empty());
     }
 
     #[test]
@@ -252,8 +271,31 @@ mod tests {
             .iter()
             .flat_map(|scheme| scheme.produced_encodings())
             .collect();
-        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
-        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
+        let compressor = BtrBlocksCompressorBuilder::default()
+            .retain_allowed_encodings(&allowed)
+            .build();
+        assert_eq!(compressor.schemes().len(), ALL_SCHEMES.len());
+    }
+
+    #[test]
+    fn unrestricted_builds_permit_everything() {
+        let compressor = BtrBlocksCompressorBuilder::default().build();
+        assert_eq!(
+            compressor.allowed_serialized_ids(),
+            &AllowedSerializedIds::All
+        );
+    }
+
+    #[test]
+    fn repeated_restrictions_intersect() {
+        let first: HashSet<ArrayId> = [FoR.id(), BitPacked.id()].into_iter().collect();
+        let second: HashSet<ArrayId> = [BitPacked.id()].into_iter().collect();
+        let compressor = BtrBlocksCompressorBuilder::default()
+            .retain_allowed_encodings(&first)
+            .retain_allowed_encodings(&second)
+            .build();
+        assert!(!compressor.has_scheme(integer::FoRScheme.id()));
+        assert!(compressor.has_scheme(integer::BitPackingScheme.id()));
     }
 
     #[test]

vortex-btrblocks/tests/scheme_modes.rs

@@ -0,0 +1,131 @@
+// 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::BitPackingScheme;
+    use vortex_compressor::scheme::CompressionEstimate;
+    use vortex_compressor::scheme::EstimateVerdict;
+    use vortex_error::VortexResult;
+    use vortex_fastlanes::BitPacked;
+    use vortex_fastlanes::Delta;
+    use vortex_session::registry::CachedId;
+
+    static NEWER_ID: CachedId = CachedId::new("test.delta_newer");
+
+    /// A scheme with two wire formats. It always writes `Delta`, and when the writer permits the
+    /// newer format it takes a different branch. The newer branch stands in for a format this test
+    /// cannot serialize, so it returns the input unchanged and the compressor falls back to
+    /// canonical output.
+    #[derive(Debug)]
+    struct TwoFormatDelta;
+
+    impl Scheme for TwoFormatDelta {
+        fn scheme_name(&self) -> &'static str {
+            "test.two_format_delta"
+        }
+
+        fn matches(&self, canonical: &Canonical) -> bool {
+            canonical.dtype().is_int()
+        }
+
+        fn produced_encodings(&self) -> Vec<ArrayId> {
+            vec![Delta.id()]
+        }
+
+        /// Children: bases=0, deltas=1.
+        fn num_children(&self) -> usize {
+            2
+        }
+
+        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<ArrayRef> {
+            if compress_ctx.allows_serialized_id(&NEWER_ID) {
+                return Ok(data.array().clone());
+            }
+            let primitive = data.array().clone().execute::<PrimitiveArray>(exec_ctx)?;
+            let (bases, deltas) = vortex_fastlanes::delta_compress(&primitive, exec_ctx)?;
+            let bases = compressor.compress_child(
+                &bases.into_array(),
+                &compress_ctx,
+                self.id(),
+                0,
+                exec_ctx,
+            )?;
+            let deltas = compressor.compress_child(
+                &deltas.into_array(),
+                &compress_ctx,
+                self.id(),
+                1,
+                exec_ctx,
+            )?;
+            Delta::try_new(bases, deltas, 0, primitive.len()).map(IntoArray::into_array)
+        }
+    }
+
+    fn compressor(allowed: &[ArrayId]) -> vortex_btrblocks::BtrBlocksCompressor {
+        BtrBlocksCompressorBuilder::empty()
+            .with_new_scheme(&TwoFormatDelta)
+            .with_new_scheme(&BitPackingScheme)
+            .retain_allowed_encodings(&allowed.iter().copied().collect())
+            .build()
+    }
+
+    /// The scheme reads the writer's permitted IDs from its context and picks its format.
+    #[rstest]
+    #[case::frozen_only(vec![Delta.id(), BitPacked.id()], true)]
+    #[case::newer_permitted(vec![Delta.id(), BitPacked.id(), *NEWER_ID], false)]
+    fn a_scheme_picks_its_format_from_the_permitted_ids(
+        #[case] allowed: Vec<ArrayId>,
+        #[case] expect_delta: bool,
+    ) -> VortexResult<()> {
+        let session = array_session();
+        vortex_fastlanes::initialize(&session);
+        let compressor = compressor(&allowed);
+        let array = PrimitiveArray::from_iter((0..65_536u32).map(|i| i / 3)).into_array();
+        let mut ctx = session.create_execution_ctx();
+        let compressed = compressor.compress(&array, &mut ctx)?;
+        assert_eq!(compressed.encoding_id() == Delta.id(), expect_delta);
+        assert_arrays_eq!(compressed, array, &mut ctx);
+        Ok(())
+    }
+
+    /// The declared format is required. Permitting only the newer one leaves the scheme out.
+    #[test]
+    fn the_declared_format_must_be_permitted() {
+        let compressor = compressor(&[*NEWER_ID, BitPacked.id()]);
+        assert!(!compressor.has_scheme(TwoFormatDelta.id()));
+        assert!(compressor.has_scheme(BitPackingScheme.id()));
+    }
+}

vortex-compressor/src/compressor/cascade.rs

@@ -59,7 +59,7 @@ impl CascadingCompressor {
 
         let canonical = array.clone().execute::<CanonicalValidity>(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());
 

vortex-compressor/src/compressor/edition_tests.rs

@@ -0,0 +1,97 @@
+// 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::EstimateVerdict;
+use crate::stats::ArrayAndStats;
+
+static V1_ID: CachedId = CachedId::new("test.format_v1");
+static V2_ID: CachedId = CachedId::new("test.format_v2");
+
+/// A scheme that always writes `test.format_v1` and, when permitted, also `test.format_v2`.
+#[derive(Debug)]
+struct ModeScheme;
+
+impl Scheme for ModeScheme {
+    fn scheme_name(&self) -> &'static str {
+        "test.mode"
+    }
+
+    fn matches(&self, canonical: &Canonical) -> bool {
+        canonical.dtype().is_int()
+    }
+
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![*V1_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<ArrayRef> {
+        Ok(data.array().clone())
+    }
+}
+
+fn only(ids: &[&CachedId]) -> AllowedSerializedIds {
+    AllowedSerializedIds::Only(ids.iter().map(|id| ***id).collect())
+}
+
+#[test]
+fn unrestricted_contexts_permit_everything() {
+    let compressor = CascadingCompressor::new(vec![&ModeScheme]);
+    assert_eq!(
+        compressor.allowed_serialized_ids(),
+        &AllowedSerializedIds::All
+    );
+    let ctx = compressor.root_context();
+    assert!(ctx.allows_serialized_id(&V1_ID));
+    assert!(ctx.allows_serialized_id(&V2_ID));
+}
+
+#[test]
+fn the_permitted_set_reaches_descendant_contexts() {
+    let compressor =
+        CascadingCompressor::new(vec![&ModeScheme]).with_allowed_serialized_ids(&only(&[&V1_ID]));
+    let root = compressor.root_context();
+    assert!(root.allows_serialized_id(&V1_ID));
+    assert!(!root.allows_serialized_id(&V2_ID));
+
+    let child = root.descend_with_scheme(ModeScheme.id(), 0);
+    assert!(child.allows_serialized_id(&V1_ID));
+    assert!(!child.allows_serialized_id(&V2_ID));
+    assert_eq!(child.allowed_serialized_ids(), &only(&[&V1_ID]));
+}
+
+#[test]
+fn repeated_restrictions_intersect() {
+    let compressor = CascadingCompressor::new(vec![&ModeScheme])
+        .with_allowed_serialized_ids(&only(&[&V1_ID, &V2_ID]))
+        .with_allowed_serialized_ids(&only(&[&V1_ID]));
+    assert_eq!(compressor.allowed_serialized_ids(), &only(&[&V1_ID]));
+    assert!(!compressor.root_context().allows_serialized_id(&V2_ID));
+
+    // Intersecting with `All` changes nothing.
+    let compressor = compressor.with_allowed_serialized_ids(&AllowedSerializedIds::All);
+    assert_eq!(compressor.allowed_serialized_ids(), &only(&[&V1_ID]));
+}