apache / apache/datafusion

GroupsAccumulatorAdapter: per-batch cost scales with total group count, not batch rows

Open
#25,116 0 comments 0 reactions 0 assignees 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?

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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.