apache / apache/datafusion

Ordered ARRAY_AGG accumulator retains one Arrow array per update_batch call, inflating per-row memory for grouped aggregation

Open
#25,199 1 comment 0 reactions 1 assignee Claimed by @TinyMurky View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.