[yashdotrv]

Which issue does this PR close?

Addresses #26705.

Rationale for this change

A partial aggregate can emit a state batch containing all its input values and then release its memory reservation. The final aggregate merges that batch before reserving memory for it. For aggregates such as median and array_agg, this leaves input-sized state unaccounted for by the memory pool.

What changes are included in this PR?

Final aggregates now reserve incoming state batches before merging them. The hash and clustered final paths can spill their current state and retry when the pool cannot admit a batch. I also updated the affected memory reservation and spill tests.

PR #26728 adds related memory accounting tests, but its cases for this issue are currently ignored.

What is the testing strategy for this PR?

I added reservation tests for the ungrouped, hash grouped, and clustered grouped final aggregate paths. The extended workspace test suite, formatting check, and Clippy pass.

I also ran the issue’s 30-million-row, four-partition median reproducer with a 300 MB memory limit. It now returns ResourcesExhausted when the incoming state cannot be reserved. Peak RSS can still exceed the configured limit; reducing temporary memory use during buffer growth is outside this PR.

Are there any user-facing changes?

Under memory pressure, a final aggregate may now spill or return ResourcesExhausted instead of continuing with unreserved state. There are no public API changes.

[jayzhan211]

@yashdotrv thanks for picking this up. Please correct the description, since most of the state is still unreserved and buffer growth is not the main cause. With the ignored #26705 cases in #26728 turned on, unreserved memory grows 121-124 MB for median (122-143 MB on main) and 500-505 MB for array_agg (531-589 MB), against a 16 MB bound. At the peak 19.1 MB of 152 MB is reserved for median, about one of the four partial states; the others appear to be queued in CoalescePartitionsExec.

A follow-up is fine because that queue is outside these streams, so #26705 should stay open. The description should also say that array_agg over 1000 groups at 30-60 MB goes from passing to ResourcesExhausted, while peak RSS at 100 MB stays near 610 MB.

[yashdotrv]

@jayzhan211 ty for your review, I'll work on the suggestions and get back.

[yashdotrv]

@jayzhan211 Thanks again for the review and detailed pointers! I have updated the branch with the requested changes:

1. Per-Row Charging Threshold: Replaced the fixed 64 KiB threshold with MAX_IN_FLIGHT_STATE_BYTES_PER_ROW = 1024 and MIN_CHARGED_STATE_BATCH_BYTES = 64 * 1024. Normal bounded batches remain exempt, avoiding unnecessary spills/refusals on queries with many aggregates.
2. Soft Limit Early Break: Kept the early break for soft group limit before try_resize in FinalHashAggregateStream.
3. Spill Logic Cleanup: Extracted the duplicate spill-and-resize logic in FinalHashAggregateStream into a shared spill_table helper.
4. Ungrouped Aggregate Guard: Guarded the calculation of incoming state batch memory in AggregateStream strictly to AggregateInputMode::Partial.
5. Call Sequence Tests Restored: Restored aggregate_memory_call_sequence.rs to match main without edits; all 8 snapshot tests pass.
6. Test Coverage:
   • Updated final_aggregate_reserves_incoming_state_batch and final_hash_reserves_incoming_state_batch to use 64-row sliced batches retaining multi-megabyte underlying buffers.
   • Added clustered_final_refuses_large_state_batch.
   • Added final_hash_spills_to_admit_large_state_batch.
   • Added clustered_final_spills_to_admit_large_state_batch to test the spill retry path under GroupClusteringMode::Partial(vec![0]).
7. PR Description: Clarified that Partial aggregate state reaches the final aggregate without a memory reservation (median, array_agg) #26705 remains open for the queued partial states in CoalescePartitionsExec.