Tell the compressor which serialized IDs the writer may emit

docs/specs/editions.md

@@ -95,8 +95,10 @@ unrestricted. It does not register missing readers, so files written this way ha
 
 Compression and edition compatibility are separate. Compressors produce current in-memory arrays and do not select a
 wire ID. The writer maps each allowed serialized ID to its current in-memory encoding and restricts the default
-BtrBlocks compressor to schemes producing those encodings. Custom compressors remain unrestricted, with serialization
-providing the final compatibility boundary when edition enforcement is enabled. At that boundary, the array plugin
+BtrBlocks compressor to schemes producing those encodings. It also hands the compressor the allowed serialized IDs
+themselves, so a scheme whose encoding has more than one wire format can produce the newest form those IDs permit
+instead of one the serializer would refuse; the compressor never reads editions directly. Custom compressors remain
+unrestricted, with serialization providing the final compatibility boundary when edition enforcement is enabled. At that boundary, the array plugin
 produces an ID, metadata, buffers, and children. The serialization context interns the returned ID and fails the write
 if the selected editions do not permit it. A serializer may emit a historical ID when the value satisfies that ID's
 frozen contract, but it does not inspect the edition allowlist. Without disabling edition enforcement, a custom layout

vortex-btrblocks/src/builder.rs

@@ -91,12 +91,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
 #[derive(Debug, Clone)]
 pub struct BtrBlocksCompressorBuilder {
     schemes: Vec<&'static dyn Scheme>,
+    allowed_serialized_ids: Option<HashSet<ArrayId>>,
 }
 
 impl Default for BtrBlocksCompressorBuilder {
     fn default() -> Self {
         Self {
             schemes: ALL_SCHEMES.to_vec(),
+            allowed_serialized_ids: None,
         }
     }
 }
@@ -108,6 +110,7 @@ impl BtrBlocksCompressorBuilder {
     pub fn empty() -> Self {
         Self {
             schemes: Vec::new(),
+            allowed_serialized_ids: None,
         }
     }
 
@@ -188,6 +191,14 @@ impl BtrBlocksCompressorBuilder {
         #[cfg(feature = "pco")]
         excluded.extend([integer::PcoScheme.id(), float::PcoScheme.id()]);
         let builder = self.exclude_schemes(excluded);
+        // CUDA kernels decode each encoding's original wire format, whose ID is the encoding's
+        // own. An encoding with newer formats therefore keeps producing the original one here.
+        let original_formats = builder
+            .schemes
+            .iter()
+            .flat_map(|scheme| scheme.produced_encodings())
+            .collect();
+        let builder = builder.allow_serialized_ids(&original_formats);
 
         #[cfg(all(feature = "zstd", feature = "unstable_encodings"))]
         let builder = builder.with_new_scheme(&binary::ZstdBuffersScheme);
@@ -214,15 +225,33 @@ impl BtrBlocksCompressorBuilder {
         self
     }
 
+    /// Restricts the output to the serialized IDs in `allowed`, intersecting with any earlier
+    /// restriction, so a scheme whose encoding has several wire formats produces the newest one
+    /// still allowed.
+    ///
+    /// The file writer passes the serialized IDs of its configured editions.
+    pub fn allow_serialized_ids(mut self, allowed: &HashSet<ArrayId>) -> Self {
+        self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() {
+            Some(existing) => existing.intersection(allowed).copied().collect(),
+            None => allowed.clone(),
+        });
+        self
+    }
+
     /// Builds the configured [`BtrBlocksCompressor`].
     pub fn build(self) -> BtrBlocksCompressor {
-        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
+        let compressor = CascadingCompressor::new(self.schemes);
+        BtrBlocksCompressor(match self.allowed_serialized_ids {
+            Some(allowed) => compressor.with_allowed_serialized_ids(allowed),
+            None => compressor,
+        })
     }
 }
 
 #[cfg(test)]
 mod tests {
     use vortex_array::VTable;
+    use vortex_array::arrays::Bool;
     use vortex_fastlanes::FoR;
 
     use super::*;
@@ -287,6 +316,20 @@ mod tests {
         }
     }
 
+    /// Every serialized ID is allowed by default. The CUDA preset allows only the original format
+    /// of each encoding its schemes produce.
+    #[test]
+    fn cuda_compatible_allows_only_original_formats() {
+        let default = BtrBlocksCompressorBuilder::default().build();
+        assert!(default.0.allows_serialized_id(Bool.id()));
+
+        let cuda = BtrBlocksCompressorBuilder::default()
+            .only_cuda_compatible()
+            .build();
+        assert!(cuda.0.allows_serialized_id(FoR.id()));
+        assert!(!cuda.0.allows_serialized_id(Bool.id()));
+    }
+
     #[test]
     fn cuda_compatible_uses_fsst_for_strings() {
         let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible();

vortex-compressor/src/compressor/mod.rs

@@ -9,6 +9,9 @@ mod sample;
 mod select;
 mod structural;
 
+use vortex_array::ArrayId;
+use vortex_utils::aliases::hash_set::HashSet;
+
 use crate::builtins::IntDictScheme;
 use crate::scheme::ChildSelection;
 use crate::scheme::DescendantExclusion;
@@ -46,6 +49,10 @@ pub struct CascadingCompressor {
     /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from
     /// list offsets).
     root_exclusions: Vec<DescendantExclusion>,
+
+    /// The serialized IDs the output may use, or `None` for no restriction. See
+    /// [`allows_serialized_id`](Self::allows_serialized_id).
+    allowed_serialized_ids: Option<HashSet<ArrayId>>,
 }
 
 impl CascadingCompressor {
@@ -63,9 +70,34 @@ impl CascadingCompressor {
         Self {
             schemes,
             root_exclusions,
+            allowed_serialized_ids: None,
         }
     }
 
+    /// Restricts the output to the serialized IDs in `allowed`, intersecting with any earlier
+    /// restriction.
+    ///
+    /// The file writer passes the serialized IDs its editions permit, so a scheme whose encoding
+    /// has several wire formats produces the newest one still allowed.
+    pub fn with_allowed_serialized_ids(mut self, allowed: HashSet<ArrayId>) -> Self {
+        self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() {
+            Some(existing) => existing.intersection(&allowed).copied().collect(),
+            None => allowed,
+        });
+        self
+    }
+
+    /// Returns whether the output may use the serialized ID `id`.
+    ///
+    /// Schemes whose encoding has several wire formats consult this to pick the newest allowed
+    /// one. Without a restriction every ID is allowed, so that is always the newest. The
+    /// serialization context still validates whatever ID a scheme ends up producing.
+    pub fn allows_serialized_id(&self, id: ArrayId) -> bool {
+        self.allowed_serialized_ids
+            .as_ref()
+            .is_none_or(|allowed| allowed.contains(&id))
+    }
+
     /// Returns whether the compressor was configured with `scheme`.
     pub fn has_scheme(&self, scheme: SchemeId) -> bool {
         self.schemes

vortex-compressor/src/compressor/tests.rs

@@ -9,11 +9,14 @@ 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::arrays::Bool;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::Constant;
 use vortex_array::arrays::Map;
 use vortex_array::arrays::NullArray;
+use vortex_array::arrays::Primitive;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::assert_arrays_eq;
 use vortex_array::builders::MapBuilder;
@@ -26,6 +29,7 @@ use vortex_array::validity::Validity;
 use vortex_buffer::buffer;
 use vortex_error::VortexResult;
 use vortex_session::VortexSession;
+use vortex_utils::aliases::hash_set::HashSet;
 
 use super::CascadingCompressor;
 use super::ROOT_SCHEME_ID;
@@ -96,6 +100,48 @@ impl Scheme for DirectRatioScheme {
     }
 }
 
+/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`.
+static SEEN_FORMAT: Mutex<Option<bool>> = Mutex::new(None);
+
+/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether
+/// the newer one is allowed and records the answer.
+#[derive(Debug)]
+struct FormatRecordingScheme;
+
+impl Scheme for FormatRecordingScheme {
+    fn scheme_name(&self) -> &'static str {
+        "test.format_recording"
+    }
+
+    fn matches(&self, canonical: &Canonical) -> bool {
+        matches_integer_primitive(canonical)
+    }
+
+    fn produced_encodings(&self) -> Vec<ArrayId> {
+        Vec::new()
+    }
+
+    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> {
+        *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id()));
+        Ok(data.array().clone())
+    }
+}
+
 #[derive(Debug)]
 struct ImmediateAlwaysUseScheme;
 
@@ -841,3 +887,38 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> {
     assert_arrays_eq!(&compressed, &array, &mut exec_ctx);
     Ok(())
 }
+
+#[test]
+fn allowed_serialized_ids_default_to_everything_and_intersect() {
+    let compressor = compressor();
+    assert!(compressor.allows_serialized_id(Constant.id()));
+    assert!(compressor.allows_serialized_id(Bool.id()));
+
+    let restricted =
+        compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()]));
+    assert!(restricted.allows_serialized_id(Constant.id()));
+    assert!(!restricted.allows_serialized_id(Bool.id()));
+
+    let narrowed =
+        restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()]));
+    assert!(narrowed.allows_serialized_id(Primitive.id()));
+    assert!(!narrowed.allows_serialized_id(Constant.id()));
+    assert!(!narrowed.allows_serialized_id(Bool.id()));
+}
+
+/// A scheme sees the restriction through the compressor it is handed: everything is allowed until
+/// the writer narrows the set to its editions.
+#[test]
+fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> {
+    let array = PrimitiveArray::from_iter(0..4096i32).into_array();
+    let mut exec_ctx = SESSION.create_execution_ctx();
+
+    let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]);
+    unrestricted.compress(&array, &mut exec_ctx)?;
+    assert_eq!(*SEEN_FORMAT.lock(), Some(true));
+
+    let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()]));
+    restricted.compress(&array, &mut exec_ctx)?;
+    assert_eq!(*SEEN_FORMAT.lock(), Some(false));
+    Ok(())
+}

vortex-file/src/writer.rs

@@ -239,7 +239,7 @@ impl VortexWriteOptions {
         let enforce_editions = !self.disable_editions;
         // The array context is built here, rather than when the options were constructed, so that
         // encodings registered on the session in between are still eligible for the file.
-        let (array_ctx, allowed_array_encodings) =
+        let (array_ctx, allowed_array_encodings, allowed_serialized_ids) =
             new_array_context(&self.session, enforce_editions);
         let ctx = LayoutWriterContext::new(array_ctx)
             .with_buffered_bytes_tracker(self.buffered_bytes.clone());
@@ -253,7 +253,8 @@ impl VortexWriteOptions {
             None => WriteStrategyBuilder::default()
                 .with_btrblocks_builder(
                     BtrBlocksCompressorBuilder::default()
-                        .retain_allowed_encodings(&allowed_array_encodings),
+                        .retain_allowed_encodings(&allowed_array_encodings)
+                        .allow_serialized_ids(&allowed_serialized_ids),
                 )
                 .build(),
         };
@@ -384,10 +385,12 @@ impl VortexWriteOptions {
     }
 }
 
+/// Returns the array context, the in-memory encodings the compressor may produce, and the
+/// serialized IDs its output may use.
 fn new_array_context(
     session: &VortexSession,
     enforce_editions: bool,
-) -> (ArrayContext, HashSet<ArrayId>) {
+) -> (ArrayContext, HashSet<ArrayId>, HashSet<ArrayId>) {
     // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated.
     // This is preferred for now over having an empty context here, because only the
     // serialised array order is deterministic. The serialisation of arrays are done
@@ -406,14 +409,17 @@ fn new_array_context(
         .filter_map(|serialized_id| arrays.registry().get(serialized_id))
         .map(|plugin| plugin.id())
         .collect();
+    // The compressor sees the same set, so an encoding with several wire formats produces the
+    // newest one the editions permit.
+    let allowed_serialized_ids: HashSet<ArrayId> = serialized_ids.iter().copied().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.
         array_ctx.with_allowed_ids(serialized_ids.into_iter().collect())
     } else {
         array_ctx
     };
-    (array_ctx, allowed_array_encodings)
+    (array_ctx, allowed_array_encodings, allowed_serialized_ids)
 }
 
 /// The ids of `kind` the enabled editions permit.
@@ -787,10 +793,12 @@ mod tests {
         session.register_edition(&DECLARATION)?;
         session.enable_edition(EDITION)?;
 
-        let (ctx, allowed_array_encodings) = new_array_context(&session, true);
+        let (ctx, allowed_array_encodings, allowed_serialized_ids) =
+            new_array_context(&session, true);
         assert_eq!(ctx.to_ids(), [Primitive.id()]);
         assert!(ctx.intern(&Bool.id()).is_none());
         assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()]));
+        assert_eq!(allowed_serialized_ids, HashSet::from([Primitive.id()]));
         Ok(())
     }
 
@@ -807,9 +815,14 @@ mod tests {
             )
         });
 
-        let (ctx, allowed_array_encodings) = new_array_context(&session, false);
+        let (ctx, allowed_array_encodings, allowed_serialized_ids) =
+            new_array_context(&session, false);
         assert_eq!(ctx.to_ids(), registered_ids);
         assert_eq!(allowed_array_encodings, registered_encodings);
+        assert_eq!(
+            allowed_serialized_ids,
+            registered_ids.iter().copied().collect::<HashSet<_>>()
+        );
         assert!(ctx.intern(&Bool.id()).is_some());
     }