vortex_array/aggregate_fn/kernels.rs
1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Pluggable aggregate function kernels used to provide encoding-specific implementations of
5//! aggregate functions.
6
7use std::fmt::Debug;
8
9use vortex_error::VortexResult;
10
11use crate::ArrayRef;
12use crate::ExecutionCtx;
13use crate::aggregate_fn::AggregateFnRef;
14use crate::aggregate_fn::GroupedArray;
15use crate::scalar::Scalar;
16
17/// A pluggable kernel for an aggregate function.
18///
19/// The provided array should be aggregated into a single scalar representing the partial state
20/// of a single group.
21pub trait DynAggregateKernel: 'static + Send + Sync + Debug {
22 fn aggregate(
23 &self,
24 aggregate_fn: &AggregateFnRef,
25 batch: &ArrayRef,
26 ctx: &mut ExecutionCtx,
27 ) -> VortexResult<Option<Scalar>>;
28}
29
30/// A pluggable kernel for batch aggregation of many groups.
31///
32/// A kernel can be registered either for an aggregate function regardless of the element encoding,
33/// or for a specific aggregate function and element encoding. Element-encoding kernels are matched
34/// on the inner array of the provided grouped array, not on the outer list encoding. This is more
35/// pragmatic than having every kernel match on the outer list encoding and having to deal with the
36/// possibility of multiple list encodings.
37///
38/// Each value in the grouped array represents a group and the result of the grouped aggregate
39/// should be an array of the same length, where each element is the aggregate state of the
40/// corresponding group.
41///
42/// Return `Ok(None)` if the kernel cannot be applied to the given aggregate function.
43pub trait DynGroupedAggregateKernel: 'static + Send + Sync + Debug {
44 /// Aggregate each group in the provided grouped array and return an array of the aggregate
45 /// states.
46 fn grouped_aggregate(
47 &self,
48 aggregate_fn: &AggregateFnRef,
49 groups: &GroupedArray,
50 ctx: &mut ExecutionCtx,
51 ) -> VortexResult<Option<ArrayRef>>;
52}