Ordered ARRAY_AGG accumulator retains one Arrow array per update_batch call, inflating per-row memory for grouped aggregation
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Is your feature request related to a problem or challenge?
Follow-up from #24392, which rewrote `OrderSensitiveArrayAggAccumulator` to retain payloads
as Arrow arrays instead of `ScalarValue`s. That is a large win when `update_batch` is called
with reasonably sized batches, but the accumulator pushes one `ArrayRef` per *call*, so the
fixed per-array cost (`ArrayData` + a 64-byte-minimum buffer + the `Arc`) is paid per call
rather than per row:
https://github.com/apache/datafusion/blob/main/datafusion/functions-aggregate/src/array_agg.rs#L1488
Ordered `ARRAY_AGG` cannot use a `GroupsAccumulator` (`groups_accumulator_supported` requires
`order_bys.is_empty()`), so grouped queries go through `GroupsAccumulatorAdapter`, which calls
`update_batch` once per group per input batch. For a high-cardinality `GROUP BY` that is
typically 1–2 rows per call, which is the worst case for this layout.
Measured with the accumulator's own `size()`, 2048 rows, `Int64` payload + `Int64` ordering key:
| rows per `update_batch` | after #24392 | before #24392 |
|---|---|---|
| 1 | 217 B/row | 152 B/row |
| 8 | 56 B/row | 152 B/row |
| 64 | 43 B/row | 152 B/row |
So `GROUP BY high_cardinality_col` with `array_agg(x ORDER BY y)` — roughly the shape of #20788 —
uses about 1.4x more memory per row than it did before, while everything else got 3–4x better.
### Describe the solution you'd like
Coalesce small inputs instead of retaining each as its own batch: below a row threshold,
`concat` into the tail batch and rewrite the affected `entries` rather than pushing a new
`ArrayRef`. That keeps the large-batch win and removes the per-call fixed cost.
Suggested fix: coalesce small inputs rather than retaining each as its own batch, e.g.
```rust
const COALESCE_ROWS: usize = 64;
// in store_batch, after compaction:
let row_count = values.len();
// ...
let start = self.entries.len();
let batch_idx = match self.batches.last() {
Some(last) if last.len() + row_count <= COALESCE_ROWS => {
let merged = arrow::compute::concat(&[last.as_ref(), values.as_ref()])?;
let idx = self.batches.len() - 1;
// rows already recorded for this batch keep their row_idx; the new rows
// start at the old length
let offset = self.batches[idx].len();
self.batches[idx] = merged;
self.entries.extend(
(0..row_count).map(|row_idx| OrderedArrayAggEntry {
batch_idx: idx,
row_idx: offset + row_idx,
}),
);
return Ok(Some(start..self.entries.len()));
}
_ => {
self.batches.push(values);
self.batches.len() - 1
}
};
```
### Describe alternatives you've considered
- A real `GroupsAccumulator` for ordered `ARRAY_AGG`, which would side-step
`GroupsAccumulatorAdapter`'s per-group slicing entirely. Bigger change, probably the
right long-term answer.
- Leaving it as is — the regression is confined to very small `update_batch` calls.
### Additional context
`benches/array_agg.rs` currently only exercises the unordered `ArrayAggAccumulator::merge_batch`,
so neither the win nor this regression is visible in CI benchmarks. A bench over the ordered path
covering both regimes (few rows per call vs. full batches, pre-ordered vs. random input) would be
worth adding alongside any fix.
Contributor guide
Assessment
This issue has not been assessed yet.