viirya opened a new issue, #26725: URL: https://github.com/apache/datafusion/issues/26725
### Describe the bug #26576 (for #26575) changed `ClusteredPartialAggregateStream` so that a refused reservation under full group clustering emits the completed groups instead of failing. The emit leaves one active group, but the table keeps the capacity it grew for the batch, so `memory_size()` barely drops. The output stage then re-reserves the table and fails again: https://github.com/apache/datafusion/blob/cd571e5124b019917960fb948dfc14c5923e7fa4/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs#L319-L329 https://github.com/apache/datafusion/blob/cd571e5124b019917960fb948dfc14c5923e7fa4/datafusion/physical-plan/src/aggregates/clustered_partial_stream.rs#L387-L392 Full clustering emits `EmitTo::First(len - 1)`, so the group values and accumulators shrink. What stays allocated: - `GroupValuesPrimitive` (single primitive key): `HashTable::retain` keeps the table capacity. - `GroupValuesClustered` (multiple keys): the `new_groups` scratch vector. - `group_indices`, in every case. With 8192 distinct keys per batch, `memory_size()` before and after the emit, compared with a table that holds only the active group: | keys | after aggregating | after emitting completed groups | one-group table | |---|---|---|---| | `Int64` | 475 KB | 344 KB | 5.6 KB | | `Int64, Int64` | 328 KB | 131 KB | 0.4 KB | | `Utf8` | 934 KB | 82 KB | 17 KB | | `Utf8, Int64` | 459 KB | 131 KB | 8.6 KB | The retained size is bounded by the groups in one batch and does not grow with the total number of groups. A limit that fits the active group but not one batch's capacity still fails. The existing test, `completed_groups_under_memory_pressure`, sets its limit to the largest reservation of an unlimited run, and that reservation already includes the retained capacity. `take_state_batch`, used under partial group clustering, calls `clear_shrink` and shrinks `group_indices`, so that path releases the memory. ### To Reproduce A partial `SUM(v) GROUP BY k` over three batches of 8192 distinct sorted `Int64` keys, with a 100 KiB `GreedyMemoryPool`: ``` Failed to allocate additional 330.8 KB for ClusteredPartialAggregateStream[0] with 5.4 KB already allocated for this reservation - 94.6 KB remain available for the total memory pool: greedy(used: 5.4 KB, pool_size: 100.0 KB) ``` With two `Int64` keys, it fails to allocate an additional 128.0 KB. From SQL, `select k0, sum(v) from t group by k0` over a `MemTable` sorted on `k0` with 4 partitions plans `AggregateExec: mode=Partial, ..., group_clustering_mode=Full`. With a shared `GreedyMemoryPool`, the query fails in this stream at some limits (for example 1.25–3 MB with 4 partitions). At those limits, the non-spilling downstream operators (`RepartitionExec[Merge]`, `SortPreservingMergeExec`, `ClusteredFinalAggregateStream`) also fail once the partial stream gets past this point. A shared limit therefore does not separate this stream. A pool that limits only `ClusteredPartialAggregateStream` reservations reproduces it end to end. ### Expected behavior After the memory-pressure emit, the table should hold only what the active group needs, so the stream finishes under a limit that fits that group. ### Additional context The fix I plan is to use `take_state_batch` under full group clustering too, as partial group clustering already does: emit every state, including the active group, and reset the table. The final stage merges the active group's two state rows, as it merges the same key from different partitions or from early emission in `GroupedHashAggregateStream`. ### Willingness to contribute I plan to contribute a fix for this bug independently -- 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]
