Require schemes to declare produced encodings

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

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

@@ -31,8 +31,8 @@ impl Scheme for ZstdScheme {
         canonical.dtype().is_binary()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![vortex_zstd::Zstd.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_zstd::Zstd.id()]
     }
 
     fn expected_compression_ratio(

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

@@ -31,8 +31,8 @@ impl Scheme for ZstdBuffersScheme {
         canonical.dtype().is_binary()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![vortex_zstd::ZstdBuffers.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_zstd::ZstdBuffers.id()]
     }
 
     fn expected_compression_ratio(

vortex-btrblocks/src/schemes/decimal.rs

@@ -40,8 +40,8 @@ impl Scheme for DecimalScheme {
         matches!(canonical, Canonical::Decimal(_))
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![DecimalByteParts.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![DecimalByteParts.id()]
     }
 
     /// Children: primitive=0.

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

@@ -42,12 +42,12 @@ impl Scheme for ALPScheme {
         canonical.dtype().is_float()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
+    fn produced_encodings(&self) -> Vec<ArrayId> {
         let mut encodings = vec![ALP.id()];
         if use_experimental_patches() {
             encodings.push(Patched.id());
         }
-        Some(encodings)
+        encodings
     }
 
     /// Children: encoded_ints=0.

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

@@ -39,8 +39,8 @@ impl Scheme for ALPRDScheme {
         canonical.dtype().is_float()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![vortex_alp::ALPRD.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_alp::ALPRD.id()]
     }
 
     fn expected_compression_ratio(

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

@@ -31,8 +31,8 @@ impl Scheme for PcoScheme {
         canonical.dtype().is_float()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![vortex_pco::Pco.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![vortex_pco::Pco.id()]
     }
 
     fn expected_compression_ratio(

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

@@ -38,8 +38,8 @@ impl Scheme for FloatRLEScheme {
         canonical.dtype().is_float()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![RLE.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![RLE.id()]
     }
 
     /// Children: values=0, indices=1, offsets=2.

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

@@ -41,8 +41,8 @@ impl Scheme for NullDominatedSparseScheme {
         canonical.dtype().is_float()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
-        Some(vec![Sparse.id()])
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        vec![Sparse.id()]
     }
 
     /// Children: indices=0.

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

@@ -40,12 +40,12 @@ impl Scheme for BitPackingScheme {
         canonical.dtype().is_int()
     }
 
-    fn produced_encodings(&self) -> Option<Vec<ArrayId>> {
+    fn produced_encodings(&self) -> Vec<ArrayId> {
         let mut encodings = vec![BitPacked.id()];
         if use_experimental_patches() {
             encodings.push(Patched.id());
         }
-        Some(encodings)
+        encodings
     }
 
     fn expected_compression_ratio(