Use ArcSwap for aggregate fn registry

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-array/src/aggregate_fn/accumulator.rs

@@ -142,21 +142,19 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
         }
 
         let session = ctx.session().clone();
-        let kernels = &session.aggregate_fns().kernels;
+        let kernels = &session.aggregate_fns().kernels.load();
 
         // 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
+            let kernel = kernels
                 .get(&(batch_id, Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(batch_id, None)))
+                .or_else(|| kernels.get(&(batch_id, None)))
                 .copied();
-            drop(kernels_r);
             if let Some(kernel) = kernel
                 && let Some(result) = kernel.aggregate(&self.aggregate_fn, batch, ctx)?
             {
@@ -187,13 +185,11 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
                 break;
             }
 
-            let kernels_r = kernels.read();
             let batch_id = batch.encoding_id();
-            let kernel = kernels_r
+            let kernel = kernels
                 .get(&(batch_id, Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(batch_id, None)))
+                .or_else(|| kernels.get(&(batch_id, None)))
                 .copied();
-            drop(kernels_r);
             if let Some(kernel) = kernel
                 && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)?
             {

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -163,17 +163,16 @@ 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;
+        let kernels = &session.aggregate_fns().grouped_kernels.load();
 
         for _ in 0..max_iterations() {
             if elements.is::<AnyCanonical>() {
                 break;
             }
 
-            let kernels_r = kernels.read();
-            if let Some(result) = kernels_r
+            if let Some(result) = kernels
                 .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(elements.encoding_id(), None)))
+                .or_else(|| kernels.get(&(elements.encoding_id(), None)))
                 .and_then(|kernel| {
                     // SAFETY: we assume that elements execution is safe
                     let groups = unsafe {
@@ -255,17 +254,16 @@ 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;
+        let kernels = &session.aggregate_fns().grouped_kernels.load();
 
         for _ in 0..64 {
             if elements.is::<AnyCanonical>() {
                 break;
             }
 
-            let kernels_r = kernels.read();
-            if let Some(result) = kernels_r
+            if let Some(result) = kernels
                 .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
-                .or_else(|| kernels_r.get(&(elements.encoding_id(), None)))
+                .or_else(|| kernels.get(&(elements.encoding_id(), None)))
                 .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,7 +4,7 @@
 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;
@@ -49,10 +49,10 @@ 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>>,
+    pub(super) kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
+    pub(super) grouped_kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
 }
 
 impl SessionVar for AggregateFnSession {
@@ -70,9 +70,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
@@ -107,15 +107,20 @@ impl Default for AggregateFnSession {
 
 impl AggregateFnSession {
     /// Returns the aggregate function registry.
-    pub fn registry(&self) -> &AggregateFnRegistry {
-        &self.registry
+    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
+        });
     }
 
     /// Register an aggregate function kernel for a specific aggregate function and array type.
@@ -125,9 +130,12 @@ 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
+        });
     }
 }