Filter compression schemes by serialized IDs instead of array IDs

vortex-btrblocks/src/builder.rs

@@ -200,10 +200,10 @@ impl BtrBlocksCompressorBuilder {
         self
     }
 
-    /// Retains only schemes whose produced encodings all belong to `allowed`.
+    /// Retains only schemes whose produced serialized IDs all belong to `allowed`.
     ///
-    /// The file writer uses this to restrict compression to the encodings of its configured
-    /// editions.
+    /// `allowed` holds serialized IDs. 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)));

vortex-compressor/src/scheme/mod.rs

@@ -124,11 +124,15 @@ pub trait Scheme: Debug + Send + Sync {
     /// Whether this scheme can compress the given canonical array.
     fn matches(&self, canonical: &Canonical) -> bool;
 
-    /// The array encodings this scheme itself may introduce into its compressed output.
+    /// The serialized IDs this scheme itself may write into its compressed output.
     ///
-    /// Cascaded children are compressed by other schemes, which declare their own encodings,
-    /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here.
-    /// Canonical arrays the scheme merely rearranges do not need to be declared.
+    /// Every declared ID must be permitted for the scheme to be used. Cascaded children are
+    /// compressed by other schemes, which declare their own IDs, so only arrays constructed
+    /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme
+    /// merely rearranges do not need to be declared.
+    ///
+    /// For most encodings this is the in-memory encoding ID. An encoding with several wire
+    /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID.
     fn produced_encodings(&self) -> Vec<ArrayId>;
 
     /// Returns the stats generation options this scheme requires. The compressor merges all

vortex-file/src/writer.rs

@@ -239,21 +239,22 @@ 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) =
-            new_array_context(&self.session, enforce_editions);
+        let array_ctx = new_array_context(&self.session, enforce_editions);
         let ctx = LayoutWriterContext::new(array_ctx)
             .with_buffered_bytes_tracker(self.buffered_bytes.clone());
         let ctx = if enforce_editions {
             ctx.with_allowed_aggregates(edition_filter(&self.session, ComponentKind::Aggregate))
         } else {
             ctx
         };
+        let allowed_serialized_ids: HashSet<ArrayId> =
+            ctx.array_ctx().to_ids().into_iter().collect();
         let strategy = match self.strategy {
             Some(strategy) => strategy,
             None => WriteStrategyBuilder::default()
                 .with_btrblocks_builder(
                     BtrBlocksCompressorBuilder::default()
-                        .retain_allowed_encodings(&allowed_array_encodings),
+                        .retain_allowed_encodings(&allowed_serialized_ids),
                 )
                 .build(),
         };
@@ -384,36 +385,30 @@ impl VortexWriteOptions {
     }
 }
 
-fn new_array_context(
-    session: &VortexSession,
-    enforce_editions: bool,
-) -> (ArrayContext, HashSet<ArrayId>) {
+fn new_array_context(session: &VortexSession, enforce_editions: bool) -> ArrayContext {
     // 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
     // parallel and with an empty context they can register their encodings to the context
     // in different order, changing the written bytes from run to run.
-    let arrays = session.arrays();
-    let serialized_ids = if enforce_editions {
+    //
+    // The seeded IDs are also what the writer may emit: callers read them back with
+    // `ArrayContext::to_ids` to restrict compression to the same set.
+    let serialized_ids: Vec<ArrayId> = if enforce_editions {
         session.enabled_component_ids(ComponentKind::Array)
     } else {
-        arrays
+        session
+            .arrays()
             .registry()
             .read(|registry| registry.keys().copied().collect())
     };
-    let allowed_array_encodings = serialized_ids
-        .iter()
-        .filter_map(|serialized_id| arrays.registry().get(serialized_id))
-        .map(|plugin| plugin.id())
-        .collect();
     let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect());
-    let array_ctx = if enforce_editions {
+    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)
+    }
 }
 
 /// The ids of `kind` the enabled editions permit.
@@ -787,29 +782,22 @@ mod tests {
         session.register_edition(&DECLARATION)?;
         session.enable_edition(EDITION)?;
 
-        let (ctx, allowed_array_encodings) = new_array_context(&session, true);
+        let ctx = 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()]));
         Ok(())
     }
 
     #[test]
     fn disabling_editions_allows_all_registered_array_ids() {
         let session = array_session();
-        let (registered_ids, registered_encodings) = session.arrays().registry().read(|registry| {
-            (
-                registry.keys().copied().sorted().collect::<Vec<_>>(),
-                registry
-                    .values()
-                    .map(|plugin| plugin.id())
-                    .collect::<HashSet<_>>(),
-            )
-        });
+        let registered_ids = session
+            .arrays()
+            .registry()
+            .read(|registry| registry.keys().copied().sorted().collect::<Vec<_>>());
 
-        let (ctx, allowed_array_encodings) = new_array_context(&session, false);
+        let ctx = new_array_context(&session, false);
         assert_eq!(ctx.to_ids(), registered_ids);
-        assert_eq!(allowed_array_encodings, registered_encodings);
         assert!(ctx.intern(&Bool.id()).is_some());
     }