refactor: store valid IDs in serialization context (#8897)

The serialization context only needs registry IDs for filtering, so it
now snapshots those IDs instead of retaining the entire registry.

---------

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

Author: joseph-isaacs

vortex-array/src/arrays/patched/vtable/mod.rs

@@ -588,7 +588,7 @@ mod tests {
         let session = array_session();
         session.arrays().register(Patched);
 
-        let ctx = ArrayContext::empty().with_registry(session.arrays().registry().clone());
+        let ctx = ArrayContext::empty().with_valid_ids(session.arrays().registry().ids());
         let serialized = array
             .serialize(&ctx, &session, &SerializeOptions::default())
             .unwrap();

vortex-array/src/lib.rs

@@ -190,4 +190,4 @@ pub fn legacy_session() -> &'static VortexSession {
     &LEGACY_SESSION
 }
 
-pub type ArrayContext = Context<ArrayPluginRef>;
+pub type ArrayContext = Context;

vortex-file/src/writer.rs

@@ -200,8 +200,8 @@ impl VortexWriteOptions {
         // 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 ctx = ArrayContext::new(ALLOWED_ENCODINGS.iter().cloned().sorted().collect())
-            // Configure a registry just to ensure only known encodings are interned.
-            .with_registry(self.session.arrays().registry().clone());
+            // Only permit encodings known to the session.
+            .with_valid_ids(self.session.arrays().registry().ids());
         let dtype = stream.dtype().clone();
 
         let (mut ptr, eof) = SequenceId::root().split();

vortex-layout/src/lib.rs

@@ -42,4 +42,4 @@ mod strategy;
 mod test;
 pub mod vtable;
 
-pub type LayoutContext = Context<LayoutEncodingRef>;
+pub type LayoutContext = Context;

vortex-session/src/registry.rs

@@ -21,6 +21,7 @@ use parking_lot::RwLock;
 use vortex_error::VortexExpect;
 use vortex_utils::aliases::DefaultHashBuilder;
 use vortex_utils::aliases::dash_map::DashMap;
+use vortex_utils::aliases::hash_set::HashSet;
 
 /// Global string interner for [`Id`] values.
 static INTERNER: LazyLock<ThreadedRodeo<Spur, DefaultHashBuilder>> =
@@ -227,32 +228,32 @@ impl ReadContext {
 /// 1. This object holds an Arc of RwLock internally because we need concurrent access from the
 ///    layout writer code path. We should update SegmentSink to take an Array rather than
 ///    ByteBuffer such that serializing arrays is done sequentially.
-/// 2. The name is terrible. `Interner<T>` is better, but I want to minimize breakage for now.
+/// 2. The name is terrible. `Interner` is better, but I want to minimize breakage for now.
 #[derive(Clone, Debug)]
-pub struct Context<T> {
+pub struct Context {
     // TODO(ngates): it's a long story, but if we make SegmentSink and SegmentSource take an
     //  enum of Segment { Array, DType, Buffer } then we don't actually need a mutable context
     //  in the LayoutWriter, therefore we don't need a RwLock here and everyone is happier.
     ids: Arc<RwLock<Vec<Id>>>,
-    // Optional registry used to filter the permissible interned items.
-    registry: Option<Registry<T>>,
+    // Optional set used to filter the permissible interned IDs.
+    valid_ids: Option<Arc<HashSet<Id>>>,
 }
 
-impl<T> Default for Context<T> {
+impl Default for Context {
     fn default() -> Self {
         Self {
             ids: Arc::new(RwLock::new(Vec::new())),
-            registry: None,
+            valid_ids: None,
         }
     }
 }
 
-impl<T: Clone> Context<T> {
+impl Context {
     /// Create a context with the given initial IDs.
     pub fn new(ids: Vec<Id>) -> Self {
         Self {
             ids: Arc::new(RwLock::new(ids)),
-            registry: None,
+            valid_ids: None,
         }
     }
 
@@ -261,18 +262,18 @@ impl<T: Clone> Context<T> {
         Self::default()
     }
 
-    /// Configure a registry to restrict the permissible set of interned items.
-    pub fn with_registry(mut self, registry: Registry<T>) -> Self {
-        self.registry = Some(registry);
+    /// Restrict the permissible set of interned IDs.
+    pub fn with_valid_ids(mut self, ids: impl IntoIterator<Item = Id>) -> Self {
+        self.valid_ids = Some(Arc::new(ids.into_iter().collect()));
         self
     }
 
     /// Intern an ID, returning its index.
     pub fn intern(&self, id: &Id) -> Option<u16> {
-        if let Some(registry) = &self.registry
-            && registry.find(id).is_none()
+        if let Some(valid_ids) = &self.valid_ids
+            && !valid_ids.contains(id)
         {
-            // ID not in registry, cannot intern.
+            // ID is not valid, cannot intern.
             return None;
         }
 
@@ -295,3 +296,24 @@ impl<T: Clone> Context<T> {
         self.ids.read().clone()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::CachedId;
+    use super::Context;
+
+    static VALID: CachedId = CachedId::new("vortex.test.valid");
+    static INVALID: CachedId = CachedId::new("vortex.test.invalid");
+
+    #[test]
+    fn context_filters_interned_ids() {
+        let valid = *VALID;
+        let invalid = *INVALID;
+        let context = Context::empty().with_valid_ids([valid]);
+
+        assert_eq!(context.intern(&valid), Some(0));
+        assert_eq!(context.intern(&valid), Some(0));
+        assert_eq!(context.intern(&invalid), None);
+        assert_eq!(context.to_ids(), [valid]);
+    }
+}