[viirya]

Which issue does this PR close?

• Closes Fully sorted partial aggregate still fails with ResourcesExhausted after emitting completed groups, because the table keeps its capacity #26725.

Rationale for this change

When a reservation is refused, the clustered aggregate streams emit their completed groups instead of failing. #26576 did this for the partial stream under full group clustering, and #26637 did it for the final and single streams. The remaining table keeps the capacity it grew for the batch, so re-reserving it can still fail with ResourcesExhausted. With 8192 distinct Int64 keys per batch, the partial table still needs 344 KB, while a table holding only the active group needs 5.6 KB. The issue has the details.

What changes are included in this PR?

• GroupValues::shrink_to_fit (new, does nothing by default) releases capacity that the stored groups do not need. It is implemented for:

  • GroupValuesPrimitive: the hash table and values.
  • GroupValuesClustered: new_groups.
  • GroupValuesColumn: the hash table and the scratch buffers of one intern call.
  • GroupValuesRows: the hash table, hashes, and rows buffers.

  GroupValuesBytes and GroupValuesBytesView already rebuild their map on EmitTo::First.

• ClusteredAggregateTable::shrink_to_fit (new) calls it and also releases group_indices.

• After a refused reservation, ClusteredPartialAggregateStream (under full group clustering), ClusteredFinalAggregateStream, and ClusteredSingleAggregateStream emit the completed groups as before and then call shrink_to_fit. The output does not change, and each stream still returns the refusal when no group is complete. The path without memory pressure does not call it.

Accumulator scratch, such as the APPROX_DISTINCT hashes, is not released yet; see #26781.

What is the testing strategy for this PR?

• clustered_partial_stream::tests::completed_groups_under_memory_pressure runs three batches of 8192 distinct keys whose boundaries share a key, under a 100 KiB limit, and compares the output exactly with an unlimited run. It covers one key each for GroupValuesPrimitive (Int64), GroupValuesClustered (Int64, Int64), GroupValuesColumn (Float64, Float64), and GroupValuesRows (Int64, List). All four cases fail on main with ResourcesExhausted. The Int64, Float64, Float64, and Int64, List cases also fail if only group_indices is released.
• completed_groups_release_batch_capacity in clustered_final_stream and clustered_single_stream runs the same input with an Int64 key under the same limit, using the new distinct_key_batches and assert_output_unchanged_under_limit helpers in clustered_test_util. Both fail on main with ResourcesExhausted.

Are there any user-facing changes?

No. GroupValues gains a method with a default implementation.

This pull request and its description were written by Isaac.

[bvolpato]

LGTM on 175d52158051. The clustered partial stream tests, SQL memory-limit regression, and aggregate memory call-sequence tests passed locally. I also checked overlapping group keys across batch and partition boundaries. Removing just the production change reproduces ResourcesExhausted under the 100 KiB limit. Restoring it passes.

GitHub currently reports merge conflicts with main. This approval covers the tested head. The conflict resolution still needs validation before merging.

[comphead]

Thanks @viirya. The change is small, and the final stage already merges split partial states, as it does for early emission in GroupedHashAggregateStream. Comments inline. Two smaller points:

• batches() (line 443) builds num_keys + 1 identical Int64Arrays per batch. Building one and using Arc::clone for each column avoids the copies.
• This PR conflicts with #26637 in the clustered_partial_stream.rs tests. The conflict is also semantic. #26637 moves this test into a shared clustered_test_util helper that compares the output exactly at held (assert_completed_groups_emitted_under_limit), and the split row added here breaks that comparison. Whichever PR lands second needs to compare the partial case by merged key.

[comphead]

Design question. The retained capacity sits in the shared ClusteredAggregateTable buffer: group_indices, the GroupValuesPrimitive hash table, and the clustered new_groups scratch. ClusteredFinalAggregateStream and ClusteredSingleAggregateStream emit through the same EmitTo::First path and cannot emit the active group. So I expect them to fail the same way under a limit below one batch of capacity, even after #26637. #26637 already hit this retention under partial clustering.

Releasing that capacity on the refusal path after take_completed_state_batch would cover all three streams with one change. For example, call group_indices.shrink_to_fit() and add a shrink hook on the group values. It would also keep one state row per group and match the legacy GroupedHashAggregateStream, which never emits the active group under full clustering. Was that considered? If the reset is preferred for Partial, it would help to say in the PR that Final and Single need a different fix.

[viirya]

Yes, switched to this in e740377. On a refusal under full clustering, the stream emits the completed groups and then calls a new ClusteredAggregateTable::shrink_to_fit. It clears and shrinks group_indices and calls a new GroupValues::shrink_to_fit, which does nothing by default. The implementations:

• GroupValuesPrimitive: the hash table and values.
• GroupValuesClustered: new_groups.
• GroupValuesColumn: the hash table and hashes_buffer. In the non-streaming case, also the emit and vectorized scratch.
• GroupValuesRows: the hash table, hashes_buffer, and rows_buffer.

GroupValuesBytes and GroupValuesBytesView rebuild their map on EmitTo::First, so they need nothing.

The stream emits one state row per group again, and it returns the refusal when no group is complete, as the legacy stream does. Final and Single can call the same method after their completed-groups emit.

[comphead]

This shows the partial stream no longer has an error path for a non-empty batch under full clustering. It finishes with one group under a 6,000-byte limit. The legacy stream returns the refusal here, because it only emits early when group_values.len() > 1. If that is intended, the name full_clustering_refusal_is_an_error no longer fits the partial case. Splitting the partial case into its own test, or renaming it, would keep the name accurate.

[viirya]

Since the stream again emits only the completed groups, the partial case returns the refusal here again. I restored the original snapshot, so the name fits again.

[comphead]

This test and completed_groups_under_memory_pressure now go through the same branch. If the ranges here overlap by one key (for example 0..8193, 8192..16385, 16384..24576), this test also covers the split state rows, under a limit that fails on main. You could then assert the extra rows, drop the older test, and drop the held sampling from run.

[viirya]

Done. completed_groups_under_memory_pressure now uses the overlapping ranges, so the active group spans each batch boundary. It runs under the 100 KiB limit, which fails on main. It covers the four group-values implementations that implement the new hook: Int64, Int64, Int64, Float64, Float64, and Int64, List. Without the hook, the Int64, Float64, Float64, and Int64, List cases still fail when only group_indices is released. Because the output no longer changes, the test compares it exactly. The older test and the held sampling are removed.

[comphead]

This duplicates memory_pressure_releases_retained_capacity. Both use 8192 distinct keys per batch, one and two Int64 keys, the 100 KiB limit, and compare the merged results. The only extra check is that the planner picks mode=Partial with group_clustering_mode=Full, and .slt plans already show that (group_by.slt, agg_func_substitute.slt). Dropping this file also removes PartialStreamLimitPool and the second batch builder.

If it stays, run plans the query twice: once in create_physical_plan and again in df.collect(). Collecting the plan that was already built avoids the second planning pass.

[viirya]

Agreed, removed.

[sunchao]

Thanks @viirya. The new SUM regressions passed on 67ce0599. I found a sort returning out-of-order rows and a remaining retained-memory failure in APPROX_DISTINCT. I’d like to see both addressed before approval. For the local reproducers I used the same compatible dependency versions for the PR and the controls because my registry lacked some of the locked versions. Fork CI test job

[sunchao]

Splitting the active group makes the advertised Unique(k) constraint incorrect even within one fully clustered partition. I used the two sorted batches from completed_groups_under_memory_pressure, with v=9 and v=1 for the two occurrences of k=256. With a 13,000-byte limit, SortExec(k ASC, sum_state ASC) directly above the partial returned (256, 9) then (256, 1). ordering_satisfy accepted the secondary ordering because k was marked unique, so SortExec passed the rows through. Restoring the base emission gave (256, 10) at the same limit. Could you keep the ordering by k but withhold group-key uniqueness for partial outputs that can split, and add a regression for the secondary sort? Constraint declaration · Sort fast path

[viirya]

Thanks, confirmed. With your inputs and the 13,000-byte limit, a SortExec(k, sum) above the partial passes (256, 9), (256, 1) through unsorted. The stream no longer resets the table under full clustering. It emits the completed groups and releases the retained capacity instead, so the active group stays in one row, and Unique(k) holds within a partition again. completed_groups_under_memory_pressure compares the output exactly under the limit, so it fails if a group is split again.

[sunchao]

The reset still leaves batch-sized scratch in APPROX_DISTINCT. I ran a fully clustered partial with APPROX_DISTINCT(k) and APPROX_DISTINCT(v), one sorted batch of 8,192 distinct Int64 keys, and v=42. Each HllGroupsAccumulator keeps a 64 KiB hashes vector after state(EmitTo::All), so reserving the empty table still fails at the PR’s 100 KiB limit. Releasing those vectors on the full drain let the same run emit all 8,192 rows. Could the pressure reset release this accumulator scratch as well, with a regression for this case? State emission · Scratch accounting

[viirya]

Confirmed. HllGroupsAccumulator::size counts the hashes scratch, and nothing releases it after state. That memory belongs to the accumulator rather than the table buffer. Releasing it needs a shrink hook on GroupsAccumulator, which reaches every accumulator that keeps per-batch scratch. I'd like to keep this PR to the group values, and I opened #26781 for the accumulators.

[viirya]

Thanks @comphead. batches() now builds one array per batch and clones it into each column. Now that the split row is gone, the conflict with #26637 is only textual.