From 2a2cde32c512c88940878486a580a374211bea82 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 16 Oct 2026 13:53:50 +0000 Subject: [PATCH] perf(array): snapshot aggregate kernels in ExecutionCtx Aggregate accumulators looked up `AggregateFnSession` on every batch and on every iteration of their execute loop, paying a session arc-swap load, two type-map lookups and a registry load each time. Snapshot the aggregate, grouped and grouped-encoding kernel registries into the `ExecutionCtx` on first use, mirroring the execute-parent kernel snapshot, and dispatch through it. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GjnvCm6kT6P79iEBaHshSN --- vortex-array/src/aggregate_fn/accumulator.rs | 51 +++++++++++-- .../src/aggregate_fn/accumulator_grouped.rs | 11 ++- vortex-array/src/aggregate_fn/session.rs | 72 +++++++++++++++++-- vortex-array/src/executor.rs | 14 ++++ 4 files changed, 128 insertions(+), 20 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index 824ea3990ac..2ab8ac676aa 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/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 Accumulator { /// 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` 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 Accumulator { 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_::(), + Some(7.0) + ); + + session + .get::() + .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_::(), + 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] diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs index af5dcaed37b..7c5cc224363 100644 --- a/vortex-array/src/aggregate_fn/accumulator_grouped.rs +++ b/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 GroupedAccumulator { 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 GroupedAccumulator { } // 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 diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 65b00d81031..669fbb2bb82 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/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>, + grouped_kernels: Arc>, + grouped_encoding_kernels: + Arc>, +} + +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, + 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. diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 71245351deb..afc5464b03e 100644 --- a/vortex-array/src/executor.rs +++ b/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, execute_parent_kernels: Arc, + // OnceLock avoids snapshotting the aggregate kernels when a context does not aggregate. + aggregate_kernels: OnceLock, #[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);