perf: make GroupsAccumulatorAdapter cost proportional to the batch, not the group count - #25123
Open
adriangb wants to merge 3 commits into
Open
perf: make GroupsAccumulatorAdapter cost proportional to the batch, not the group count#25123adriangb wants to merge 3 commits into
adriangb wants to merge 3 commits into
Conversation
…the batch `GroupsAccumulatorAdapter` routed each input batch through one scratch `Vec<u32>` of row indices per group. Building the `take` index meant walking every group that exists and skipping the empty ones, so each batch cost one iteration per group whatever it touched: with 1.5M groups in a partition and 8192-row batches, about 180 iterations per input row before any aggregation happened. The vectors also held their capacity between batches, so scratch space was retained per group for the lifetime of the group. Route the batch by counting sort instead. Each state carries a single `u32`: pass one counts the rows of each group and records a group the first time it is seen, pass two turns those counts into the offsets at which each group's rows start and leaves every cursor at the start of its own range, and pass three scatters the rows into the `take` index and walks the cursors back to zero for the next batch. Every pass is over the rows of the batch or over the groups the batch touches, and the only scratch that outlives a batch is `groups_with_rows` and `offsets`, both bounded by the batch size and both charged in `size()`. Rows keep their relative order within a group, so each accumulator sees exactly the rows, in the order, it saw before. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`covar_samp` has no native `GroupsAccumulator`, so it runs through `GroupsAccumulatorAdapter`. The aggregate itself is cheap, which leaves the adapter's routing as the dominant cost and makes these queries a direct measure of it. q15 groups by `"UserID"` (about 17.6M groups) and q16 by `"RegionID"` (about 9,000), so the pair covers both ends of the cardinality range that the adapter has to stay fast at. `corr` over the same two columns is the native-`GroupsAccumulator` control. The outer `MAX` keeps the stored result small. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
There was no benchmark that isolated what `GroupsAccumulatorAdapter` costs, so a change to its routing could only be judged through a whole query, where the parquet scan and the aggregate itself hide the effect. Sweep the group count from 64 to 1M with two accumulators. `routing` wraps an accumulator that only counts the rows it is handed, so what it measures is the adapter and nothing else: the upper bound on what a change to the routing can move. `covar_samp` wraps a real aggregate with no native `GroupsAccumulator`, so it shows how much of that upper bound a query sees. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
adriangb
force-pushed
the
claude/datafusion-issue-25116-75677f
branch
from
September 9, 2026 16:28
afc4aa0 to
3238cd6
Compare
Contributor
Author
|
run benchmark clickbench_extended clickbench_partitioned external_aggr tpch |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark external_aggrResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark clickbench_partitionedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark clickbench_extendedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark clickbench_partitionedCPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark clickbench_extendedCPU Details (lscpu)Details
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/datafusion-issue-25116-75677f (3238cd6) to 4048898 (merge-base) diff Run configurationrun benchmark external_aggrCPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageexternal_aggr — base (merge-base)
external_aggr — branch
File an issue against this benchmark runner |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Rationale for this change
Some aggregate functions do not have a
GroupsAccumulator. The built-inexamples are
covar_samp,covar_pop,regr_*,approx_percentile_cont,approx_median,nth_valueandany_value. A user-defined aggregate that hasonly an
Accumulatoris also an example. For all of these, DataFusion usesGroupsAccumulatorAdapter. The adapter keeps oneAccumulatorfor each group.It sends the rows of each input batch to the correct accumulators.
If you group by a column that has many different values, these queries are slow.
This query on the ClickBench
hitsdata is an example:The cause is the adapter. For each input batch, the adapter looked at every
group that exists. It did this even when the batch had rows for only a few of
those groups. There are 17.6 million different
UserIDvalues, thusapproximately 1.5 million groups in each of the 12 partitions. With 8192 rows in
a batch, the adapter did approximately 180 steps for each input row before it
started the aggregate work.
The adapter also kept a scratch
Vec<u32>of row indexes for each group. Thatmemory stayed for the full life of the group.
What changes are included in this PR?
1. The adapter sorts the rows of a batch by group with a count
Each group now keeps one
u32value in place of the scratchVec<u32>. Thevalue is 0 between batches. For each batch, the adapter does three steps:
it finds a row for that group.
group start in the
takeindex. Each count becomes that start position.back to 0 for the next batch.
Each step goes through the rows of the batch, or through the groups that the
batch has rows for. No step goes through all the groups that exist. The only
scratch memory that stays between batches is the list of groups with rows and
the list of start positions. Both are not larger than one batch.
size()countsboth of them.
The rows of a group stay in the same order as in the input. Thus each
Accumulatorgets the same rows, in the same order, as before this change.This change also makes #24858
unnecessary. There is no more scratch memory for each group to account for. The
two memory limit tests in that PR show that the adapter spills because of that
memory, thus they do not apply after this change.
2. Two new
clickbench_extendedqueriesQ15 groups
covar_sampby"UserID"(17.6 million groups) and Q16 groups it by"RegionID"(9,040 groups). The aggregate is cheap, thus the adapter is thelargest part of the time. An outer
MAXkeeps the result small.3. A new benchmark for the adapter
datafusion/functions-aggregate/benches/groups_accumulator_adapter.rsmoves thegroup count from 64 to 1,000,000. It uses two accumulators, because the cost of
the adapter and the cost of the aggregate move in opposite directions as the
group count increases:
adapter_routinguses an accumulator that only counts the rows that it gets.What it measures is the adapter and nothing else.
adapter_covar_sampuses a real aggregate that has noGroupsAccumulator. Itshows how much of that a query gets.
Benchmark results
ClickBench
hits_partitioned, 100 million rows, 12 partitions, warm page cache, releasebuild, M-series laptop with 12 cores. This machine moves by approximately 10%
between sequential runs. Thus the two binaries ran one after the other for each
single measurement, and the numbers below are medians of 12 measurements for
"RegionID"and 6 for"UserID".corron the same two columns is the control:it has a
GroupsAccumulator, thus this PR does not change it.covar_sampby"UserID"(17.6M groups)corrby"UserID"(control)covar_sampby"RegionID"(9k groups)corrby"RegionID"(control)The two controls give a noise band of approximately 6%. The
"RegionID"resultis inside that band, thus this PR does not make the low group count slower.
The new adapter benchmark
adapter_routingadapter_covar_sampAt 64 groups the adapter alone is approximately 7% slower. The old code copied
one contiguous block of row indexes for each group, and the new code writes each
row index on its own. A real accumulator hides that cost:
adapter_covar_sampat 64 and at 1,024 groups shows no difference that is statistically significant.
A second code path above a group count limit would remove those 7%. This PR does
not add one. One design is sufficient at both ends, and a second path adds a
value to tune.
What is the testing strategy for this PR?
Two new unit tests in
datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs:adapter_routes_rows_to_their_own_groupsends five batches, which include anempty batch, through the adapter. The batches interleave the groups, they go
back to groups that they already used, and they leave some groups with no
rows. The test then compares all the groups against a result that it counts
itself.
adapter_routes_filtered_rowsshows that a filter keeps rows away from theaccumulators. This includes a group where the filter removes all of the rows.
To make sure that the first test catches an error, I removed the step that sets
the positions back to 0. The test failed. I then put the step back.
The full extended test suite is green: 11,042 tests pass.
Are there any user-facing changes?
Queries that group by a column with many different values, and that use an
aggregate with no
GroupsAccumulator, are faster. There are no changes to anypublic API.
GroupsAccumulatorAdapter::size()reports a different number. It no longercounts scratch memory for each group, because there is none. It now counts the
scratch memory of one batch. A memory pool sees a smaller number for the same
query.
🤖 Generated with Claude Code