perf(array): snapshot aggregate kernels in ExecutionCtx

vortex-array/src/aggregate_fn/accumulator.rs

@@ -16,7 +16,6 @@ use crate::aggregate_fn::AggregateDTypes;
 use crate::aggregate_fn::AggregateFn;
 use crate::aggregate_fn::AggregateFnRef;
 use crate::aggregate_fn::AggregateFnVTable;
-use crate::aggregate_fn::session::AggregateFnSessionExt;
 use crate::columnar::AnyColumnar;
 use crate::dtype::DType;
 use crate::executor::max_iterations;
@@ -99,16 +98,14 @@ impl<V: AggregateFnVTable> Accumulator<V> {
     /// Accumulate `batch` into the partial state through kernel dispatch, the vtable's
     /// short-circuit, or execution to a columnar array, without consulting cached statistics.
     fn accumulate_batch(&mut self, batch: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
-        let session = ctx.session().clone();
-
         // 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 kernel = session
-                .aggregate_fns()
+            let kernel = ctx
+                .aggregate_kernels()
                 .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)?
@@ -144,8 +141,8 @@ impl<V: AggregateFnVTable> Accumulator<V> {
                 break;
             }
 
-            if let Some(kernel) = session
-                .aggregate_fns()
+            if let Some(kernel) = ctx
+                .aggregate_kernels()
                 .find_aggregate_kernel(batch.encoding_id(), self.aggregate_fn.id())
                 && let Some(result) = kernel.aggregate(&self.aggregate_fn, &batch, ctx)?
             {
@@ -507,6 +504,46 @@ mod tests {
         Ok(())
     }
 
+    /// An execution context snapshots the aggregate kernels on first use, so a kernel registered
+    /// afterwards only fires through a new context.
+    #[test]
+    fn execution_ctx_snapshots_aggregate_kernels_on_first_use() -> VortexResult<()> {
+        static KERNEL: SentinelMeanPartialKernel = SentinelMeanPartialKernel;
+        let session = fresh_session();
+        let mut ctx = session.create_execution_ctx();
+
+        let mut acc = mean_f64_accumulator()?;
+        acc.accumulate(&dict_of_seven(), &mut ctx)?;
+        assert_eq!(
+            acc.flush()?
+                .as_struct()
+                .field("sum")
+                .unwrap()
+                .as_primitive()
+                .as_::<f64>(),
+            Some(7.0)
+        );
+
+        session
+            .get::<AggregateFnSession>()
+            .register_aggregate_kernel(Dict.id(), Some(Mean::combined().id()), &KERNEL);
+
+        acc.accumulate(&dict_of_seven(), &mut ctx)?;
+        assert_eq!(
+            acc.flush()?
+                .as_struct()
+                .field("sum")
+                .unwrap()
+                .as_primitive()
+                .as_::<f64>(),
+            Some(7.0)
+        );
+
+        acc.accumulate(&dict_of_seven(), &mut session.create_execution_ctx())?;
+        assert_eq!(acc.flush()?, sentinel_partial());
+        Ok(())
+    }
+
     /// Kernel returns `Ok(None)` => dispatch falls through to `Combined::try_accumulate`'s
     /// natural fan-out. The natural partial is `{sum: 7.0, count: 1}`.
     #[test]

vortex-array/src/aggregate_fn/accumulator_grouped.rs

@@ -21,7 +21,6 @@ use crate::aggregate_fn::AggregateFn;
 use crate::aggregate_fn::AggregateFnRef;
 use crate::aggregate_fn::AggregateFnVTable;
 use crate::aggregate_fn::DynAccumulator;
-use crate::aggregate_fn::session::AggregateFnSessionExt;
 use crate::arrays::ChunkedArray;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::ListViewArray;
@@ -289,12 +288,10 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<()> {
         let mut elements = groups.elements().clone();
-        let session = ctx.session().clone();
-
         for _ in 0..max_iterations() {
             // Try a registered grouped kernel for the current element encoding.
-            if let Some(kernel) = session
-                .aggregate_fns()
+            if let Some(kernel) = ctx
+                .aggregate_kernels()
                 .find_grouped_encoding_kernel(elements.encoding_id(), self.aggregate_fn.id())
             {
                 // SAFETY: we assume that elements execution is safe
@@ -307,8 +304,8 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
             }
 
             // Try a grouped kernel for the current aggregate regardless of element encoding.
-            if let Some(kernel) = session
-                .aggregate_fns()
+            if let Some(kernel) = ctx
+                .aggregate_kernels()
                 .find_grouped_kernel(self.aggregate_fn.id())
             {
                 // SAFETY: we preserve the grouped shape and validity while replacing the

vortex-array/src/aggregate_fn/session.rs

@@ -8,6 +8,7 @@ use vortex_session::ArcSwapMap;
 use vortex_session::SessionExt;
 use vortex_session::SessionGuard;
 use vortex_session::SessionVar;
+use vortex_utils::aliases::hash_map::HashMap;
 
 use crate::aggregate_fn::AggregateFnId;
 use crate::aggregate_fn::AggregateFnPluginRef;
@@ -180,12 +181,8 @@ impl AggregateFnSession {
     ) -> 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()
-        })
+        self.kernels
+            .read(|kernels| lookup_aggregate_kernel(kernels, id, fn_id))
     }
 
     /// Registers an aggregate kernel for an array encoding.
@@ -252,6 +249,69 @@ impl AggregateFnSession {
         let fn_id = agg_fn_id.into();
         self.grouped_encoding_kernels.insert((id, fn_id), kernel)
     }
+
+    /// Captures the currently registered aggregate kernels, so repeated lookups skip the session
+    /// and registry loads.
+    pub(crate) fn kernels_snapshot(&self) -> AggregateKernels {
+        AggregateKernels {
+            kernels: self.kernels.snapshot(),
+            grouped_kernels: self.grouped_kernels.snapshot(),
+            grouped_encoding_kernels: self.grouped_encoding_kernels.snapshot(),
+        }
+    }
+}
+
+/// A point-in-time snapshot of the aggregate kernel registries of an [`AggregateFnSession`].
+///
+/// Lookups follow the same rules as the corresponding `find_*` methods on
+/// [`AggregateFnSession`].
+#[derive(Clone, Debug)]
+pub(crate) struct AggregateKernels {
+    kernels: Arc<HashMap<AggregateKernelKey, &'static dyn DynAggregateKernel>>,
+    grouped_kernels: Arc<HashMap<AggregateFnId, &'static dyn DynGroupedAggregateKernel>>,
+    grouped_encoding_kernels:
+        Arc<HashMap<GroupedEncodingKernelKey, &'static dyn DynGroupedAggregateKernel>>,
+}
+
+impl AggregateKernels {
+    /// See [`AggregateFnSession::find_aggregate_kernel`].
+    pub(crate) fn find_aggregate_kernel(
+        &self,
+        array_id: ArrayId,
+        agg_fn_id: AggregateFnId,
+    ) -> Option<&'static dyn DynAggregateKernel> {
+        lookup_aggregate_kernel(&self.kernels, array_id, agg_fn_id)
+    }
+
+    /// See [`AggregateFnSession::find_grouped_kernel`].
+    pub(crate) fn find_grouped_kernel(
+        &self,
+        agg_fn_id: AggregateFnId,
+    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
+        self.grouped_kernels.get(&agg_fn_id).copied()
+    }
+
+    /// See [`AggregateFnSession::find_grouped_encoding_kernel`].
+    pub(crate) fn find_grouped_encoding_kernel(
+        &self,
+        array_id: ArrayId,
+        agg_fn_id: AggregateFnId,
+    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
+        self.grouped_encoding_kernels
+            .get(&(array_id, agg_fn_id))
+            .copied()
+    }
+}
+
+fn lookup_aggregate_kernel(
+    kernels: &HashMap<AggregateKernelKey, &'static dyn DynAggregateKernel>,
+    array_id: ArrayId,
+    agg_fn_id: AggregateFnId,
+) -> Option<&'static dyn DynAggregateKernel> {
+    kernels
+        .get(&(array_id, Some(agg_fn_id)))
+        .or_else(|| kernels.get(&(array_id, None)))
+        .copied()
 }
 
 /// Extension trait for accessing aggregate function session data.

vortex-array/src/executor.rs

@@ -34,6 +34,8 @@ use vortex_session::VortexSession;
 use crate::AnyCanonical;
 use crate::ArrayRef;
 use crate::IntoArray;
+use crate::aggregate_fn::session::AggregateFnSessionExt;
+use crate::aggregate_fn::session::AggregateKernels;
 use crate::array::ArrayId;
 use crate::builders::ArrayBuilder;
 use crate::builders::builder_with_capacity_in;
@@ -357,6 +359,8 @@ pub struct ExecutionCtx {
     // OnceLock avoids cloning the session allocator when a context does not allocate.
     allocator: OnceLock<BufferAllocatorRef>,
     execute_parent_kernels: Arc<ParentExecutionKernels>,
+    // OnceLock avoids snapshotting the aggregate kernels when a context does not aggregate.
+    aggregate_kernels: OnceLock<AggregateKernels>,
     #[cfg(debug_assertions)]
     id: usize,
     #[cfg(debug_assertions)]
@@ -375,6 +379,7 @@ impl ExecutionCtx {
             session,
             allocator: OnceLock::new(),
             execute_parent_kernels,
+            aggregate_kernels: OnceLock::new(),
             #[cfg(debug_assertions)]
             id: {
                 static EXEC_CTX_ID: AtomicUsize = AtomicUsize::new(0);
@@ -395,6 +400,15 @@ impl ExecutionCtx {
         self.allocator.get_or_init(|| self.session.allocator())
     }
 
+    /// Get the aggregate kernels for this execution context.
+    ///
+    /// The session's aggregate kernel registries are snapshotted on first use. Kernels registered
+    /// after that are not visible to this context.
+    pub(crate) fn aggregate_kernels(&self) -> &AggregateKernels {
+        self.aggregate_kernels
+            .get_or_init(|| self.session.aggregate_fns().kernels_snapshot())
+    }
+
     /// Set the allocator for this execution context.
     pub fn with_allocator(mut self, allocator: BufferAllocatorRef) -> Self {
         self.allocator = OnceLock::from(allocator);