Skip to content

fix: reserve a hash aggregate's emit before materializing it, and copy collect_list runs without slicing - #26

Merged
msafonov merged 1 commit into
branch-1.1from
joom/1.1-agg-emit-accounting
Oct 18, 2026
Merged

msafonov merged 1 commit into
branch-1.1from
joom/1.1-agg-emit-accounting

Conversation

@msafonov

Copy link
Copy Markdown
Collaborator

Which issue does this PR close?

Closes #. (Joom fork, OOM diagnosis of night 10-17 -> 10-18, fix B)

Rationale for this change

dim_device_loader on em32 (stages 37/39: partial and final collect_list(struct of 16 fields)) runs allocated 12-13.8 GiB per executor with 3.9 GiB reserved. Measured locally with the accounting allocator (collect_list over a 16-field struct, 512 x 1024-row fresh input batches, 181 MiB, Comet's skip-partial off), joom-1.1-12:

case peak allocated peak reserved
partial, 65536 groups 273.9 MiB 136.4 MiB
final, 65536 groups (runs of 8) 412.4 MiB 131.3 MiB
partial, 262144 groups (scattered) 284.7 MiB 143.2 MiB

The hash aggregate materializes all groups at once (EmitTo::All), and CollectListGroupsAccumulator::evaluate copies everything it retained, outside any reservation. The final case additionally builds a slice per run of the 16-field struct (a few KB each) and grows string buffers by doubling.

What changes are included in this PR?

  • AggregateHashTable::emit_size() (vendored physical-plan): size of accumulators whose state/value types copy, of group keys converted back from rows, plus 24 B per group; 0 for primitive state with column keys, so common aggregates are unchanged. Partial and final streams reserve it together with the table: on refusal the partial emits early and the final spills, so materializing runs in granted memory.
  • The materialized states of an early emit (EmittingOnMemoryPressure) and the materialized output are recorded in the reservation with the infallible resize until handed over (were an ignored try_resize in partial, hash_stream.rs:803, and an error in final).
  • collect_list long runs: copied with MutableArrayData, structs field by field, string/binary capacities counted up front, instead of concat over one slice per run.

How are these changes tested?

  • New core/tests/aggregate_emit_memory.rs (own test binary, since the allocator balance is process-wide): partial and final collect_list must allocate at most 1.1x their peak reservation + 8 MiB, and a final aggregate with half the memory must spill and stay within the limit with the same rows and values. On the base it fails (partial allocated 287189924 with 143057884 reserved); with the fix: partial 273.9 vs 272.0 MiB, final 261.4 vs 261.6 MiB, final at half the pool 133.7 MiB allocated for a 130.8 MiB limit.
  • New collect_list_copies_struct_runs_with_their_nulls (spark-expr).
  • Local (aarch64 container): vendored physical-plan 1761 passed / 0 failed; workspace green (core 661 + new test, spark-expr 973, shuffle 190, jni-bridge 28, common 70); fmt and clippy --all-targets --workspace -D warnings clean.
  • Cluster validation (dim_device collect_list stages on the datasci-ab playground vs vanilla) follows.

🤖 Generated with Claude Code

…y collect_list runs without slicing

A hash aggregate materializes all its groups at once (EmitTo::All) for its output, an early
emit or a spill. Accumulators whose state or value is variable-size copy what they hold
then; collect_list copies every value it retained. None of that was reserved: measured with
the accounting allocator on collect_list(struct of 16 fields), a partial aggregate peaked at
2.0x its reservation and a final one at 3.1x, the extra for the final one being a slice per
run of the 16-field struct (a few KB each) and string buffers grown by doubling.
dim_device_loader's collect_list stages run 8-10 GiB per executor outside the pool on em32.

- AggregateHashTable::emit_size: the size of accumulators whose state or value types copy,
  of group keys converted back from rows, and 24 B per group for the gather indices; 0 for a
  table of primitive state and column keys. Streams reserve it with the table: a partial
  aggregate emits early and a final one spills when the pool refuses it.
- The materialized states of an early emit and the materialized output stay reserved until
  handed over; they are recorded with the infallible resize instead of an ignored try_resize
  (partial) or an error (final), since they are allocated already.
- collect_list copies long runs with MutableArrayData, a struct field by field, with string
  and binary capacities counted up front, instead of concatenating one slice per run.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@github-actions github-actions Bot added bug Something isn't working area:aggregation labels Oct 18, 2026
@msafonov
msafonov merged commit 981e26d into branch-1.1 Oct 18, 2026
36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:aggregation bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant