Add ArcSwapMap and use it throughout our session registries

vortex-array/src/aggregate_fn/accumulator.rs

@@ -142,21 +142,16 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
         }
 
         let session = ctx.session().clone();
-        let kernels = &session.aggregate_fns().kernels;
 
         // 1. Kernel registry first: a registered `(encoding, aggregate_fn)` kernel is strictly
         //    more specific than the vtable's `try_accumulate` short-circuit. Checking the
         //    registry first gives kernels for `Combined<V>` aggregates a chance to fire —
         //    `Combined::try_accumulate` always returns true, so a later kernel check would be
         //    unreachable.
         {
-            let kernels_r = kernels.read();
-            let batch_id = batch.encoding_id();
-            let kernel = kernels_r
-                .get(&(batch_id, Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(batch_id, None)))
-                .copied();
-            drop(kernels_r);
+            let kernel = session
+                .aggregate_fns()
+                .find_aggregate_kernel(batch.encoding_id(), self.aggregate_fn.id());
             if let Some(kernel) = kernel
                 && let Some(result) = kernel.aggregate(&self.aggregate_fn, batch, ctx)?
             {
@@ -187,14 +182,9 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
                 break;
             }
 
-            let kernels_r = kernels.read();
-            let batch_id = batch.encoding_id();
-            let kernel = kernels_r
-                .get(&(batch_id, Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(batch_id, None)))
-                .copied();
-            drop(kernels_r);
-            if let Some(kernel) = kernel
+            if let Some(kernel) = session
+                .aggregate_fns()
+                .find_aggregate_kernel(batch.encoding_id(), self.aggregate_fn.id())
                 && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)?
             {
                 vortex_ensure!(

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -163,17 +163,15 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
         let mut elements = groups.elements().clone();
         let groups_validity = groups.validity()?;
         let session = ctx.session().clone();
-        let kernels = &session.aggregate_fns().grouped_kernels;
 
         for _ in 0..max_iterations() {
             if elements.is::<AnyCanonical>() {
                 break;
             }
 
-            let kernels_r = kernels.read();
-            if let Some(result) = kernels_r
-                .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(elements.encoding_id(), None)))
+            if let Some(result) = session
+                .aggregate_fns()
+                .find_grouped_kernel(elements.encoding_id(), self.aggregate_fn.id())
                 .and_then(|kernel| {
                     // SAFETY: we assume that elements execution is safe
                     let groups = unsafe {
@@ -255,17 +253,15 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
         let mut elements = groups.elements().clone();
         let groups_validity = groups.validity()?;
         let session = ctx.session().clone();
-        let kernels = &session.aggregate_fns().grouped_kernels;
 
         for _ in 0..64 {
             if elements.is::<AnyCanonical>() {
                 break;
             }
 
-            let kernels_r = kernels.read();
-            if let Some(result) = kernels_r
-                .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(elements.encoding_id(), None)))
+            if let Some(result) = session
+                .aggregate_fns()
+                .find_grouped_kernel(elements.encoding_id(), self.aggregate_fn.id())
                 .and_then(|kernel| {
                     // SAFETY: we assume that elements execution is safe
                     let groups = unsafe {

vortex-array/src/aggregate_fn/proto.rs

@@ -36,7 +36,7 @@ impl AggregateFnRef {
     /// Note: the serialization format is not stable and may change between versions.
     pub fn from_proto(proto: &pb::AggregateFn, session: &VortexSession) -> VortexResult<Self> {
         let agg_fn_id: AggregateFnId = AggregateFnId::new(proto.id.as_str());
-        let agg_fn = if let Some(plugin) = session.aggregate_fns().registry().find(&agg_fn_id) {
+        let agg_fn = if let Some(plugin) = session.aggregate_fns().find_plugin(&agg_fn_id) {
             plugin.deserialize(proto.metadata(), session)?
         } else if session.allows_unknown() {
             new_foreign_aggregate_fn(agg_fn_id, proto.metadata().to_vec())

vortex-array/src/aggregate_fn/session.rs

@@ -4,12 +4,9 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use parking_lot::RwLock;
 use vortex_session::Ref;
 use vortex_session::SessionExt;
 use vortex_session::SessionVar;
-use vortex_session::registry::Registry;
-use vortex_utils::aliases::hash_map::HashMap;
 
 use crate::aggregate_fn::AggregateFnId;
 use crate::aggregate_fn::AggregateFnPluginRef;
@@ -34,6 +31,7 @@ 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;
@@ -43,16 +41,17 @@ use crate::arrays::dict::compute::is_constant::DictIsConstantKernel;
 use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel;
 use crate::arrays::dict::compute::min_max::DictMinMaxKernel;
 
-/// Registry of aggregate function vtables.
-pub type AggregateFnRegistry = Registry<AggregateFnPluginRef>;
-
-/// Session state for aggregate function vtables.
+/// Session state for aggregate functions and encoding-specific aggregate kernels.
+///
+/// The default session registers the built-in aggregate functions and kernels. Additional
+/// aggregate functions and kernels may be registered by extensions when they are added to a
+/// [`VortexSession`](vortex_session::VortexSession).
 #[derive(Debug)]
 pub struct AggregateFnSession {
-    registry: AggregateFnRegistry,
+    registry: ArcSwapMap<AggregateFnId, AggregateFnPluginRef>,
 
-    pub(super) kernels: RwLock<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
-    pub(super) grouped_kernels: RwLock<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
+    kernels: ArcSwapMap<KernelKey, &'static dyn DynAggregateKernel>,
+    grouped_kernels: ArcSwapMap<KernelKey, &'static dyn DynGroupedAggregateKernel>,
 }
 
 impl SessionVar for AggregateFnSession {
@@ -70,9 +69,9 @@ type KernelKey = (ArrayId, Option<AggregateFnId>);
 impl Default for AggregateFnSession {
     fn default() -> Self {
         let this = Self {
-            registry: AggregateFnRegistry::default(),
-            kernels: RwLock::new(HashMap::default()),
-            grouped_kernels: RwLock::new(HashMap::default()),
+            registry: ArcSwapMap::default(),
+            kernels: ArcSwapMap::default(),
+            grouped_kernels: ArcSwapMap::default(),
         };
 
         // Register the built-in aggregate functions
@@ -106,34 +105,88 @@ impl Default for AggregateFnSession {
 }
 
 impl AggregateFnSession {
-    /// Returns the aggregate function registry.
-    pub fn registry(&self) -> &AggregateFnRegistry {
-        &self.registry
+    /// Returns the aggregate function plugin registered for `id`, if any.
+    pub fn find_plugin(&self, id: &AggregateFnId) -> Option<AggregateFnPluginRef> {
+        self.registry.get(id)
     }
 
     /// Register an aggregate function vtable in the session, replacing any existing vtable with
     /// the same ID.
     pub fn register<V: AggregateFnVTable>(&self, vtable: V) {
-        self.registry
-            .register(vtable.id(), Arc::new(vtable) as AggregateFnPluginRef);
+        let id = vtable.id();
+        let pluginref = Arc::new(vtable) as AggregateFnPluginRef;
+        self.registry.insert(id, pluginref);
     }
 
-    /// Register an aggregate function kernel for a specific aggregate function and array type.
+    /// Returns the aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
+    ///
+    /// Lookup first checks for a kernel registered for the exact aggregate function, then falls
+    /// back to a kernel registered for all aggregate functions on the same array encoding.
+    pub fn find_aggregate_kernel(
+        &self,
+        array_id: impl Into<ArrayId>,
+        agg_fn_id: impl Into<AggregateFnId>,
+    ) -> Option<&'static dyn DynAggregateKernel> {
+        let id = array_id.into();
+        let fn_id = agg_fn_id.into();
+        self.kernels.read(|kernels| {
+            kernels
+                .get(&(id, Some(fn_id)))
+                .or_else(|| kernels.get(&(id, None)))
+                .copied()
+        })
+    }
+
+    /// Registers an aggregate kernel for an array encoding.
+    ///
+    /// When `agg_fn_id` is `Some`, the kernel is used only for that aggregate function. When
+    /// `agg_fn_id` is `None`, the kernel is used as the fallback for aggregate functions on the
+    /// array encoding that do not have a more specific kernel.
     pub fn register_aggregate_kernel(
         &self,
         array_id: impl Into<ArrayId>,
         agg_fn_id: Option<impl Into<AggregateFnId>>,
         kernel: &'static dyn DynAggregateKernel,
     ) {
-        self.kernels
-            .write()
-            .insert((array_id.into(), agg_fn_id.map(|id| id.into())), kernel);
+        let id = (array_id.into(), agg_fn_id.map(|id| id.into()));
+        self.kernels.insert(id, kernel);
+    }
+
+    /// Returns the grouped aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
+    ///
+    /// Lookup first checks for a kernel registered for the exact aggregate function, then falls
+    /// back to a kernel registered for all aggregate functions on the same array encoding.
+    pub fn find_grouped_kernel(
+        &self,
+        array_id: impl Into<ArrayId>,
+        agg_fn_id: impl Into<AggregateFnId>,
+    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
+        let id = array_id.into();
+        let fn_id = agg_fn_id.into();
+        self.grouped_kernels.read(|kernels| {
+            kernels
+                .get(&(id, Some(fn_id)))
+                .or_else(|| kernels.get(&(id, None)))
+                .copied()
+        })
+    }
+
+    /// Registers a grouped aggregate kernel for a specific aggregate function and array encoding.
+    pub fn register_grouped_kernel(
+        &self,
+        array_id: impl Into<ArrayId>,
+        agg_fn_id: impl Into<AggregateFnId>,
+        kernel: &'static dyn DynGroupedAggregateKernel,
+    ) {
+        let id = array_id.into();
+        let fn_id = agg_fn_id.into();
+        self.grouped_kernels.insert((id, Some(fn_id)), kernel)
     }
 }
 
 /// Extension trait for accessing aggregate function session data.
 pub trait AggregateFnSessionExt: SessionExt {
-    /// Returns the aggregate function vtable registry.
+    /// Returns the aggregate function session data.
     fn aggregate_fns(&self) -> Ref<'_, AggregateFnSession> {
         self.get::<AggregateFnSession>()
     }