GroupsAccumulatorAdapter: per-batch cost scales with total group count, not batch rows
- 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?
Aggregates that have no native `GroupsAccumulator` get slow when a `GROUP BY` has many groups. Built-in examples are `covar_samp`, `covar_pop`, `regr_*`, `approx_percentile_cont`, `approx_median`, `nth_value` and `any_value`, plus any user-defined aggregate that only implements `Accumulator`.
#### Example query
ClickBench `hits`, grouped by a key with many distinct values:
```sql
SELECT "UserID", covar_samp("ResolutionWidth", "ResolutionHeight") AS c
FROM hits
GROUP BY "UserID";
```
#### How DataFusion plans it
```
ProjectionExec: expr=[UserID@0 as UserID, covar_samp(...)@1 as c]
AggregateExec: mode=FinalPartitioned, gby=[UserID@0 as UserID], aggr=[covar_samp(hits.ResolutionWidth,hits.ResolutionHeight)]
RepartitionExec: partitioning=Hash([UserID@0], 12), input_partitions=12
AggregateExec: mode=Partial, gby=[UserID@0 as UserID], aggr=[covar_samp(hits.ResolutionWidth,hits.ResolutionHeight)]
DataSourceExec: file_groups={12 groups: [...hits_partitioned...]}, projection=[UserID, ResolutionWidth, ResolutionHeight]
```
`covar_samp` has no `GroupsAccumulator`, so both `AggregateExec`s wrap its `Accumulator` in `GroupsAccumulatorAdapter`. The adapter keeps one `Accumulator` per group and routes each input batch to them.
#### Current performance
Same query with three keys of increasing cardinality. `corr` over the same two columns is the control: it has a native `GroupsAccumulator`, so it goes through the fast path.
| `GROUP BY` | groups | `covar_samp` (adapter) | `corr` (native) | slowdown |
| --- | --- | --- | --- | --- |
| `"RegionID"` | 9,040 | 0.25 s | 0.18 s | 1.4x |
| `"ClientIP"` | 9.8M | 5.9 s | 0.92 s | 6.4x |
| `"UserID"` | 17.6M | 9.4 s | 1.1 s | 8.6x |
`hits_partitioned`, 100M rows, `datafusion-cli` 55.0.0 release build, 12 partitions, warm page cache, single runs on an M-series laptop with 12 cores. The gap is small with few groups and grows with the number of groups.
#### Root cause
`GroupsAccumulatorAdapter` (`datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs`) routes each batch through one scratch `Vec` of row indices per group:
1. Push each row's index into its group's vector.
2. **Walk every group that exists**, skip the empty ones, and concatenate the rest into a `take` index.
3. `take` the values once, then call each touched group's `Accumulator` with its slice.
4. Clear the vectors. They keep their capacity.
Step 2 costs one iteration per existing group on every batch, whatever the batch touches. With 1.5M groups per partition (`UserID` / 12) and 8192-row batches, that is about 180 iterations per input row before any aggregation happens. Step 3 adds one dynamic dispatch and one array slice per touched group, which at high cardinality is one per row.
The per-group vectors also hold memory between batches. That is a smaller effect, and https://github.com/apache/datafusion/pull/24858 fixes the accounting for it.
### Describe the solution you'd like
Make the adapter's per-batch cost proportional to the batch, not to the number of groups. Two steps, either of which stands on its own:
1. **Track touched groups.** Record the groups seen in step 1 and walk only those in step 2. Removes the scan over idle groups, keeps the memory layout.
2. **Drop the per-group vectors.** Sort the batch's `(group, row)` pairs, or bucket them through a map of touched groups, and build the `take` index and offsets from that. Scratch memory becomes the size of one batch and nothing is retained per group, which makes the accounting in https://github.com/apache/datafusion/pull/24858 unnecessary.
Switching to a different implementation above a cardinality threshold would also work, but it adds a second code path and a tuning knob for a problem one design can solve at both ends.
#### Benchmark
`clickbench_extended` q05 already runs `APPROX_PERCENTILE_CONT` through the adapter grouped by `"ClientIP", "WatchID"`, with q04 (`MEDIAN`, native) as its control. The query above is a cleaner target because `covar_samp` is cheap, so the adapter dominates. I suggest adding it to `clickbench_extended` with `"UserID"` and a `"RegionID"` twin, wrapped in an outer `MAX`/`COUNT` so the stored result stays small. Success is the `"UserID"` query approaching `corr` with no regression on `"RegionID"`.
Contributor guide
Research direction
Start by reading datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs and trace how GroupsAccumulatorAdapter collects and processes group rows for each batch. Run the existing clickbench_extended q05 and q04 benchmarks, then add the proposed UserID and RegionID comparison; done means per-batch work no longer scans all groups, the UserID case approaches corr, and RegionID does not regress.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust, sql
- Domain
- databases, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 58/100