[viirya]

Which issue does this PR close?

• Closes Fully sorted partial aggregate keeps per-batch accumulator scratch reserved after emitting completed groups (APPROX_DISTINCT) #26781.

Rationale for this change

#26726 lets a fully clustered aggregate (partial, final and single) continue under memory pressure: it emits the completed groups, then releases the capacity that the group values and group_indices grew for the batch, so that reserving the table fits again. Accumulators can also keep per-batch scratch that size() counts, and GroupsAccumulator had no way to release it. Such a query still fails with ResourcesExhausted, although the table holds only the active group.

For example, with 8192-row batches and a 100 KiB pool:

• APPROX_DISTINCT(k), APPROX_DISTINCT(v) GROUP BY k: each HllGroupsAccumulator keeps a 64 KiB hashes buffer, and the stream fails to reserve ~123 KB.
• ANY_VALUE(v) GROUP BY k, or any aggregate without a native GroupsAccumulator: GroupsAccumulatorAdapter keeps its per-batch Scratch (groups_with_rows, offsets, batch_indices, ~160 KiB). It releases that only when every group is emitted, but the clustered emit keeps the active group. The stream fails to reserve ~283 KB.

What changes are included in this PR?

• Add GroupsAccumulator::shrink_to_fit, which does nothing by default, documented like GroupValues::shrink_to_fit.
• ClusteredAggregateTable::shrink_to_fit calls it for every accumulator, through HashAggregateAccumulator::shrink_to_fit. The clustered partial, final and single streams all call ClusteredAggregateTable::shrink_to_fit after the memory-pressure emit, so all three release accumulator scratch.
• ClusteredAggregateTable::take_state_batch (partial group clustering emits every state) and the non-clustered AggregateHashTable::take_state_batch (partial hash early emit, single/final hash spills) call it too, since the accumulators otherwise keep their scratch reserved after every state is emitted.
• HllGroupsAccumulator drops hashes.
• GroupsAccumulatorAdapter drops its Scratch through a shared free_scratch(), and the empty row indices of the remaining groups. Both update allocation_bytes.
• FirstLastGroupsAccumulator (ordered FIRST_VALUE / LAST_VALUE) rebuilds its per-batch scoreboard bitmap at the remaining group count, since take_state truncates it but keeps its capacity.

The other GroupsAccumulator implementations keep only per-group state (or no counted per-batch scratch), so they keep the default. ForeignGroupsAccumulator (FFI) also keeps the default, because forwarding the call would need a new function pointer in FFI_GroupsAccumulator. That can be a follow-up if a foreign accumulator needs it.

What is the testing strategy for this PR?

New tests, each of which fails with ResourcesExhausted without its part of the change, under a GreedyMemoryPool (64 KiB unless noted) that fits the active groups but not the scratch:

• clustered_partial_stream::tests::completed_groups_release_accumulator_scratch: full group clustering, each batch boundary shares a key, exact comparison with an unbounded run. Covers APPROX_DISTINCT (HllGroupsAccumulator) and ANY_VALUE (GroupsAccumulatorAdapter).
• clustered_partial_stream::tests::active_group_releases_adapter_indices: the active group holds most rows of the batch, so the adapter must release its indices (16 KiB pool).
• clustered_partial_stream::tests::completed_groups_release_first_value_bitmap: FIRST_VALUE(v ORDER BY v) over 131,072 groups (16 KiB pool).
• clustered_partial_stream::tests::emitted_states_release_accumulator_scratch: partial group clustering (take_state_batch).
• hash_stream::tests::emitted_states_release_accumulator_scratch: unsorted input, partial hash early emit (AggregateHashTable::take_state_batch), checked through early_emit_count.

clustered_test_util::sum_over_sorted_input now delegates to a new aggregate_over_sorted_input that takes the aggregate function, which in turn uses aggregate_expr_over_sorted_input to configure the expression (for example ORDER BY v). A new collect_under_limit runs a plan under a limit.

Are there any user-facing changes?

A new provided method, GroupsAccumulator::shrink_to_fit, with a no-op default. Existing implementations don't need to change.

This pull request and its description were written by Isaac.

[alamb]

Looks good to me. Thanks @viirya

I also had the agent do an ablation test and confirm that tests fail without this chaneg

---- aggregates::clustered_partial_stream::tests::accumulator_scratch_under_memory_pressure stdout ----
Error: Context("aggregates approx_distinct", ResourcesExhausted("Failed to allocate additional 122.9 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)"))

---- aggregates::clustered_single_stream::tests::completed_groups_release_accumulator_scratch stdout ----
Error: ResourcesExhausted("Failed to allocate additional 58.9 KB for ClusteredSingleAggregateStream[0] with 5.4 KB already allocated for this reservation - 58.6 KB remain available for the total memory pool: greedy(used: 5.4 KB, pool_size: 64.0 KB)")

[alamb]

I think this is the actual change to call shrink-to-fit on the accumulators

[comphead]

take_state_batch below has the same gap. The partial stream calls it in the OOM branch for partial group clustering. It clears the key and index buffers but leaves HllGroupsAccumulator::hashes counted, and Outputting then reserves memory_size() again. From reading the code (I haven't run it), APPROX_DISTINCT(k0), APPROX_DISTINCT(v) GROUP BY k0, k1 with the input sorted on k0 only should still fail under the 100 KiB pool, because the two hashes buffers alone are about 128 KB. Could take_state_batch run the same accumulator loop?

[viirya]

Good catch, confirmed. With APPROX_DISTINCT(v) GROUP BY k, k2 over input sorted on k only, a 64 KiB limit failed after the emit with "Failed to allocate additional 64.0 KB ... with 600.0 B already allocated": the table was empty, but hashes was still counted.

The non-clustered AggregateHashTable::take_state_batch had the same gap. It is used by the partial hash stream's early emit and by the single/final hash streams before spilling. A partial APPROX_DISTINCT(v) GROUP BY k over unsorted input failed the same way under 64 KiB.

In 32b641f, both take_state_batch now call shrink_to_fit on every accumulator. New tests, each of which fails without its fix:

• clustered_partial_stream::tests::emitted_states_release_accumulator_scratch (partial group clustering)
• hash_stream::tests::emitted_states_release_accumulator_scratch (no clustering)

[comphead]

Thanks @viirya. Releasing scratch through a no-op trait default mirrors GroupValues::shrink_to_fit and leaves existing implementations untouched, which looks right to me.

Design questions:

• The next update_batch regrows the scratch, so under sustained pressure every batch allocates and frees it again, and the real peak during update_batch stays the same. Only the reservation between batches fits. That matches how #26726 treats group_indices, so it is fine for this PR. For HllGroupsAccumulator, hashing in fixed-size chunks (slices of array are zero-copy) would bound hashes regardless of batch size and make its override unnecessary. Worth a follow-up?
• The shrink only runs after a completed group is emitted, because take_completed_state_batch()?.ok_or(oom)? comes first. A batch that holds one long-running key still fails under the same limit, although shrinking alone would make the table fit. Was shrinking before giving up considered?

Nit:

• groups.shrink_to_fit() in HllGroupsAccumulator and states.shrink_to_fit() in the adapter look like no-ops on this path. With full clustering only the active group stays, and split_vec_min_alloc already moves it into a fresh allocation when most groups are emitted. Dropping them would keep both overrides to the scratch they actually release.

[comphead]

free_scratch_if_empty already has the take-and-free logic. Could both call a small free_scratch() helper, so the scratch accounting lives in one place?

[viirya]

Done. free_scratch_if_empty and shrink_to_fit both call a new free_scratch().

[comphead]

This could reuse the shared helpers that completed_groups_release_accumulator_scratch already uses. That drops the run_aggregates split, the local build, the Aggregates alias and three imports:

for udaf in [approx_distinct_udaf(), any_value_udaf()] {
    let plan = aggregate_over_sorted_input(
        AggregateMode::Partial,
        &["k"],
        distinct_key_batches(),
        udaf,
    )?;
    assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full);
    assert_output_unchanged_under_limit(&plan, 64 * 1024).await?;
}

It builds one APPROX_DISTINCT instead of two, so the limit has to sit below a single hashes buffer, as in the single-mode test. I expect 64 KiB to work for both cases but haven't run it.

All three streams go through the same ClusteredAggregateTable::shrink_to_fit, and the single stream's call site is already covered by completed_groups_release_batch_capacity. So with the loop above, completed_groups_release_accumulator_scratch repeats the same path and one of the two tests could go.

[viirya]

Done, thanks. 64 KiB works for both APPROX_DISTINCT and ANY_VALUE. Without GroupsAccumulator::shrink_to_fit, they fail asking for 58.9 KB and 283.0 KB. The test is now completed_groups_release_accumulator_scratch. I dropped run_aggregates and the single-mode test.

[viirya]

Thanks @alamb for the review and the ablation.

@comphead, thanks for the careful review. I pushed 32b641f:

• The take_state_batch gap is fixed for both the clustered and the non-clustered hash tables (details inline).
• Dropped groups.shrink_to_fit() and states.shrink_to_fit(). You're right that they are no-ops on these paths: with full clustering, split_vec_min_alloc already moves the active group into a fresh allocation, and take_state_batch emits EmitTo::All.
• Added the free_scratch() helper and consolidated the tests.

On the design questions:

• Chunked hashing in HllGroupsAccumulator sounds good as a follow-up: it would bound hashes regardless of batch size and avoid regrowing it under sustained pressure. The hook is still needed for GroupsAccumulatorAdapter, whose per-batch indices also scale with the batch.
• Shrinking before giving up: confirmed. With three 8192-row batches of a single key, the partial stream fails a 64 KiB limit asking for 132.1 KB, but the table is about 4.3 KB after shrink_to_fit. The final and single streams have their own no-completed-groups paths (spill or error) that would change too, so I'd like to handle all three in a follow-up: Clustered aggregation: release retained capacity before failing or spilling when no group is complete #26911.

[sunchao]

The three new regression tests passed. I found two sorted inputs where the stream completes a group, tries to release scratch, and still exhausts the pool. Releasing the specific buffer identified in each comment let the same input finish under the same 16 KiB limit.

I ran 348 repository aggregate tests against the production source at 32b641fb7. All passed. The PR build and the control build used Rust 1.99.0 and the same compatible dependency lock with Arrow 60.0.0 because the local registry lacked the pinned uuid 1.29.0.

[sunchao]

[P2] Release the active group’s empty row-index buffer as well

free_scratch() releases the adapter-wide Scratch, but each surviving AccumulatorState also has an indices: Vec<u32>. invoke_per_accumulator charges its capacity as it grows, then calls indices.clear() after consuming the rows. The new hook leaves that empty allocation charged.

I reproduced this with a fully sorted partial ANY_VALUE(v) GROUP BY k, one 8,192-row batch, two keys, and a 16 KiB greedy pool. With keys [0, 1, 1, …, 1], group 0 completes and the active group retains 8,192 index slots. The stream then fails reserving another 27.1 KiB. Reversing the frequencies to [0, 0, …, 0, 1] returns both groups under the same limit. Independently, an adapter-only probe measured 33,032 bytes after shrinking, including 32,768 bytes of empty indices. There is a completed group to emit here, which distinguishes this from #26911. In a control build that released those empty indices and updated the accounting, the original input returned both groups at 16 KiB.

Could shrink_to_fit() release unused index capacity in the remaining states and subtract the released bytes from allocation_bytes? A skewed active-group case would complement the current distinct-key fixture.

[viirya]

Thanks, confirmed: your skewed input ([0, 1, 1, …, 1]) failed under 16 KiB asking for 27.1 KB. In a8b5fb3, GroupsAccumulatorAdapter::shrink_to_fit also releases the empty indices of the remaining states and subtracts them from allocation_bytes. It skips non-empty indices, which only a failed update leaves behind. New test: clustered_partial_stream::tests::active_group_releases_adapter_indices.

[sunchao]

[P2] Ordered FIRST_VALUE / LAST_VALUE also retain counted batch scratch

FirstLastGroupsAccumulator uses the default hook. Its take_state truncates extreme_of_each_group_buf.1, a BooleanBufferBuilder used to track each batch's winners. Truncation preserves capacity, and size() still charges that capacity. The bitmap reached 16 KiB in my test even though only one group remained.

I reproduced this with a fully clustered partial FIRST_VALUE(v ORDER BY v) GROUP BY k over one batch of 131,072 distinct sorted Int64 keys (v = k). The 16 KiB greedy pool refused another 11.1 KiB after the completed prefix was emitted. With 32 KiB it returned all 131,072 rows. The 8,192-key control also succeeded at 16 KiB. Rebuilding just the winner bitmap at the surviving group count let the original 131,072-key input return all rows at 16 KiB.

Could the native first/last accumulator release its batch scratch through this hook and cover the ordered case? The ORDER BY matters because it selects the native implementation.

[viirya]

Good find, thanks. I missed this in the sweep. BooleanBufferBuilder::truncate keeps the capacity, and size() counts it. Your 131,072-key input failed under 16 KiB asking for 11.1 KB.

In a8b5fb3, FirstLastGroupsAccumulator::shrink_to_fit rebuilds the scoreboard bitmap at the remaining group count, all false, and clears extreme_buf_dirty. Both users of the bitmap accept that state: the pre-ordered path expects it all false between batches, and get_filtered_extreme_of_each_group resets it when it starts. The other fields don't need it: extreme_of_each_group_buf.0 goes through take_needed, and is_sets is rebuilt by finish(). New test: clustered_partial_stream::tests::completed_groups_release_first_value_bitmap, using FIRST_VALUE(v ORDER BY v).

[sunchao]

Optional test cleanup. The new collect_under_limit(&aggregate, 64 * 1024) already constructs a migration-enabled context and collects the stream. Could this reuse it and assert early_emit_count > 0, as the existing hash-stream test does? The row count alone could also pass if a future skip-aggregation threshold sends the input straight through.

[viirya]

Done. It now uses collect_under_limit and asserts early_emit_count > 0.

[sunchao]

Thanks for adding both 16 KiB reproductions. At a8b5fb3be, the skewed ANY_VALUE case releases the surviving group's empty row indices and updates the accounting. The ordered FIRST_VALUE case releases the old winner bitmap while preserving the active group's result. The hash-stream test also now checks that early emission actually happened.

All 350 physical-plan aggregate tests and 44 first/last tests passed locally. I used Rust 1.99.0 and a compatible dependency lock because the local registry lacked the pinned smallvec version. The fork Rust workflow passed on this commit.

One separate follow-up for the preordered first/last fast path. After 131,072 groups, the empty touched_groups_buf can still hold 1 MiB while size() reports 296 bytes for one remaining group. That undercount predates this PR. Tracking its capacity in size() and releasing it alongside the bitmap would make the pool reflect the actual allocation. The two reported reservation failures are fixed.

[viirya]

Thanks @alamb @sunchao @comphead