vortex_array/aggregate_fn/
vtable.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5use std::fmt::Debug;
6use std::fmt::Display;
7use std::fmt::Formatter;
8use std::hash::Hash;
9
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_session::VortexSession;
13
14use crate::ArrayRef;
15use crate::Canonical;
16use crate::DynArray;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::aggregate_fn::AggregateFn;
20use crate::aggregate_fn::AggregateFnId;
21use crate::aggregate_fn::AggregateFnRef;
22use crate::arrays::ConstantArray;
23use crate::dtype::DType;
24use crate::scalar::Scalar;
25
26/// Defines the interface for aggregate function vtables.
27///
28/// This trait is non-object-safe and allows the implementer to make use of associated types
29/// for improved type safety, while allowing Vortex to enforce runtime checks on the inputs and
30/// outputs of each function.
31///
32/// The [`AggregateFnVTable`] trait should be implemented for a struct that holds global data across
33/// all instances of the aggregate. In almost all cases, this struct will be an empty unit
34/// struct, since most aggregates do not require any global state.
35pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
36    /// Options for this aggregate function.
37    type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash;
38
39    /// The partial accumulator state for a single group.
40    type Partial: 'static + Send;
41
42    /// Returns the ID of the aggregate function vtable.
43    fn id(&self) -> AggregateFnId;
44
45    /// Serialize the options for this aggregate function.
46    ///
47    /// Should return `Ok(None)` if the function is not serializable, and `Ok(vec![])` if it is
48    /// serializable but has no metadata.
49    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
50        _ = options;
51        Ok(None)
52    }
53
54    /// Deserialize the options of this aggregate function.
55    fn deserialize(
56        &self,
57        _metadata: &[u8],
58        _session: &VortexSession,
59    ) -> VortexResult<Self::Options> {
60        vortex_bail!("Aggregate function {} is not deserializable", self.id());
61    }
62
63    /// The return [`DType`] of the aggregate.
64    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType>;
65
66    /// DType of the intermediate partial accumulator state.
67    ///
68    /// Use a struct dtype when multiple fields are needed
69    /// (e.g., Mean: `Struct { sum: f64, count: u64 }`).
70    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType>;
71
72    /// Return the partial accumulator state for an empty group.
73    fn empty_partial(
74        &self,
75        options: &Self::Options,
76        input_dtype: &DType,
77    ) -> VortexResult<Self::Partial>;
78
79    /// Combine partial scalar state into the accumulator.
80    fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()>;
81
82    /// Flush the partial aggregate for the given accumulator state.
83    ///
84    /// The returned scalar must have the same DType as specified by `state_dtype` for the
85    /// options and input dtype used to construct the state.
86    ///
87    /// The internal state of the accumulator is reset to the empty state after flushing.
88    fn flush(&self, partial: &mut Self::Partial) -> VortexResult<Scalar>;
89
90    /// Is the partial accumulator state is "saturated", i.e. has it reached a state where the
91    /// final result is fully determined.
92    fn is_saturated(&self, state: &Self::Partial) -> bool;
93
94    /// Accumulate a new canonical array into the accumulator state.
95    fn accumulate(
96        &self,
97        state: &mut Self::Partial,
98        batch: &Canonical,
99        ctx: &mut ExecutionCtx,
100    ) -> VortexResult<()>;
101
102    /// Finalize an array of accumulator states into an array of aggregate results.
103    ///
104    /// The provides `states` array has dtype as specified by `state_dtype`, the result array
105    /// must have dtype as specified by `return_dtype`.
106    fn finalize(&self, states: ArrayRef) -> VortexResult<ArrayRef>;
107
108    /// Finalize a scalar accumulator state into an aggregate result.
109    ///
110    /// The provided `state` has dtype as specified by `state_dtype`, the result scalar must have
111    /// dtype as specified by `return_dtype`.
112    fn finalize_scalar(&self, state: Scalar) -> VortexResult<Scalar> {
113        let array = ConstantArray::new(state, 1).into_array();
114        let result = self.finalize(array)?;
115        result.scalar_at(0)
116    }
117}
118
119#[derive(Clone, Debug, PartialEq, Eq, Hash)]
120pub struct EmptyOptions;
121impl Display for EmptyOptions {
122    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
123        write!(f, "")
124    }
125}
126
127/// Factory functions for aggregate vtables.
128pub trait AggregateFnVTableExt: AggregateFnVTable {
129    /// Bind this vtable with the given options into an [`AggregateFnRef`].
130    fn bind(&self, options: Self::Options) -> AggregateFnRef {
131        AggregateFn::new(self.clone(), options).erased()
132    }
133}
134impl<V: AggregateFnVTable> AggregateFnVTableExt for V {}
vortex_array/aggregate_fn/vtable.rs

vortex_array/aggregate_fn/
vtable.rs
