more

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

Author: robert3005

vortex-array/src/aggregate_fn/accumulator.rs

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

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -163,16 +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.load();
+        let aggregate_fns = session.aggregate_fns();
 
         for _ in 0..max_iterations() {
             if elements.is::<AnyCanonical>() {
                 break;
             }
 
-            if let Some(result) = kernels
-                .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
-                .or_else(|| kernels.get(&(elements.encoding_id(), None)))
+            if let Some(result) = aggregate_fns
+                .find_groupped_kernel(elements.encoding_id(), self.aggregate_fn.id())
                 .and_then(|kernel| {
                     // SAFETY: we assume that elements execution is safe
                     let groups = unsafe {
@@ -254,16 +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.load();
+        let aggregate_fns = session.aggregate_fns();
 
         for _ in 0..64 {
             if elements.is::<AnyCanonical>() {
                 break;
             }
 
-            if let Some(result) = kernels
-                .get(&(elements.encoding_id(), Some(self.aggregate_fn.id())))
-                .or_else(|| kernels.get(&(elements.encoding_id(), None)))
+            if let Some(result) = aggregate_fns
+                .find_groupped_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/session.rs

@@ -8,7 +8,6 @@ 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: ArcSwap<HashMap<AggregateFnId, AggregateFnPluginRef>>,
 
-    pub(super) kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
-    pub(super) grouped_kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
+    kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
+    grouped_kernels: ArcSwap<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
 }
 
 impl SessionVar for AggregateFnSession {
@@ -106,7 +102,7 @@ impl Default for AggregateFnSession {
 }
 
 impl AggregateFnSession {
-    /// Returns the aggregate function 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()
     }
@@ -123,6 +119,20 @@ impl AggregateFnSession {
         });
     }
 
+    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.
     pub fn register_aggregate_kernel(
         &self,
@@ -137,6 +147,20 @@ impl AggregateFnSession {
             existing
         });
     }
+
+    pub fn find_groupped_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()
+    }
 }
 
 /// Extension trait for accessing aggregate function session data.