comphead opened a new issue, #26575:
URL: https://github.com/apache/datafusion/issues/26575
### Describe the bug
A `Partial` aggregate over input sorted on all grouping keys
(`ClusteredPartialAggregateStream`, `OrderedPartialAggregateStream` in 55.x)
reserves memory for the whole batch it just aggregated before emitting the
completed groups. If that reservation is refused under `GroupClustering::Full`,
`Aggregating::handle_stage` returns `Err(oom)`. The memory is already
allocated, and emitting the completed groups would shrink the table to the one
active group. So in a shared pool the query fails although it could continue.
This is a regression from 54. `GroupedHashAggregateStream`, still used with
`enable_migration_aggregate = false`, emits early for every `Partial`
aggregate, and for fully ordered input `GroupOrdering::oom_emit_to` emits only
the completed groups.
### To Reproduce
In DataFusion Comet on 55.2 (87 sorted grouping keys, a spillable shuffle
writer in the same pool), 39 of 40 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.
```
The stream held 474 KB between batches and asked for 11,959,496 bytes after
the next one.
### Expected behavior
On a refused reservation with full clustering, emit
`take_completed_state_batch()` and keep only the active group, as the legacy
stream does. Return the error only when no group is complete.
`Outputting::handle_stage` already reserves just the retained table when it
cannot hold the output batch, and the emitted rows are the ones the success
path emits, so results and order stay the same.
### Additional context
Related: #22710, #25637, apache/datafusion-comet#7146.
--
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]
