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::Columnar;
16use crate::ExecutionCtx;
17use crate::aggregate_fn::AggregateFn;
18use crate::aggregate_fn::AggregateFnId;
19use crate::aggregate_fn::AggregateFnRef;
20use crate::dtype::DType;
21use crate::scalar::Scalar;
22
23/// Defines the interface for aggregate function vtables.
24///
25/// This trait is non-object-safe and allows the implementer to make use of associated types
26/// for improved type safety, while allowing Vortex to enforce runtime checks on the inputs and
27/// outputs of each function.
28///
29/// The [`AggregateFnVTable`] trait should be implemented for a struct that holds global data across
30/// all instances of the aggregate. In almost all cases, this struct will be an empty unit
31/// struct, since most aggregates do not require any global state.
32pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
33 /// Options for this aggregate function.
34 type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash;
35
36 /// The partial accumulator state for a single group.
37 type Partial: 'static + Send;
38
39 /// Returns the ID of the aggregate function vtable.
40 fn id(&self) -> AggregateFnId;
41
42 /// Serialize the options for this aggregate function.
43 ///
44 /// Should return `Ok(None)` if the function is not serializable, and `Ok(vec![])` if it is
45 /// serializable but has no metadata.
46 fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
47 _ = options;
48 Ok(None)
49 }
50
51 /// Deserialize the options of this aggregate function.
52 fn deserialize(
53 &self,
54 _metadata: &[u8],
55 _session: &VortexSession,
56 ) -> VortexResult<Self::Options> {
57 vortex_bail!("Aggregate function {} is not deserializable", self.id());
58 }
59
60 /// Coerce the input type for this aggregate function.
61 ///
62 /// This is optionally used by Vortex users when performing type coercion over a Vortex
63 /// expression. The default implementation returns the input type unchanged.
64 fn coerce_args(&self, options: &Self::Options, input_dtype: &DType) -> VortexResult<DType> {
65 let _ = options;
66 Ok(input_dtype.clone())
67 }
68
69 /// The return [`DType`] of the aggregate.
70 ///
71 /// Returns `None` if the aggregate function cannot be applied to the input dtype.
72 fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType>;
73
74 /// DType of the intermediate partial accumulator state.
75 ///
76 /// Use a struct dtype when multiple fields are needed
77 /// (e.g., Mean: `Struct { sum: f64, count: u64 }`).
78 ///
79 /// Returns `None` if the aggregate function cannot be applied to the input dtype.
80 fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType>;
81
82 /// Return the partial accumulator state for an empty group.
83 fn empty_partial(
84 &self,
85 options: &Self::Options,
86 input_dtype: &DType,
87 ) -> VortexResult<Self::Partial>;
88
89 /// Combine partial scalar state into the accumulator.
90 fn combine_partials(&self, partial: &mut Self::Partial, other: Scalar) -> VortexResult<()>;
91
92 /// Convert the partial state into a partial scalar.
93 ///
94 /// The returned scalar must have the same DType as specified by `partial_dtype` for the
95 /// options and input dtype used to construct the state.
96 fn to_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar>;
97
98 /// Reset the state of the accumulator to an empty group.
99 fn reset(&self, partial: &mut Self::Partial);
100
101 /// Is the partial accumulator state is "saturated", i.e. has it reached a state where the
102 /// final result is fully determined.
103 fn is_saturated(&self, state: &Self::Partial) -> bool;
104
105 /// Try to accumulate the raw array before decompression.
106 ///
107 /// Returns `true` if the array was handled, `false` to fall through to
108 /// the default kernel dispatch and canonicalization path.
109 ///
110 /// This is useful for aggregates that only depend on array metadata (e.g., validity)
111 /// rather than the encoded data, avoiding unnecessary decompression.
112 fn try_accumulate(
113 &self,
114 _state: &mut Self::Partial,
115 _batch: &ArrayRef,
116 _ctx: &mut ExecutionCtx,
117 ) -> VortexResult<bool> {
118 Ok(false)
119 }
120
121 /// Accumulate a new canonical array into the accumulator state.
122 fn accumulate(
123 &self,
124 state: &mut Self::Partial,
125 batch: &Columnar,
126 ctx: &mut ExecutionCtx,
127 ) -> VortexResult<()>;
128
129 /// Finalize an array of accumulator states into an array of aggregate results.
130 ///
131 /// The provides `states` array has dtype as specified by `state_dtype`, the result array
132 /// must have dtype as specified by `return_dtype`.
133 fn finalize(&self, states: ArrayRef) -> VortexResult<ArrayRef>;
134
135 /// Finalize a scalar accumulator state into an aggregate result.
136 ///
137 /// The provided `state` has dtype as specified by `state_dtype`, the result scalar must have
138 /// dtype as specified by `return_dtype`.
139 fn finalize_scalar(&self, partial: &Self::Partial) -> VortexResult<Scalar>;
140}
141
142#[derive(Clone, Debug, PartialEq, Eq, Hash)]
143pub struct EmptyOptions;
144impl Display for EmptyOptions {
145 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
146 write!(f, "")
147 }
148}
149
150/// Factory functions for aggregate vtables.
151pub trait AggregateFnVTableExt: AggregateFnVTable {
152 /// Bind this vtable with the given options into an [`AggregateFnRef`].
153 fn bind(&self, options: Self::Options) -> AggregateFnRef {
154 AggregateFn::new(self.clone(), options).erased()
155 }
156}
157impl<V: AggregateFnVTable> AggregateFnVTableExt for V {}