Replace all usages of Registry<T> with ArcSwapMap

docs/developer-guide/internals/session.md

@@ -15,10 +15,11 @@ The session is built on two primitives from the `vortex-session` crate:
 - **`VortexSession`** -- a cloneable, thread-safe map from Rust `TypeId` to a shared
   (`Arc`-wrapped) value. Any type that is `Clone + Send + Sync + Debug + 'static` can be stored
   as a session variable.
-- **`Registry<T>`** -- a concurrent map from string IDs to values of type `T`, used by each
-  component to look up registered plugins at runtime.
+- **`ArcSwapMap<K, V>`** -- a concurrent copy-on-write map, used both as the backing store of
+  `VortexSession` itself and (keyed by interned string `Id`s) as the plugin registries each
+  component uses to look up registered plugins at runtime.
 
-Because `VortexSession` is backed by an `ArcSwap`, cloning is cheap and all clones share the
+Because both primitives are backed by an `ArcSwap`, cloning is cheap and all clones share the
 same state, with lock-free reads and copy-on-write writes. This makes it safe to hand the
 session to multiple threads, tasks, or I/O operations without coordination.
 

encodings/zstd/src/zstd_buffers.rs

@@ -109,9 +109,10 @@ impl ZstdBuffers {
         buffer_handles: &[BufferHandle],
         session: &VortexSession,
     ) -> VortexResult<ArrayRef> {
-        let registry = session.arrays().registry().clone();
-        let inner_vtable = registry
-            .find(&array.data().inner_encoding_id)
+        let inner_vtable = session
+            .arrays()
+            .registry()
+            .get(&array.data().inner_encoding_id)
             .ok_or_else(|| {
                 vortex_err!("Unknown inner encoding: {}", array.data().inner_encoding_id)
             })?;

vortex-array/src/aggregate_fn/session.rs

@@ -4,6 +4,7 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
@@ -35,7 +36,6 @@ use crate::aggregate_fn::fns::sum::Sum;
 use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
 use crate::aggregate_fn::kernels::DynAggregateKernel;
 use crate::aggregate_fn::kernels::DynGroupedAggregateKernel;
-use crate::arc_swap_map::ArcSwapMap;
 use crate::array::ArrayId;
 use crate::array::VTable;
 use crate::arrays::Chunked;

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_valid_ids(session.arrays().registry().ids());
+        let ctx = ArrayContext::empty().with_allowed_ids(session.arrays().registry().keys());
         let serialized = array
             .serialize(&ctx, &session, &SerializeOptions::default())
             .unwrap();

vortex-array/src/dtype/serde/flatbuffers.rs

@@ -270,7 +270,7 @@ impl TryFrom<ViewedDType> for DType {
                         vortex_err!("failed to parse extension metadata from flatbuffer")
                     })?
                     .bytes();
-                let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().find(&id) {
+                let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().get(&id) {
                     vtable.deserialize(metadata, storage_dtype)?
                 } else if vfdt.session.allows_unknown() {
                     ForeignExtDType::from_parts(id, metadata.to_vec(), storage_dtype)?

vortex-array/src/dtype/serde/proto.rs

@@ -116,7 +116,7 @@ impl DType {
                         .ok_or_else(|| vortex_err!("Extension DType missing storage proto"))?,
                     session,
                 )?;
-                let ext_dtype = if let Some(vtable) = session.dtypes().registry().find(&id) {
+                let ext_dtype = if let Some(vtable) = session.dtypes().registry().get(&id) {
                     vtable.deserialize(e.metadata(), storage_dtype)?
                 } else if session.allows_unknown() {
                     ForeignExtDType::from_parts(id, e.metadata().to_vec(), storage_dtype)?

vortex-array/src/dtype/serde/serde.rs

@@ -722,7 +722,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, ExtDTypeRef> {
                     storage_dtype.ok_or_else(|| de::Error::missing_field("storage_dtype"))?;
                 let metadata = metadata.ok_or_else(|| de::Error::missing_field("metadata"))?;
 
-                if let Some(vtable) = self.session.dtypes().registry().find(&id) {
+                if let Some(vtable) = self.session.dtypes().registry().get(&id) {
                     vtable.deserialize(&metadata, storage_dtype).map_err(|e| {
                         de::Error::custom(format!(
                             "failed to deserialize extension dtype {}: {}",

vortex-array/src/dtype/session.rs

@@ -6,10 +6,11 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
-use vortex_session::registry::Registry;
+use vortex_session::registry::Id;
 
 use crate::dtype::extension::ExtDTypePluginRef;
 use crate::dtype::extension::ExtVTable;
@@ -18,19 +19,16 @@ use crate::extension::datetime::Time;
 use crate::extension::datetime::Timestamp;
 use crate::extension::uuid::Uuid;
 
-/// Registry for extension dtypes.
-pub type ExtDTypeRegistry = Registry<ExtDTypePluginRef>;
-
 /// Session for managing extension dtypes.
 #[derive(Clone, Debug)]
 pub struct DTypeSession {
-    registry: ExtDTypeRegistry,
+    registry: ArcSwapMap<Id, ExtDTypePluginRef>,
 }
 
 impl Default for DTypeSession {
     fn default() -> Self {
         let this = Self {
-            registry: Registry::default(),
+            registry: ArcSwapMap::default(),
         };
 
         // Register built-in temporal extension dtypes
@@ -57,11 +55,11 @@ impl DTypeSession {
     /// Register an extension DType with the Vortex session.
     pub fn register<V: ExtVTable>(&self, vtable: V) {
         self.registry
-            .register(vtable.id(), Arc::new(vtable) as ExtDTypePluginRef);
+            .insert(vtable.id(), Arc::new(vtable) as ExtDTypePluginRef);
     }
 
     /// Return the registry of extension dtypes.
-    pub fn registry(&self) -> &ExtDTypeRegistry {
+    pub fn registry(&self) -> &ArcSwapMap<Id, ExtDTypePluginRef> {
         &self.registry
     }
 }

vortex-array/src/expr/proto.rs

@@ -47,7 +47,7 @@ impl Expression {
             .map(|e| Expression::from_proto(e, session))
             .collect::<VortexResult<Vec<_>>>()?;
 
-        let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().find(&expr_id) {
+        let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().get(&expr_id) {
             vtable.deserialize(expr.metadata(), session)?
         } else if session.allows_unknown() {
             ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len())

vortex-array/src/lib.rs

@@ -4,7 +4,7 @@
 //!
 //! At the heart of Vortex are [arrays](ArrayRef).
 //!
-//! Arrays are typed views of memory buffers that hold [scalars](crate::scalar::Scalar). These
+//! Arrays are typed views of memory buffers that hold [scalars](scalar::Scalar). These
 //! buffers can be held in a number of physical encodings to perform lightweight compression that
 //! exploits the particular data distribution of the array's values.
 //!
@@ -51,10 +51,10 @@
 //!
 //! # Nulls and Scalars
 //!
-//! [`Validity`](crate::validity::Validity) separates nullness from values. It can be a cheap
+//! [`Validity`](validity::Validity) separates nullness from values. It can be a cheap
 //! constant state (`NonNullable`, `AllValid`, `AllInvalid`) or a boolean array that may itself be
-//! encoded. [`Scalar`](crate::scalar::Scalar) is the single-value counterpart: it pairs a
-//! [`DType`] with an optional [`ScalarValue`](crate::scalar::ScalarValue).
+//! encoded. [`Scalar`](scalar::Scalar) is the single-value counterpart: it pairs a
+//! [`DType`] with an optional [`ScalarValue`](scalar::ScalarValue).
 //!
 //! # Extending Vortex
 //!
@@ -66,16 +66,16 @@
 //! - [`OperationsVTable`] provides scalar access.
 //! - [`ValidityVTable`] exposes validity only for nullable arrays.
 //!
-//! New logical extension dtypes implement [`ExtVTable`](crate::dtype::extension::ExtVTable) and
+//! New logical extension dtypes implement [`ExtVTable`](dtype::extension::ExtVTable) and
 //! store values in an ordinary Vortex storage dtype.
 //!
-//! [`PrimitiveArray`]: crate::arrays::PrimitiveArray
-//! [`DType`]: crate::dtype::DType
-//! [`ChunkedArray`]: crate::arrays::ChunkedArray
-//! [`ConstantArray`]: crate::arrays::ConstantArray
-//! [`FilterArray`]: crate::arrays::FilterArray
-//! [`SliceArray`]: crate::arrays::SliceArray
-//! [`ScalarFnArray`]: crate::arrays::ScalarFnArray
+//! [`PrimitiveArray`]: arrays::PrimitiveArray
+//! [`DType`]: dtype::DType
+//! [`ChunkedArray`]: arrays::ChunkedArray
+//! [`ConstantArray`]: arrays::ConstantArray
+//! [`FilterArray`]: arrays::FilterArray
+//! [`SliceArray`]: arrays::SliceArray
+//! [`ScalarFnArray`]: arrays::ScalarFnArray
 
 extern crate self as vortex_array;
 
@@ -92,7 +92,7 @@ pub use smallvec;
 pub use vortex_array_macros::array_slots;
 use vortex_session::SessionExt;
 use vortex_session::VortexSession;
-use vortex_session::registry::Context;
+use vortex_session::registry::Interner;
 
 use crate::aggregate_fn::session::AggregateFnSession;
 use crate::dtype::session::DTypeSession;
@@ -105,7 +105,6 @@ use crate::stats::session::StatsSession;
 pub mod aggregate_fn;
 #[doc(hidden)]
 pub mod aliases;
-pub mod arc_swap_map;
 mod array;
 pub mod arrays;
 pub mod buffer;
@@ -149,7 +148,7 @@ pub mod flatbuffers {
 }
 
 /// Register vortex-array's built-in session-scoped kernels into the active
-/// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry.
+/// [`ArrayKernels`](optimizer::kernels::ArrayKernels) registry.
 ///
 /// If the session contains a [`KernelSession`], this registers into its registry. Sessions that use
 /// [`KernelSession::default`] already receive these built-in kernels.
@@ -190,4 +189,4 @@ pub fn legacy_session() -> &'static VortexSession {
     &LEGACY_SESSION
 }
 
-pub type ArrayContext = Context;
+pub type ArrayContext = Interner;

vortex-array/src/optimizer/kernels.rs

@@ -32,6 +32,7 @@ use std::sync::Arc;
 use std::sync::LazyLock;
 
 use vortex_error::VortexResult;
+use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
@@ -42,7 +43,6 @@ use vortex_utils::aliases::hash_map::HashMap;
 
 use crate::ArrayRef;
 use crate::ExecutionCtx;
-use crate::arc_swap_map::ArcSwapMap;
 use crate::array::VTable;
 use crate::arrays::Struct;
 use crate::arrays::struct_::compute::rules::struct_cast_reduce_parent;

vortex-array/src/scalar_fn/session.rs

@@ -4,10 +4,11 @@
 use std::any::Any;
 use std::sync::Arc;
 
+use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
-use vortex_session::registry::Registry;
+use vortex_session::registry::Id;
 
 use crate::scalar_fn::ScalarFnPluginRef;
 use crate::scalar_fn::ScalarFnVTable;
@@ -32,31 +33,28 @@ use crate::scalar_fn::fns::select::Select;
 use crate::scalar_fn::fns::stat::StatFn;
 use crate::scalar_fn::fns::variant_get::VariantGet;
 
-/// Registry of scalar function vtables.
-pub type ScalarFnRegistry = Registry<ScalarFnPluginRef>;
-
 /// Session state for scalar function vtables and rewrite rules.
 #[derive(Clone, Debug)]
 pub struct ScalarFnSession {
-    registry: ScalarFnRegistry,
+    registry: ArcSwapMap<Id, ScalarFnPluginRef>,
 }
 
 impl ScalarFnSession {
-    pub fn registry(&self) -> &ScalarFnRegistry {
+    pub fn registry(&self) -> &ArcSwapMap<Id, ScalarFnPluginRef> {
         &self.registry
     }
 
     /// Register a scalar function vtable in the session, replacing any existing vtable with the same ID.
     pub fn register<V: ScalarFnVTable>(&self, vtable: V) {
         self.registry
-            .register(vtable.id(), Arc::new(vtable) as ScalarFnPluginRef);
+            .insert(vtable.id(), Arc::new(vtable) as ScalarFnPluginRef);
     }
 }
 
 impl Default for ScalarFnSession {
     fn default() -> Self {
         let this = Self {
-            registry: ScalarFnRegistry::default(),
+            registry: ArcSwapMap::default(),
         };
 
         // Register built-in expressions.

vortex-array/src/serde.rs

@@ -325,7 +325,7 @@ impl SerializedArray {
         let encoding_id = ctx
             .resolve(encoding_idx)
             .ok_or_else(|| vortex_err!("Unknown encoding index: {}", encoding_idx))?;
-        let Some(plugin) = session.arrays().registry().find(&encoding_id) else {
+        let Some(plugin) = session.arrays().registry().get(&encoding_id) else {
             if session.allows_unknown() {
                 return self.decode_foreign(encoding_id, dtype, len, ctx);
             }

vortex-array/src/session/mod.rs

@@ -6,10 +6,11 @@ use std::sync::Arc;
 
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
-use vortex_session::registry::Registry;
+use vortex_session::registry::Id;
 
 use crate::ArrayRef;
 use crate::array::ArrayPlugin;
@@ -33,36 +34,34 @@ use crate::arrays::VarBin;
 use crate::arrays::VarBinView;
 use crate::arrays::Variant;
 
-pub type ArrayRegistry = Registry<ArrayPluginRef>;
-
 #[derive(Clone, Debug)]
 pub struct ArraySession {
     /// The set of registered array encodings.
-    registry: ArrayRegistry,
+    registry: ArcSwapMap<Id, ArrayPluginRef>,
 }
 
 impl ArraySession {
     pub fn empty() -> ArraySession {
         Self {
-            registry: ArrayRegistry::default(),
+            registry: ArcSwapMap::default(),
         }
     }
 
-    pub fn registry(&self) -> &ArrayRegistry {
+    pub fn registry(&self) -> &ArcSwapMap<Id, ArrayPluginRef> {
         &self.registry
     }
 
     /// Register a new array encoding, replacing any existing encoding with the same ID.
     pub fn register<P: ArrayPlugin>(&self, plugin: P) {
         self.registry
-            .register(plugin.id(), Arc::new(plugin) as ArrayPluginRef);
+            .insert(plugin.id(), Arc::new(plugin) as ArrayPluginRef);
     }
 }
 
 impl Default for ArraySession {
     fn default() -> Self {
         let this = ArraySession {
-            registry: ArrayRegistry::default(),
+            registry: ArcSwapMap::default(),
         };
 
         // Register the canonical encodings.
@@ -110,7 +109,7 @@ pub trait ArraySessionExt: SessionExt {
 
     /// Serialize an array using a plugin from the registry.
     fn array_serialize(&self, array: &ArrayRef) -> VortexResult<Option<Vec<u8>>> {
-        let Some(plugin) = self.arrays().registry.find(&array.encoding_id()) else {
+        let Some(plugin) = self.arrays().registry.get(&array.encoding_id()) else {
             vortex_bail!(
                 "Array {} is not registered for serializations",
                 array.encoding_id()
@@ -136,13 +135,13 @@ mod tests {
     fn array_session_default_registers_encodings() {
         let session = VortexSession::empty().with::<ArraySession>();
 
-        assert!(session.arrays().registry().find(&Bool.id()).is_some());
+        assert!(session.arrays().registry().contains_key(&Bool.id()));
     }
 
     #[test]
     fn empty_array_session_registers_no_encodings() {
         let session = VortexSession::empty().with_some(ArraySession::empty());
 
-        assert!(session.arrays().registry().find(&Bool.id()).is_none());
+        assert!(!session.arrays().registry().contains_key(&Bool.id()));
     }
 }

vortex-arrow/src/session.rs

@@ -39,7 +39,6 @@ use tracing::trace;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
-use vortex_array::arc_swap_map::ArcSwapMap;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::ListArray;
 use vortex_array::arrays::ListViewArray;
@@ -56,6 +55,7 @@ use vortex_array::validity::Validity;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
+use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;