Use ArcSwap for aggregate fn registry

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,11 +4,10 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use parking_lot::RwLock;
+use arc_swap::ArcSwap;
 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;
@@ -43,16 +42,13 @@ 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.
 #[derive(Debug)]
 pub struct AggregateFnSession {
-    registry: AggregateFnRegistry,
+    registry: ArcSwap<HashMap<AggregateFnId, AggregateFnPluginRef>>,
 
-    pub(super) kernels: RwLock<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
-    pub(super) grouped_kernels: RwLock<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
+    kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
+    grouped_kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
 }
 
 impl SessionVar for AggregateFnSession {
@@ -70,9 +66,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: ArcSwap::from_pointee(HashMap::default()),
+            kernels: ArcSwap::from_pointee(HashMap::default()),
+            grouped_kernels: ArcSwap::from_pointee(HashMap::default()),
         };
 
         // Register the built-in aggregate functions
@@ -106,16 +102,35 @@ impl Default for AggregateFnSession {
 }
 
 impl AggregateFnSession {
-    /// Returns the aggregate function registry.
-    pub fn registry(&self) -> &AggregateFnRegistry {
-        &self.registry
+    /// Find plugin in the registry for the given id
+    pub fn find_plugin(&self, id: &AggregateFnId) -> Option<AggregateFnPluginRef> {
+        self.registry.load().get(id).cloned()
     }
 
     /// 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.rcu(move |registry| {
+            let mut existing = registry.as_ref().clone();
+            existing.insert(id, Arc::clone(&pluginref));
+            existing
+        });
+    }
+
+    pub fn find_aggregate_kernel(
+        &self,
+        array_id: impl Into<ArrayId>,
+        agg_fn_id: impl Into<AggregateFnId>,
+    ) -> Option<&'static dyn DynAggregateKernel> {
+        let loaded = self.kernels.load();
+        let id = array_id.into();
+        let fn_id = agg_fn_id.into();
+        loaded
+            .get(&(id, Some(fn_id)))
+            .or_else(|| loaded.get(&(id, None)))
+            .copied()
     }
 
     /// Register an aggregate function kernel for a specific aggregate function and array type.
@@ -125,9 +140,26 @@ impl AggregateFnSession {
         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.rcu(move |registry| {
+            let mut existing = registry.as_ref().clone();
+            existing.insert(id, kernel);
+            existing
+        });
+    }
+
+    pub fn find_grouped_kernel(
+        &self,
+        array_id: impl Into<ArrayId>,
+        agg_fn_id: impl Into<AggregateFnId>,
+    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
+        let loaded = self.grouped_kernels.load();
+        let id = array_id.into();
+        let fn_id = agg_fn_id.into();
+        loaded
+            .get(&(id, Some(fn_id)))
+            .or_else(|| loaded.get(&(id, None)))
+            .copied()
     }
 }