Repository navigation
fix: reserve a hash aggregate's emit before materializing it, and copy collect_list runs without slicing - #26
Merged
Conversation
…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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:The hash aggregate materializes all groups at once (
EmitTo::All), andCollectListGroupsAccumulator::evaluatecopies 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.EmittingOnMemoryPressure) and the materialized output are recorded in the reservation with the infallibleresizeuntil handed over (were an ignoredtry_resizein partial, hash_stream.rs:803, and an error in final).collect_listlong runs: copied withMutableArrayData, structs field by field, string/binary capacities counted up front, instead ofconcatover one slice per run.How are these changes tested?
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.collect_list_copies_struct_runs_with_their_nulls(spark-expr).clippy --all-targets --workspace -D warningsclean.🤖 Generated with Claude Code