Skip to main content

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::arrays::FixedSizeListArray;
15use crate::arrays::ListViewArray;
16use crate::scalar::Scalar;
17
18/// A pluggable kernel for an aggregate function.
19///
20/// The provided array should be aggregated into a single scalar representing the state of a single
21/// group.
22pub trait DynAggregateKernel: 'static + Send + Sync + Debug {
23    fn aggregate(
24        &self,
25        aggregate_fn: &AggregateFnRef,
26        batch: &ArrayRef,
27        ctx: &mut ExecutionCtx,
28    ) -> VortexResult<Option<Scalar>>;
29}
30
31/// A pluggable kernel for batch aggregation of many groups.
32///
33/// The kernel is matched on the encoding of the _elements_ array, which is the inner array of the
34/// provided `ListViewArray`. This is more pragmatic than having every kernel match on the outer
35/// list encoding and having to deal with the possibility of multiple list encodings.
36///
37/// Each element of the list array represents a group and the result of the grouped aggregate
38/// should be an array of the same length, where each element is the aggregate state of the
39/// corresponding group.
40///
41/// Return `Ok(None)` if the kernel cannot be applied to the given aggregate function.
42pub trait DynGroupedAggregateKernel: 'static + Send + Sync + Debug {
43    /// Aggregate each group in the provided `ListViewArray` and return an array of the
44    /// aggregate states.
45    fn grouped_aggregate(
46        &self,
47        aggregate_fn: &AggregateFnRef,
48        groups: &ListViewArray,
49    ) -> VortexResult<Option<ArrayRef>>;
50
51    /// Aggregate each group in the provided `FixedSizeListArray` and return an array of the
52    /// aggregate states.
53    fn grouped_aggregate_fixed_size(
54        &self,
55        aggregate_fn: &AggregateFnRef,
56        groups: &FixedSizeListArray,
57    ) -> VortexResult<Option<ArrayRef>> {
58        // TODO(ngates): we could automatically delegate to `grouped_aggregate` if SequenceArray
59        //  was in the vortex-array crate
60        let _ = (aggregate_fn, groups);
61        Ok(None)
62    }
63}