datafusion_expr_common/accumulator.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Accumulator module contains the trait definition for aggregation function's accumulators.
19
20use arrow::array::ArrayRef;
21use datafusion_common::{Result, ScalarValue, internal_err};
22use std::fmt::Debug;
23
24/// Tracks an aggregate function's state.
25///
26/// `Accumulator`s are stateful objects that implement a single group. They
27/// aggregate values from multiple rows together into a final output aggregate.
28///
29/// [`GroupsAccumulator]` is an additional more performant (but also complex) API
30/// that manages state for multiple groups at once.
31///
32/// An accumulator knows how to:
33/// * update its state from inputs via [`update_batch`]
34///
35/// * compute the final value from its internal state via [`evaluate`]
36///
37/// * retract an update to its state from given inputs via
38/// [`retract_batch`] (when used as a window aggregate [window
39/// function])
40///
41/// * convert its internal state to a vector of aggregate values via
42/// [`state`] and combine the state from multiple accumulators
43/// via [`merge_batch`], as part of efficient multi-phase grouping.
44///
45/// [`update_batch`]: Self::update_batch
46/// [`retract_batch`]: Self::retract_batch
47/// [`state`]: Self::state
48/// [`evaluate`]: Self::evaluate
49/// [`merge_batch`]: Self::merge_batch
50/// [window function]: https://en.wikipedia.org/wiki/Window_function_(SQL)
51pub trait Accumulator: Send + Sync + Debug + std::any::Any {
52 /// Updates the accumulator's state from its input.
53 ///
54 /// `values` contains the arguments to this aggregate function.
55 ///
56 /// For example, the `SUM` accumulator maintains a running sum,
57 /// and `update_batch` adds each of the input values to the
58 /// running sum.
59 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()>;
60
61 /// Returns the final aggregate value.
62 ///
63 /// For example, the `SUM` accumulator maintains a running sum,
64 /// and `evaluate` will produce that running sum as its output.
65 ///
66 /// This function gets `&mut self` to allow for the accumulator to build
67 /// arrow-compatible internal state that can be returned without copying
68 /// when possible (for example distinct strings).
69 ///
70 /// ## Correctness
71 ///
72 /// This function must not consume the internal state, as it is also used in window
73 /// aggregate functions where it can be executed multiple times depending on the
74 /// current window frame. Consuming the internal state can cause the next invocation
75 /// to have incorrect results.
76 ///
77 /// - Even if this accumulator doesn't implement [`retract_batch`] it may still be used
78 /// in window aggregate functions where the window frame is
79 /// `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`
80 ///
81 /// It is fine to modify the state (e.g. re-order elements within internal state vec) so long
82 /// as this doesn't cause an incorrect computation on the next call of evaluate.
83 ///
84 /// [`retract_batch`]: Self::retract_batch
85 fn evaluate(&mut self) -> Result<ScalarValue>;
86
87 /// Returns the allocated size required for this accumulator, in
88 /// bytes, including `Self`.
89 ///
90 /// This value is used to calculate the memory used during
91 /// execution so DataFusion can stay within its allotted limit.
92 ///
93 /// "Allocated" means that for internal containers such as `Vec`,
94 /// the `capacity` should be used not the `len`.
95 ///
96 /// May be expensive; check the implementation before calling on hot paths.
97 fn size(&self) -> usize;
98
99 /// Returns the intermediate state of the accumulator, consuming the
100 /// intermediate state.
101 ///
102 /// This function should not be called twice, otherwise it will
103 /// result in potentially non-deterministic behavior.
104 ///
105 /// This function gets `&mut self` to allow for the accumulator to build
106 /// arrow-compatible internal state that can be returned without copying
107 /// when possible (for example distinct strings).
108 ///
109 /// Intermediate state is used for "multi-phase" grouping in
110 /// DataFusion, where an aggregate is computed in parallel with
111 /// multiple `Accumulator` instances, as described below:
112 ///
113 /// # Multi-Phase Grouping
114 ///
115 /// ```text
116 /// ▲
117 /// │ evaluate() is called to
118 /// │ produce the final aggregate
119 /// │ value per group
120 /// │
121 /// ┌─────────────────────────┐
122 /// │GroupBy │
123 /// │(AggregateMode::Final) │ state() is called for each
124 /// │ │ group and the resulting
125 /// └─────────────────────────┘ RecordBatches passed to the
126 /// Final GroupBy via merge_batch()
127 /// ▲
128 /// │
129 /// ┌────────────────┴───────────────┐
130 /// │ │
131 /// │ │
132 /// ┌─────────────────────────┐ ┌─────────────────────────┐
133 /// │ GroupBy │ │ GroupBy │
134 /// │(AggregateMode::Partial) │ │(AggregateMode::Partial) │
135 /// └─────────────────────────┘ └─────────────────────────┘
136 /// ▲ ▲
137 /// │ │ update_batch() is called for
138 /// │ │ each input RecordBatch
139 /// .─────────. .─────────.
140 /// ,─' '─. ,─' '─.
141 /// ; Input : ; Input :
142 /// : Partition 0 ; : Partition 1 ;
143 /// ╲ ╱ ╲ ╱
144 /// '─. ,─' '─. ,─'
145 /// `───────' `───────'
146 /// ```
147 ///
148 /// The partial state is serialized as `Arrays` and then combined
149 /// with other partial states from different instances of this
150 /// Accumulator (that ran on different partitions, for example).
151 ///
152 /// The state can be and often is a different type than the output
153 /// type of the [`Accumulator`] and needs different merge
154 /// operations (for example, the partial state for `COUNT` needs
155 /// to be summed together)
156 ///
157 /// Some accumulators can return multiple values for their
158 /// intermediate states. For example, the average accumulator
159 /// tracks `sum` and `n`, and this function should return a vector
160 /// of two values, sum and n.
161 ///
162 /// Note that [`ScalarValue::List`] can be used to pass multiple
163 /// values if the number of intermediate values is not known at
164 /// planning time (e.g. for `MEDIAN`)
165 ///
166 /// # Multi-phase repartitioned Grouping
167 ///
168 /// Many multi-phase grouping plans contain a Repartition operation
169 /// as well as shown below:
170 ///
171 /// ```text
172 /// ▲ ▲
173 /// │ │
174 /// │ │
175 /// │ │
176 /// │ │
177 /// │ │
178 /// ┌───────────────────────┐ ┌───────────────────────┐ 4. Each AggregateMode::Final
179 /// │GroupBy │ │GroupBy │ GroupBy has an entry for its
180 /// │(AggregateMode::Final) │ │(AggregateMode::Final) │ subset of groups (in this case
181 /// │ │ │ │ that means half the entries)
182 /// └───────────────────────┘ └───────────────────────┘
183 /// ▲ ▲
184 /// │ │
185 /// └─────────────┬────────────┘
186 /// │
187 /// │
188 /// │
189 /// ┌─────────────────────────┐ 3. Repartitioning by hash(group
190 /// │ Repartition │ keys) ensures that each distinct
191 /// │ HASH(x) │ group key now appears in exactly
192 /// └─────────────────────────┘ one partition
193 /// ▲
194 /// │
195 /// ┌───────────────┴─────────────┐
196 /// │ │
197 /// │ │
198 /// ┌─────────────────────────┐ ┌──────────────────────────┐ 2. Each AggregateMode::Partial
199 /// │ GroupBy │ │ GroupBy │ GroupBy has an entry for *all*
200 /// │(AggregateMode::Partial) │ │ (AggregateMode::Partial) │ the groups
201 /// └─────────────────────────┘ └──────────────────────────┘
202 /// ▲ ▲
203 /// │ │
204 /// │ │
205 /// .─────────. .─────────.
206 /// ,─' '─. ,─' '─.
207 /// ; Input : ; Input : 1. Since input data is
208 /// : Partition 0 ; : Partition 1 ; arbitrarily or RoundRobin
209 /// ╲ ╱ ╲ ╱ distributed, each partition
210 /// '─. ,─' '─. ,─' likely has all distinct
211 /// `───────' `───────'
212 /// ```
213 ///
214 /// This structure is used so that the `AggregateMode::Partial` accumulators
215 /// reduces the cardinality of the input as soon as possible. Typically,
216 /// each partial accumulator sees all groups in the input as the group keys
217 /// are evenly distributed across the input.
218 ///
219 /// The final output is computed by repartitioning the result of
220 /// [`Self::state`] from each Partial aggregate and `hash(group keys)` so
221 /// that each distinct group key appears in exactly one of the
222 /// `AggregateMode::Final` GroupBy nodes. The outputs of the final nodes are
223 /// then unioned together to produce the overall final output.
224 ///
225 /// Here is an example that shows the distribution of groups in the
226 /// different phases
227 ///
228 /// ```text
229 /// ┌─────┐ ┌─────┐
230 /// │ 1 │ │ 3 │
231 /// ├─────┤ ├─────┤
232 /// │ 2 │ │ 4 │ After repartitioning by
233 /// └─────┘ └─────┘ hash(group keys), each distinct
234 /// ┌─────┐ ┌─────┐ group key now appears in exactly
235 /// │ 1 │ │ 3 │ one partition
236 /// ├─────┤ ├─────┤
237 /// │ 2 │ │ 4 │
238 /// └─────┘ └─────┘
239 ///
240 ///
241 /// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─
242 ///
243 /// ┌─────┐ ┌─────┐
244 /// │ 2 │ │ 2 │
245 /// ├─────┤ ├─────┤
246 /// │ 1 │ │ 2 │
247 /// ├─────┤ ├─────┤
248 /// │ 3 │ │ 3 │
249 /// ├─────┤ ├─────┤
250 /// │ 4 │ │ 1 │
251 /// └─────┘ └─────┘ Input data is arbitrarily or
252 /// ... ... RoundRobin distributed, each
253 /// ┌─────┐ ┌─────┐ partition likely has all
254 /// │ 1 │ │ 4 │ distinct group keys
255 /// ├─────┤ ├─────┤
256 /// │ 4 │ │ 3 │
257 /// ├─────┤ ├─────┤
258 /// │ 1 │ │ 1 │
259 /// ├─────┤ ├─────┤
260 /// │ 4 │ │ 3 │
261 /// └─────┘ └─────┘
262 ///
263 /// group values group values
264 /// in partition 0 in partition 1
265 /// ```
266 fn state(&mut self) -> Result<Vec<ScalarValue>>;
267
268 /// Updates the accumulator's state from an `Array` containing one
269 /// or more intermediate values.
270 ///
271 /// For some aggregates (such as `SUM`), merge_batch is the same
272 /// as `update_batch`, but for some aggregates (such as `COUNT`)
273 /// the operations differ. See [`Self::state`] for more details on how
274 /// state is used and merged.
275 ///
276 /// The `states` array passed was formed by concatenating the
277 /// results of calling [`Self::state`] on zero or more other
278 /// `Accumulator` instances.
279 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()>;
280
281 /// Retracts (removed) an update (caused by the given inputs) to
282 /// accumulator's state.
283 ///
284 /// This is the inverse operation of [`Self::update_batch`] and is used
285 /// to incrementally calculate window aggregates where the `OVER`
286 /// clause defines a bounded window.
287 ///
288 /// # Example
289 ///
290 /// For example, given the following input partition
291 ///
292 /// ```text
293 /// │ current │
294 /// window
295 /// │ │
296 /// ┌────┬────┬────┬────┬────┬────┬────┬────┬────┐
297 /// Input │ A │ B │ C │ D │ E │ F │ G │ H │ I │
298 /// partition └────┴────┴────┴────┼────┴────┴────┴────┼────┘
299 ///
300 /// │ next │
301 /// window
302 /// ```
303 ///
304 /// First, [`Self::evaluate`] will be called to produce the output
305 /// for the current window.
306 ///
307 /// Then, to advance to the next window:
308 ///
309 /// First, [`Self::retract_batch`] will be called with the values
310 /// that are leaving the window, `[B, C, D]` and then
311 /// [`Self::update_batch`] will be called with the values that are
312 /// entering the window, `[F, G, H]`.
313 fn retract_batch(&mut self, _values: &[ArrayRef]) -> Result<()> {
314 // TODO add retract for all accumulators
315 internal_err!(
316 "Retract should be implemented for aggregate functions when used with custom window frame queries"
317 )
318 }
319
320 /// Does the accumulator support incrementally updating its value
321 /// by *removing* values.
322 ///
323 /// If this function returns true, [`Self::retract_batch`] will be
324 /// called for sliding window functions such as queries with an
325 /// `OVER (ROWS BETWEEN 1 PRECEDING AND 2 FOLLOWING)`
326 fn supports_retract_batch(&self) -> bool {
327 false
328 }
329}