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

[adriangb]

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:

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 AggregateExecs 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<u32> 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 #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 fix: charge retained scratch indices capacity in GroupsAccumulatorAdapter #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".