DataFusion's Aggregate Split: One UDF, Two Accumulator Traits, and a Serializable State Format

TL;DR: An aggregate in DataFusion is not a function, it is a state machine: update_batch folds rows in, state() serializes the machine into Arrow arrays, merge_batch folds another machine’s state back in, and evaluate() produces the answer. That split is what makes partial/final two-phase aggregation possible. GroupsAccumulator is the vectorized sibling: one object holds the state of all groups, so there is no per-group virtual call and no per-group take. GroupsAccumulatorAdapter is the compatibility shim that lets a plain Accumulator pretend to be a GroupsAccumulator — at the price of copying rows into per-group contiguous buffers on every batch.

The problem: an aggregate is two different computations

Write SELECT avg(x) GROUP BY k over 100 million rows in 8 partitions. Two completely different jobs hide behind that one line of SQL:

Folding: for each row, update the state of the group the row belongs to.
Collapsing: for each group, turn its accumulated state into a single value.

The naive implementation does (1) with a function call per row, holding a HashMap<k, State> in memory. That is fine on a toy dataset and catastrophic on a real one, because the per-row call crosses a dynamic-dispatch boundary, the HashMap thrashes, and the state cannot be shipped anywhere.

DataFusion’s answer is to make the aggregate a value that can be serialized, shipped over a network, and merged — and then to compute it in a way that never calls a function per row.

The `Accumulator` trait: one group, one object

The whole contract is 200 lines in datafusion/expr-common/src/accumulator.rs:

rust
// datafusion/expr-common/src/accumulator.rs:123
pub trait Accumulator: Send + Sync + Debug + std::any::Any {
    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()>;   // :140
    fn evaluate(&mut self) -> Result<ScalarValue>;                    // :193
    fn size(&self) -> usize;                                          // :205
    fn state(&mut self) -> Result<Vec<ScalarValue>>;                   // :374
    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()>;      // :387
    fn retract_batch(&mut self, _values: &[ArrayRef]) -> Result<()> {  // :421
        internal_err!("Retract should be implemented ...")
    }
}

Read it as a little algebra. update_batch is the operation, state() is the encoding, merge_batch is the homomorphism over that encoding, and evaluate is the projection of the state into a final value. For SUM the operation and the merge are the same thing (state is the running sum). For COUNT they are not: the state is a count, and merging counts means adding them, not re-counting the underlying rows. For AVG the state is two values — sum and count — and the merge adds both, then evaluate divides. The trait documentation spells this out at accumulator.rs:260-266:

The state can be and often is a different type than the output type of the Accumulator and needs different merge operations (for example, the partial state for COUNT needs to be summed together)

That sentence is the whole reason the trait exists. If SUM(x) were just a function f(accumulator, row) -> (), there would be no way to run it in two phases. Because state() is a separate, serializable representation, the planner can insert a Partial aggregate before a shuffle and a Final aggregate after it, and the two halves only need to agree on the state format.

The two-phase flow

text

             ┌──────────────────────────┐
             │  Partial (per partition) │
             │  update_batch() per row  │
             │  state()  -> Arrow arrays│
             └────────────┬─────────────┘
                          │  (shuffle / repartition)
             ┌────────────▼─────────────┐
             │  Final                   │
             │  merge_batch(states)     │  <- fold partition states together
             │  evaluate() -> value     │
             └──────────────────────────┘

The state is plain Arrow data, so it crosses a RecordBatch boundary for free. AVG emits two state columns (sum, count); MEDIAN cannot know its state arity at planning time, so it packs everything into one ScalarValue::List (accumulator.rs:270-273).

The subtlety is that the final phase is not “run the function again on the state”. It is “fold state into state”. merge_batch receives the state()-encoded arrays of other accumulators, not raw rows. If you write an aggregate whose merge_batch is not the algebraic dual of update_batch, you get answers that depend on how many partitions the data had — the classic “works on my laptop” aggregate bug.

`GroupsAccumulator`: the fast path that avoids per-group calls

Now the performance trick. Accumulator handles one group. To aggregate 10 million groups you would need 10 million Box<dyn Accumulator> objects, and every row would cost you one virtual call plus a hash-map-ish lookup.

GroupsAccumulator (datafusion/expr-common/src/groups_accumulator.rs:193) flips this: one object holds the state of all groups, and the caller hands it a vector of group indices:

rust
// datafusion/expr-common/src/groups_accumulator.rs:221
fn update_batch(
    &mut self,
    values: &[ArrayRef],
    group_indices: &[usize],
    opt_filter: Option<&BooleanArray>,
    total_num_groups: usize,
) -> Result<()>;

There is no per-group call at all. group_indices[i] says “row i belongs to group group_indices[i]”, and the implementation is expected to fan that out with a tight loop over primitive values. The canonical example is PrimitiveGroupsAccumulator (datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs:43):

rust
// datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs:102
self.values.resize(total_num_groups, self.starting_value);
self.null_state.accumulate(
    group_indices,
    values,
    opt_filter,
    total_num_groups,
    |group_index, new_value| {
        // SAFETY: group_index is guaranteed to be in bounds
        let value = unsafe { self.values.get_unchecked_mut(group_index) };
        (self.prim_fn)(value, new_value);
    },
);

Three things make this fast: the state is a flat Vec<T::Native> indexed by group (cache friendly, no hashing, no allocation per group); prim_fn is a monomorphized closure, so SUM on Int64 compiles to a bare += with no vtable lookup; and the null/filter handling is hoisted out of the inner loop by accumulate (groups_accumulator/accumulate.rs:421), which has four specialized paths (no nulls/no filter, nulls/no filter, no nulls/filter, both) and even processes the null bitmap 64 bits at a time.

NullState (accumulate.rs:134) solves the awkward part: a group that has never seen a value must still produce a row (NULL for SUM, 0 for COUNT). Rather than every aggregate reinventing “did this group see anything”, the shared NullState tracks it and produces the null buffer at evaluate time.

SUM opts in with two lines (datafusion/functions-aggregate/src/sum.rs:301):

rust
fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool {
    !args.is_distinct
}

fn create_groups_accumulator(&self, args: AccumulatorArgs) -> Result<Box<dyn GroupsAccumulator>> {
    // ... PrimitiveGroupsAccumulator::<$t, _>::new(&$dt, |x, y| *x = x.add_wrapping(y))
}

The adapter: making old aggregates look new

Not every aggregate has a GroupsAccumulator, and writing one is genuinely hard (the state must be a columnar layout, EmitTo must be honoured, memory must be tracked). So DataFusion ships a bridge: GroupsAccumulatorAdapter (datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs:95).

It holds a Vec<AccumulatorState>, one Box<dyn Accumulator> per group, and bridges the vectorized API to the per-group one by reordering the rows:

text

 input rows            after take_arrays        per-group call
 ┌─────┐              ┌─────────┐
 │  2  │              │  200    │  ──┐
 │  0  │   take ────▶ │  300    │  ─┤ group 0
 │  1  │              │  200    │  ─┘
 │  0  │              │  100    │  ──┐
 └─────┘              │  NULL   │  ─┘ group 1
                     └─────────┘

The implementation (groups_accumulator.rs:281-470) builds a permutation array of u32 row indices, calls Arrow’s compute::take_arrays to produce per-group-contiguous values, and then calls the accumulator once per group through Box<dyn Accumulator>. So the adapter pays:

one take (a full copy of the batch’s values) per update batch, and
one virtual call per group per batch, plus one Accumulator allocation per group.

That is still vastly better than a virtual call per row, but it is dramatically worse than the GroupsAccumulator fast path. The selection between the two happens once, at stream construction (datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs:64):

rust
pub(in crate::aggregates) fn create_group_accumulator(
    agg_expr: &Arc<AggregateFunctionExpr>,
    metrics: Arc<dyn AggregateMetrics>,
) -> Result<Box<dyn GroupsAccumulator>> {
    if agg_expr.groups_accumulator_supported() {
        agg_expr.create_groups_accumulator_with_metrics(metrics)
    } else {
        debug!("Creating GroupsAccumulatorAdapter for {}: {agg_expr:?}", agg_expr.name());
        // ... GroupsAccumulatorAdapter::new(move || agg_expr.create_accumulator())
    }
}

The debug! line is a hint to anyone profiling: if you see it, you are on the slow path.

`EmitTo`, `take_state`, and the “emit the first N groups” dance

Aggregates cannot always emit everything at once — the hash table drains in batches to respect batch_size. EmitTo (groups_accumulator.rs:31) encodes this:

rust
pub enum EmitTo {
    All,
    /// Emit only the first `n` groups and shift all existing group indexes down by `n`.
    First(usize),
}

EmitTo::take_needed (groups_accumulator.rs:52) implements the shift with std::mem::take for All (no copy) and split_vec_min_alloc for First(n). This is why GroupsAccumulator::state takes an emit_to while Accumulator::state does not: the grouped version has to renumber the groups it did not emit, and that renumbering has to be reflected in the hash table’s own group vector. A subtle, easy-to-get-wrong contract, and one more reason the grouped API is opt-in.

first_last has its own take_state helper (datafusion/functions-aggregate/src/first_last.rs:559) that splits three parallel buffers (values, orderings, “is set” bitmap) in lockstep — the same “emit N, keep the rest” problem, but with three columns instead of one.

`convert_to_state`: skipping the hash table entirely

There is one more trick worth stealing. Normally a partial aggregate builds a hash table: rows come in, get hashed, and land in a group slot. But sometimes you already know the grouping key will not reduce anything — e.g. a partial aggregate before a GROUP BY that is only there to pre-compute a state for a downstream aggregate.

GroupsAccumulator::convert_to_state (groups_accumulator.rs:370) converts rows directly into state rows, one state row per input row, with no hash table at all:

rust
// datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs:669
fn convert_to_state(&self, values: &[ArrayRef], opt_filter: Option<&BooleanArray>)
    -> Result<Vec<ArrayRef>>
{
    // ... for each row: create a fresh accumulator, feed it that one row,
    //     then harvest accumulator.state()
}

The PartialSkip mode uses it (datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs:128, datafusion/physical-plan/src/aggregates/hash_stream.rs:516). The cost is obvious in hindsight: one accumulator allocation and one state serialization per row. The benefit: no hash table, no group vector, no group-index bookkeeping, and the output is already in the shape the next phase wants.

Trade-offs and failure modes

Two implementations to keep in sync. Every aggregate that wants the fast path must implement both Accumulator and GroupsAccumulator, and the two must agree. A GroupsAccumulator that is subtly wrong produces different answers than the Accumulator path — and the optimizer may pick either one depending on whether DISTINCT is present (see sum.rs:301).
The adapter’s copy is real. The take per batch is O(rows) memory traffic; for wide batches of expensive-to-copy types (strings, lists) this is the dominant cost of the “slow” path, not the virtual call.
evaluate must not consume state. The trait docs are emphatic (accumulator.rs:179-190): evaluate is called repeatedly for window functions with different frames, so it may reorder internal state but must not consume it. Violating this gives correct answers for GROUP BY and garbage for OVER (...).
size() is a hot path. The adapter deliberately maintains allocation_bytes incrementally rather than calling size() per group per batch, with a comment saying exactly that (“We saw size() being a bottleneck in earlier implementations when there were many distinct groups”, groups_accumulator.rs:101-106). If you implement Accumulator, make size() cheap or return a constant.
State format is a public contract. state_fields defines the Arrow schema of the intermediate state, and it is versioned with the engine. If you change a state format you break distributed plans that were planned by an older binary. approx_percentile_cont packs its state into six scalars plus a list (datafusion/functions-aggregate/src/approx_percentile_cont.rs:250) precisely because the state must round-trip through Arrow.

The takeaway: the split into update/state/merge/evaluate is not ceremony, it is what turns “aggregate function” into “algebraic data type” — something you can compute in parallel, ship across a network, and merge without ever knowing the original rows.

The problem: an aggregate is two different computations

The Accumulator trait: one group, one object