adriangb opened a new pull request, #25123: URL: https://github.com/apache/datafusion/pull/25123
## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/25116. ## Rationale for this change Some aggregate functions do not have a `GroupsAccumulator`. The built-in examples are `covar_samp`, `covar_pop`, `regr_*`, `approx_percentile_cont`, `approx_median`, `nth_value` and `any_value`. A user-defined aggregate that has only an `Accumulator` is also an example. For all of these, DataFusion uses `GroupsAccumulatorAdapter`. The adapter keeps one `Accumulator` for 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 `hits` data is an example: ```sql SELECT "UserID", covar_samp("ResolutionWidth", "ResolutionHeight") AS c FROM hits GROUP BY "UserID"; ``` 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 `UserID` values, thus approximately 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. That memory 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 `u32` value in place of the scratch `Vec<u32>`. The value is 0 between batches. For each batch, the adapter does three steps: 1. It counts the rows of each group. It also records a group the first time that it finds a row for that group. 2. It adds those counts together to get the position where the rows of each group start in the `take` index. Each count becomes that start position. 3. It writes each row index at the position for its group. The positions then go 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()` counts both of them. The rows of a group stay in the same order as in the input. Thus each `Accumulator` gets the same rows, in the same order, as before this change. This change also makes https://github.com/apache/datafusion/pull/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_extended` queries Q14 groups `covar_samp` by `"UserID"` (17.6 million groups) and Q15 groups it by `"RegionID"` (9,040 groups). The aggregate is cheap, thus the adapter is the largest part of the time. An outer `MAX` keeps the result small. ### 3. A new benchmark for the adapter `datafusion/functions-aggregate/benches/groups_accumulator_adapter.rs` moves the group 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_routing` uses an accumulator that only counts the rows that it gets. What it measures is the adapter and nothing else. - `adapter_covar_samp` uses a real aggregate that has no `GroupsAccumulator`. It shows how much of that a query gets. ## Benchmark results ### ClickBench `hits_partitioned`, 100 million rows, 12 partitions, warm page cache, release build, 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"`. `corr` on the same two columns is the control: it has a `GroupsAccumulator`, thus this PR does not change it. | Query | main | this PR | Ratio | | ----------------------------------------- | ------- | ------- | ----- | | `covar_samp` by `"UserID"` (17.6M groups) | 7278 ms | 5226 ms | 1.39x | | `corr` by `"UserID"` (control) | 663 ms | 708 ms | 0.94x | | `covar_samp` by `"RegionID"` (9k groups) | 267 ms | 275 ms | 0.97x | | `corr` by `"RegionID"` (control) | 189 ms | 183 ms | 1.04x | The two controls give a noise band of approximately 6%. The `"RegionID"` result is inside that band, thus this PR does not make the low group count slower. ### The new adapter benchmark | Groups | `adapter_routing` | `adapter_covar_samp` | | --------- | ----------------- | -------------------- | | 64 | 0.93x | 1.00x (p = 0.22) | | 1,024 | 1.06x | 1.04x (p = 0.18) | | 16,384 | 1.09x | 1.05x | | 262,144 | 1.48x | 1.24x | | 1,000,000 | 1.78x | 1.58x | At 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_samp` at 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_group` sends five batches, which include an empty 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_rows` shows that a filter keeps rows away from the accumulators. 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 any public API. `GroupsAccumulatorAdapter::size()` reports a different number. It no longer counts 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](https://claude.com/claude-code) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
