comphead opened a new issue, #7146:
URL: https://github.com/apache/datafusion-comet/issues/7146
### What is the problem the feature request solves?
Spark plans `SortAggregateExec` when an aggregate buffer cannot be updated
in place in an `UnsafeRow`, for example `min`/`max` over strings
(`UnsafeRow.isMutable` excludes strings, checked in Spark 3.4.3 and 4.1.3).
`SortAggregateExec.requiredChildOrdering` is the grouping keys, so
`EnsureRequirements` adds a `SortExec` below every sort aggregate, partial and
final.
Comet converts both. The aggregate serializes to the same native
`HashAggregate` that `HashAggregateExec` uses, but the `CometSort` runs in the
same native plan, so DataFusion sees input sorted on every grouping key and
picks its ordered aggregate streams. Spark needs that sort only because of its
hash map limitation. In Comet it costs memory and CPU (DataFusion 55.2):
1. `ExternalSorter` reserves each buffered batch at its memory size plus its
sliced size, about 2x. When it spills, `sort_and_spill_in_mem_batches` frees
the `sort_spill_reservation_bytes` headroom, then asks again for the row-format
keys and the merge, and another task can take that memory first
(apache/datafusion#24739, fixed on DataFusion main by apache/datafusion#24740,
not in 55.x).
2. On input sorted on every grouping key, the partial aggregate
(`OrderedPartialAggregateStream`) registers as non-spillable and returns
`ResourcesExhausted` when a reservation is refused. The hash partial
(`PartialHashAggregateStream`) emits its table early instead.
3. The sort compares every row on every grouping key.
Example: a job grouping by 87 columns, most of them `bigint`, with `sum` and
string `min`/`max`. About 3M rows per task, 12 tasks per executor, 32 GiB
off-heap. The plan is
```
CometScan -> CometSort -> CometSortAggregate (Partial) -> CometExchange ->
CometSort -> CometSortAggregate (Final)
```
39 of the 40 map tasks failed with
```
Additional allocation failed for OrderedPartialAggregateStream[0] with top
memory consumers (across reservations) as:
ShuffleRepartitioner[0]#12(can spill: true) consumed 2.9 GB, peak 2.9 GB,
ExternalSorterMerge[0]#1(can spill: false) consumed 256.0 MB, peak 3.0 GB,
OrderedPartialAggregateStream[0]#6(can spill: false) consumed 473.9 KB,
peak 11.9 MB.
```
and 18 retries failed in `ExternalSorterMerge[0]` during the sort's first
spill. Both failures come from the sort.
### Describe the potential solution
When a `SortAggregateExec`'s child is the `SortExec` on its grouping keys,
plan the native aggregate over the sort's child, so DataFusion runs its hash
streams.
1. Partial aggregates whose parent is a shuffle exchange: drop the sort, set
`ordered_by_grouping_keys` to false, and report no output ordering. The
exchange hides the order.
2. Final and Complete aggregates: drop the input sort, and keep the order
wherever something above may rely on it.
- `planner/aggregate.rs` already sorts the aggregate's output when
`ordered_by_grouping_keys` is set and the input is not sorted. That sorts
aggregated rows instead of input rows, which is cheaper whenever the aggregate
reduces rows.
- `RewriteJoin.restoreRequiredOrdering` re-adds a local sort under any
operator whose child no longer satisfies its `requiredChildOrdering`. It cannot
see a sort that Spark already removed because of the aggregate's reported
ordering (the comment on `CometSortAggregateExec.orderedByGroupingKeys`
describes this case). So emit unordered output only when the parent is an
exchange, and sort the output otherwise.
3. Keep `first`, `last` and `collect_list` on today's path at first, since
their results follow input order.
4. Put it behind a config, off by default until benchmarks show where hash
aggregation wins.
Tests: plan checks that no `CometSort` remains under a partial aggregate.
Ordering checks where a parent relies on the aggregate's order, for example
`sortWithinPartitions` on the keys or a sort-merge join on them. Spill runs
with a small pool. Benchmarks over high and low key cardinality, with string
`min`/`max`.
### Additional context
- Apache Auron already does this. `AuronConverters.convertSortAggregateExec`
runs a `SortAggregateExec` as `NativeAggBase.HashAgg` over
`exec.child.children.head` when the child is a sort and the aggregate's
`childOrderingRequiredTag` is not set. `AuronConvertStrategy` sets that tag
from each parent's `requiredChildOrdering`, and from a write's
`requiredOrdering`, and passes it down until a `SortExec`.
- #6936: string and binary `min`/`max` under-report memory in hash
aggregates. This change makes that path more common, so it should land first or
together.
- Distinct aggregates add `PartialMerge` sort aggregates between exchanges.
Without the sort they would run on unsorted input, which may also sidestep
#6880, but their parents need the same ordering check.
- Related: #1994, #4565, #6909, apache/datafusion#25804.
--
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]