andygrove opened a new pull request, #7222: URL: https://github.com/apache/datafusion-comet/pull/7222
## Which issue does this PR close? Closes #6880. ## Rationale for this change #6867 runs a native `PartialMerge` as DataFusion `Single` so that it spills rather than emit a key twice. When the `PartialMerge`'s input is sorted on a prefix of its grouping keys, DataFusion 55.2 doesn't run `SingleHashAggregateStream` for `Single`. It falls back to the legacy `GroupedHashAggregateStream`, whose spill asks for more than its share of the pool to sort the batch it writes, and the task fails with `Failed to reserve memory for sort during spill`. This is a regression since 1.1.0, which ran this shape as `Partial` and returned a wrong `count(DISTINCT)` instead (#6857). ## What changes are included in this PR? - The aggregate planner runs a pure `PartialMerge` as DataFusion `Final` when DataFusion sees its input sorted on some of its keys (`InputOrderMode` is not `Linear`). `Final` runs `OrderedFinalAggregateStream` there, which spills and reads its spill files back. `MergeAsPartial` evaluates to its state under `Final` as it does under `Single`, so the projection that unpacks the states doesn't change. - `Final` reads each state from the columns right after the grouping keys, and it writes its spill files with the input's column types. So the planner switches only when the grouping keys are the first input columns, `initial_input_buffer_offset` equals the number of keys, and the input's state columns have the native state types. Otherwise it keeps `Single`. For example, a Spark partial's `collect_list` buffers arrive as binary. - `SpillReplayPool` already covers `OrderedFinalAggregateStream`, so reading the spill files back past the stream's share works as it does for the #6254 final aggregates. Doc comments in `merge_as_partial.rs` and `spill_replay.rs` now mention `Final`. When Comet moves to DataFusion 56 (#6410), which replaces this fallback with `OrderedSingleAggregateStream`, the switch can probably be dropped. ## How are these changes tested? - A new `CometAggregateSuite` test runs the issue's repro: `repartition(1, g)`, `sortWithinPartitions(g)`, `groupBy(g).agg(max(v), count_distinct(k))` with 2 MiB and 8 MiB pools. It checks the answer and that the `PartialMerge`, planned over a partial aggregate over a native sort, spilled. Without the planner change, both pool sizes fail with the issue's error. With a 32 MiB pool the ordered final aggregate fits in memory and doesn't spill, so the test doesn't use it. - A new planner unit test checks that a `PartialMerge` over a native sort on its first key plans as `Final` with `PartiallySorted` input and emits the merged max and count states. It also checks that a `collect_list` `PartialMerge` over binary Spark buffers stays `Single`. - `CometAggregateSuite` (163 tests) and the native core library tests pass locally. -- 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]
