feat[compressor]: require schemes to declare their produced encodings

vortex-btrblocks/src/builder.rs

@@ -3,6 +3,7 @@
 
 //! Builder for configuring `BtrBlocksCompressor` instances.
 
+use vortex_array::ArrayId;
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::BtrBlocksCompressor;
@@ -208,6 +209,16 @@ impl BtrBlocksCompressorBuilder {
         self
     }
 
+    /// Retains only schemes whose produced encodings all belong to `allowed`.
+    ///
+    /// The file writer uses this to restrict compression to the encodings of its configured
+    /// editions.
+    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
+    }
+
     /// Builds the configured [`BtrBlocksCompressor`].
     pub fn build(self) -> BtrBlocksCompressor {
         BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
@@ -216,6 +227,9 @@ impl BtrBlocksCompressorBuilder {
 
 #[cfg(test)]
 mod tests {
+    use vortex_array::VTable;
+    use vortex_fastlanes::FoR;
+
     use super::*;
 
     #[test]
@@ -230,6 +244,27 @@ mod tests {
         assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
     }
 
+    #[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 none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new());
+        assert!(none.schemes.is_empty());
+    }
+
+    #[test]
+    fn retaining_all_declared_outputs_keeps_every_scheme() {
+        let allowed: HashSet<ArrayId> = ALL_SCHEMES
+            .iter()
+            .flat_map(|scheme| scheme.produced_encodings())
+            .collect();
+        let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed);
+        assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
+    }
+
     #[test]
     fn cuda_compatible_excludes_alprd() {
         let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();

vortex-btrblocks/src/schemes/binary/zstd.rs

@@ -3,10 +3,12 @@
 
 //! Zstd compression for binary arrays.
 
+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_compressor::scheme::CompressionEstimate;
 use vortex_compressor::scheme::DeferredEstimate;
 use vortex_error::VortexResult;
@@ -29,6 +31,10 @@ impl Scheme for ZstdScheme {
         canonical.dtype().is_binary()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_zstd::Zstd.id()]
+    }
+
     fn expected_compression_ratio(
         &self,
         _data: &ArrayAndStats,

vortex-btrblocks/src/schemes/binary/zstd_buffers.rs

@@ -3,10 +3,12 @@
 
 //! Zstd buffer-level binary compression preserving array layout for GPU decompression.
 
+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_compressor::scheme::CompressionEstimate;
 use vortex_compressor::scheme::DeferredEstimate;
 use vortex_error::VortexResult;
@@ -29,6 +31,10 @@ impl Scheme for ZstdBuffersScheme {
         canonical.dtype().is_binary()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_zstd::ZstdBuffers.id()]
+    }
+
     fn expected_compression_ratio(
         &self,
         _data: &ArrayAndStats,

vortex-btrblocks/src/schemes/decimal.rs

@@ -3,10 +3,12 @@
 
 //! Decimal compression scheme using byte-part decomposition.
 
+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::arrays::DecimalArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::decimal::narrowed_decimal;
@@ -38,6 +40,10 @@ impl Scheme for DecimalScheme {
         matches!(canonical, Canonical::Decimal(_))
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![DecimalByteParts.id()]
+    }
+
     /// Children: primitive=0.
     fn num_children(&self) -> usize {
         1

vortex-btrblocks/src/schemes/float/alp.rs

@@ -7,10 +7,12 @@ use vortex_alp::ALP;
 use vortex_alp::ALPArrayExt;
 use vortex_alp::ALPArraySlotsExt;
 use vortex_alp::alp_encode;
+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::arrays::Patched;
 use vortex_array::arrays::patched::use_experimental_patches;
 use vortex_array::arrays::primitive::PrimitiveArrayExt;
@@ -40,6 +42,14 @@ impl Scheme for ALPScheme {
         canonical.dtype().is_float()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        let mut encodings = vec![ALP.id()];
+        if use_experimental_patches() {
+            encodings.push(Patched.id());
+        }
+        encodings
+    }
+
     /// Children: encoded_ints=0.
     fn num_children(&self) -> usize {
         1

vortex-btrblocks/src/schemes/float/alprd.rs

@@ -6,10 +6,12 @@
 use vortex_alp::ALPRDArrayExt;
 use vortex_alp::ALPRDArrayOwnedExt;
 use vortex_alp::RDEncoder;
+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::arrays::primitive::PrimitiveArrayExt;
 use vortex_array::dtype::PType;
 use vortex_compressor::scheme::CompressionEstimate;
@@ -37,6 +39,10 @@ impl Scheme for ALPRDScheme {
         canonical.dtype().is_float()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_alp::ALPRD.id()]
+    }
+
     fn expected_compression_ratio(
         &self,
         data: &ArrayAndStats,

vortex-btrblocks/src/schemes/float/pco.rs

@@ -3,10 +3,12 @@
 
 //! Pco (pcodec) float compression.
 
+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_compressor::scheme::CompressionEstimate;
 use vortex_compressor::scheme::DeferredEstimate;
 use vortex_error::VortexResult;
@@ -29,6 +31,10 @@ impl Scheme for PcoScheme {
         canonical.dtype().is_float()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_pco::Pco.id()]
+    }
+
     fn expected_compression_ratio(
         &self,
         _data: &ArrayAndStats,

vortex-btrblocks/src/schemes/float/rle.rs

@@ -3,15 +3,18 @@
 
 //! Run-length float encoding.
 
+use vortex_array::ArrayId;
 use vortex_array::ArrayRef;
 use vortex_array::Canonical;
 use vortex_array::ExecutionCtx;
+use vortex_array::VTable;
 use vortex_compressor::scheme::AncestorExclusion;
 use vortex_compressor::scheme::CompressionEstimate;
 use vortex_compressor::scheme::DeferredEstimate;
 use vortex_compressor::scheme::DescendantExclusion;
 use vortex_compressor::scheme::EstimateVerdict;
 use vortex_error::VortexResult;
+use vortex_fastlanes::RLE;
 
 use crate::ArrayAndStats;
 use crate::CascadingCompressor;
@@ -35,6 +38,10 @@ impl Scheme for FloatRLEScheme {
         canonical.dtype().is_float()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![RLE.id()]
+    }
+
     /// Children: values=0, indices=1, offsets=2.
     fn num_children(&self) -> usize {
         3

vortex-btrblocks/src/schemes/float/sparse.rs

@@ -3,10 +3,12 @@
 
 //! Sparse encoding for null-dominated float arrays.
 
+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::arrays::PrimitiveArray;
 use vortex_array::arrays::primitive::PrimitiveArrayExt;
 use vortex_compressor::scheme::ChildSelection;
@@ -39,6 +41,10 @@ impl Scheme for NullDominatedSparseScheme {
         canonical.dtype().is_float()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![Sparse.id()]
+    }
+
     /// Children: indices=0.
     fn num_children(&self) -> usize {
         1

vortex-btrblocks/src/schemes/integer/bitpacking.rs

@@ -3,10 +3,12 @@
 
 //! BitPacking integer encoding.
 
+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::arrays::Patched;
 use vortex_array::arrays::patched::use_experimental_patches;
 use vortex_array::arrays::primitive::PrimitiveArrayExt;
@@ -38,6 +40,14 @@ impl Scheme for BitPackingScheme {
         canonical.dtype().is_int()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        let mut encodings = vec![BitPacked.id()];
+        if use_experimental_patches() {
+            encodings.push(Patched.id());
+        }
+        encodings
+    }
+
     fn expected_compression_ratio(
         &self,
         data: &ArrayAndStats,

vortex-btrblocks/src/schemes/integer/delta.rs

@@ -3,10 +3,12 @@
 
 //! FastLanes Delta integer encoding.
 
+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::arrays::PrimitiveArray;
 use vortex_compressor::builtins::BinaryDictScheme;
 use vortex_compressor::builtins::FloatDictScheme;
@@ -78,6 +80,10 @@ impl Scheme for DeltaScheme {
         canonical.dtype().is_int()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![Delta.id()]
+    }
+
     fn num_children(&self) -> usize {
         2
     }

vortex-btrblocks/src/schemes/integer/for_.rs

@@ -3,10 +3,12 @@
 
 //! Frame of Reference integer encoding.
 
+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::arrays::PrimitiveArray;
 use vortex_compressor::builtins::BinaryDictScheme;
 use vortex_compressor::builtins::FloatDictScheme;
@@ -41,6 +43,10 @@ impl Scheme for FoRScheme {
         canonical.dtype().is_int()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![FoR.id()]
+    }
+
     /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op.
     fn ancestor_exclusions(&self) -> Vec<AncestorExclusion> {
         vec![

vortex-btrblocks/src/schemes/integer/pco.rs

@@ -3,10 +3,12 @@
 
 //! Pco (pcodec) integer compression.
 
+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_compressor::scheme::CompressionEstimate;
 use vortex_compressor::scheme::DeferredEstimate;
 use vortex_compressor::scheme::EstimateVerdict;
@@ -30,6 +32,10 @@ impl Scheme for PcoScheme {
         canonical.dtype().is_int()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_pco::Pco.id()]
+    }
+
     fn expected_compression_ratio(
         &self,
         data: &ArrayAndStats,

vortex-btrblocks/src/schemes/integer/rle.rs

@@ -3,10 +3,12 @@
 
 //! Run-length integer encoding and shared RLE compression helpers.
 
+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::arrays::PrimitiveArray;
 use vortex_array::arrays::primitive::PrimitiveArrayExt;
 use vortex_compressor::scheme::AncestorExclusion;
@@ -156,6 +158,10 @@ impl Scheme for IntRLEScheme {
         canonical.dtype().is_int()
     }
 
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![RLE.id()]
+    }
+
     /// Children: values=0, indices=1, offsets=2.
     fn num_children(&self) -> usize {
         3