Take the permitted serialized IDs from the array context

The array context already holds the allowed set as its interning guard. Interner::allowed_ids reads it back, so the writer builds the set once and hands the compressor the same one; a write with editions disabled leaves the compressor unrestricted instead of passing it the whole registry.

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

vortex-file/src/writer.rs

@@ -17,7 +17,6 @@ use futures::pin_mut;
 use futures::select;
 use itertools::Itertools;
 use vortex_array::ArrayContext;
-use vortex_array::ArrayId;
 use vortex_array::ArrayRef;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldPath;
@@ -239,24 +238,30 @@ 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_serialized_ids) =
-            new_array_context(&self.session, enforce_editions);
+        let array_ctx = new_array_context(&self.session, enforce_editions);
+        let strategy = match self.strategy {
+            Some(strategy) => strategy,
+            None => {
+                // The compressor sees the serialized IDs the context permits: it keeps the schemes
+                // that can write one of them, and an encoding with several wire formats produces
+                // the newest one permitted.
+                let builder = BtrBlocksCompressorBuilder::default();
+                let builder = match array_ctx.allowed_ids() {
+                    Some(allowed) => builder.allow_serialized_ids(allowed),
+                    None => builder,
+                };
+                WriteStrategyBuilder::default()
+                    .with_btrblocks_builder(builder)
+                    .build()
+            }
+        };
         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 strategy = match self.strategy {
-            Some(strategy) => strategy,
-            None => WriteStrategyBuilder::default()
-                .with_btrblocks_builder(
-                    BtrBlocksCompressorBuilder::default()
-                        .allow_serialized_ids(&allowed_serialized_ids),
-                )
-                .build(),
-        };
         let dtype = stream.dtype().clone();
         if enforce_editions {
             validate_dtype_editions(&self.session, &dtype)?;
@@ -384,11 +389,9 @@ impl VortexWriteOptions {
     }
 }
 
-/// Returns the array context and the serialized IDs the compressor may write its output under.
-fn new_array_context(
-    session: &VortexSession,
-    enforce_editions: bool,
-) -> (ArrayContext, HashSet<ArrayId>) {
+/// Returns the array context for a write, restricted to the serialized IDs of the enabled
+/// editions when they are enforced.
+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
@@ -402,17 +405,13 @@ fn new_array_context(
             .registry()
             .read(|registry| registry.keys().copied().collect())
     };
-    // The compressor sees the same set: it keeps the schemes that can write one of these IDs, and
-    // an encoding with several wire formats produces the newest one permitted.
-    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 {
+    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_serialized_ids)
+    }
 }
 
 /// The ids of `kind` the enabled editions permit.
@@ -786,10 +785,10 @@ mod tests {
         session.register_edition(&DECLARATION)?;
         session.enable_edition(EDITION)?;
 
-        let (ctx, allowed_serialized_ids) = 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_serialized_ids, HashSet::from([Primitive.id()]));
+        assert_eq!(ctx.allowed_ids(), Some(&HashSet::from([Primitive.id()])));
         Ok(())
     }
 
@@ -801,12 +800,9 @@ mod tests {
             .registry()
             .read(|registry| registry.keys().copied().sorted().collect::<Vec<_>>());
 
-        let (ctx, allowed_serialized_ids) = new_array_context(&session, false);
+        let ctx = new_array_context(&session, false);
         assert_eq!(ctx.to_ids(), registered_ids);
-        assert_eq!(
-            allowed_serialized_ids,
-            registered_ids.iter().copied().collect::<HashSet<_>>()
-        );
+        assert!(ctx.allowed_ids().is_none());
         assert!(ctx.intern(&Bool.id()).is_some());
     }
 

vortex-session/src/registry.rs

@@ -207,6 +207,11 @@ impl Interner {
         self
     }
 
+    /// The permissible set of interned IDs, or `None` when any ID may be interned.
+    pub fn allowed_ids(&self) -> Option<&HashSet<Id>> {
+        self.allowed.as_deref()
+    }
+
     /// Intern an ID, returning its index.
     pub fn intern(&self, id: &Id) -> Option<u16> {
         if self