[yashaggarwal85d]

Which issue does this PR close?

• Closes Final and Single aggregates over fully sorted input fail with ResourcesExhausted instead of emitting completed groups #26628.

Rationale for this change

A Final or Single aggregate over input sorted on all grouping keys fails with ResourcesExhausted under a memory limit that the legacy GroupedHashAggregateStream handles. After each input batch, ClusteredFinalAggregateStream and ClusteredSingleAggregateStream reserve memory for the whole table before emitting completed groups. If that reservation is refused and the stream cannot spill (full group clustering, or temporary files disabled), it returns the error, even though emitting the completed groups would leave only the active group in memory.

This is the same problem as #26575 for the partial stream, fixed in #26576.

What changes are included in this PR?

• ClusteredFinalAggregateStream: when the reservation is refused and there is no spill context, take the completed groups and move to the output stage, then resume aggregation. The error is returned only when no group is complete.
• ClusteredSingleAggregateStream: the same behavior, emitting through the existing start_outputting path and resuming ReadingInput.
• Updated doc comments to describe the new behavior.

The spilling path (partial group clustering with temporary files enabled) is unchanged.

What is the testing strategy for this PR?

New unit tests, modeled on the one in #26576:

• clustered_final_stream::tests::completed_groups_under_memory_pressure
• clustered_single_stream::tests::fully_sorted_completed_groups_under_memory_pressure

Each runs SUM(k) GROUP BY k over two batches sorted on k that share a boundary key. It first runs with an unbounded pool to measure the most memory the stream holds when it hands off output, then reruns with a GreedyMemoryPool of that size and checks that the results match. Both tests fail on main with ResourcesExhausted and pass with this change. All aggregates unit tests in datafusion-physical-plan pass, and cargo fmt and cargo clippy --all-targets --all-features -D warnings on the crate are clean.

Are there any user-facing changes?

No API changes. Queries that used to fail with ResourcesExhausted in this situation now complete.

---

This solution was developed with AI assistance; I reviewed and tested the changes.

[github-actions]

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

     Cloning apache/main
    Building datafusion-physical-plan v55.1.0 (current)
       Built [  44.212s] (current)
     Parsing datafusion-physical-plan v55.1.0 (current)
      Parsed [   0.251s] (current)
    Building datafusion-physical-plan v55.1.0 (baseline)
       Built [  43.410s] (baseline)
     Parsing datafusion-physical-plan v55.1.0 (baseline)
      Parsed [   0.258s] (baseline)
    Checking datafusion-physical-plan v55.1.0 -> v55.1.0 (no change; assume patch)
     Checked [   1.078s] 229 checks: 228 pass, 1 fail, 0 warn, 31 skip

--- failure inherent_method_missing: pub method removed or renamed ---

Description:
A publicly-visible method or associated fn is no longer available under its prior name. It may have been renamed or removed entirely.
        ref: https://doc.rust-lang.org/cargo/reference/semver.html#item-remove
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.51.0/src/lints/inherent_method_missing.ron

Failed in:
  AggregateExec::input_keeps_groups_local, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/c91879df6bf841ecef350e933fd580c2a37d4b29/datafusion/physical-plan/src/aggregates/mod.rs:1663

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  91.186s] datafusion-physical-plan

[comphead]

Thanks @yashaggarwal85d I'll check this today, also @jayzhan211 FYI as you also reviewed #26576

[comphead]

Thanks @yashaggarwal85d. The new branch emits the same rows the success path already emits after a successful reservation, and it keeps the error when no group is complete. I did not find a correctness problem. Points below, most important first.

• The cargo-semver-checks failure (AggregateExec::input_keeps_groups_local removed) is not caused by this PR. The branch is based on 8ce9f6038 and does not contain that method, which exists on main (aggregates/mod.rs:1663). Merging or rebasing on main should clear it and the auto detected api change label.
• Test coverage gaps are in the inline comments (partial clustering without temp files, the still-fails case).
• Fuzz: the invariant is "output under a pool limit equals the unlimited output". One deterministic case with two batches and one boundary key checks one shape. A small randomized variant would cover more boundary shapes: a batch made entirely of the boundary key (zero completed groups), a batch with a single new key, several batches per key, multi-column keys and string keys. It could extend the existing aggregate fuzz cases rather than add a framework.
• Naming nit: the helper is run_fully_sorted here and run in clustered_partial_stream.rs. The single stream test is fully_sorted_completed_groups_under_memory_pressure while the same module already has completed_groups_under_memory_pressure. Use one naming scheme across the three modules.
• Comment nit: the new behavior is described in the type docs, the stage list and again inline in both streams. Keeping the type doc and a one-line inline why is enough.

[comphead]

The new None branch repeats the tail of the success path below: take_completed_result_batch, reservation_size_for_table, start_outputting(.., ReadingInput { table, spill_context }). It also adds two match arms that both call break_with_err, with the second e shadowing the OOM error.

With spill_context == None the has_spills check below is always false, so the success path already does the right thing. Keeping the refused error in an Option and falling through would remove the duplicate block. The only extra logic is in the Ok(None) arm, which returns the saved error instead of continuing. The same shape applies to ClusteredFinalAggregateStream::Aggregating::handle_stage, where it saves a third copy of the Outputting { batch, resume: Some(self) } construction.

[comphead]

Design question for both streams. The fix uses a refused try_resize as the trigger and then emits. The issue points at the legacy GroupedHashAggregateStream, which emits completed groups before it updates the reservation. For the non-spilling case, taking the completed groups first and then sizing the reservation once would not need an error path, and a failure would then mean the active group itself does not fit. Is there a reason to prefer the error-driven fallback, for example keeping the success path untouched? If so, a short comment saying why would help. Note that neither variant lowers the real peak, since aggregate_batch has already allocated the new groups. It only makes the accounting consistent.

[comphead]

Both new tests assert GroupClusteringMode::Full. The changed branch is spill_context == None, which is also true for Partial clustering with temporary files disabled (DiskManager without tmp files). That case reaches the new code with an active key range defined by the sorted prefix, and it is not covered. A case with a two-column key sorted on the first column and tmp files disabled would cover it, for the final and the single stream.

[comphead]

The test only checks that a tight limit now succeeds. Nothing checks the other half of the contract in the docs: the error is still returned when no group can be released. A rerun with a limit just below held (or a single-key input) that expects ResourcesExhausted would show the fix does not swallow genuine exhaustion and exercises the ok_or(oom) and Ok(None) arms.

[comphead]

This driver (context with RuntimeEnvBuilder, the held sampling loop, unlimited run then GreedyMemoryPool::new(held) run, batches_to_string comparison) is now the third copy, after clustered_partial_stream.rs on main and run_fully_sorted in clustered_final_stream.rs. Only the plan construction differs, and the final one needs partial-state input. AggregateExec::execute already dispatches to the clustered stream by mode, since enable_migration_aggregate defaults to true. One shared helper that takes Arc<dyn ExecutionPlan>, runs it on both pools and returns the output and held would leave each module with only its plan builder. The helper could live next to the existing aggregate test utilities in aggregates/mod.rs.

The unit-test form makes sense here since a SQL memory limit cannot target one stream's table size. If a SQL form is feasible, aggregate_memory_spill.slt already sets a 1M limit and notes that sorted generate_series takes the streaming path. I have not tried it, so it may be too fragile.