Account for transformer-retained batches in symmetric hash join - #24364
Open
kosiew wants to merge 25 commits into
Open
Account for transformer-retained batches in symmetric hash join#24364kosiew wants to merge 25 commits into
kosiew wants to merge 25 commits into
Conversation
…d batches, deduplicate shared buffers, and update reservation lifecycle** - RecordBatchMemoryCounter now tracks transformer retained batches. - Symmetric stream deduplicates input/transformer shared buffers. - Reservation updates occur on retain and release operations. - Added tests for noop, splitter, shared buffers, and reservation lifecycle.
…te helper (no public API/behavior change)
* Preserve RecordBatch array-object/non-buffer overhead while deduplicating Arrow buffer capacities in memory reservations. * Ensure transformer-held batches are fully accounted for in bounded-memory enforcement. * Add wide/nested retained-batch regression coverage for non-buffer allocation overhead. * Add execution-level bounded-memory tests using SymmetricHashJoinExec with a tight GreedyMemoryPool. * Verify NoopBatchTransformer and BatchSplitter retained batches correctly affect reservations and memory-limit behavior.
…age for children (Struct, List, Map, Union, Dict, RunEnd), plus regression test for shared nested StructArray child
…pBatchTransformer and BatchSplitter
…p contract; add regression-boundary rationale comment for 2_400 test limit
…n, Dictionary, RunEnd and improve shared allocation delta checks
…ArrayRef clones - Introduces a private assertion helper for testing. - Removes unnecessary ArrayRef clones after construction, improving efficiency.
github-actions
Bot
added
common
…n and memory usage - Polls transformer output to exhaustion. - Asserts reservation + stream_memory_usage grow on retain, return baseline on release. - Covers Noop + BatchSplitter.
kosiew
force-pushed
the
memcalc-03-23393
branch
from
August 14, 2026 12:26
ec4a627 to
792fa41
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24364 +/- ##
==========================================
+ Coverage 81.26% 81.46% +0.20%
==========================================
Files 1112 1119 +7
Lines 391578 400942 +9364
Branches 391578 400942 +9364
==========================================
+ Hits 318197 326631 +8434
- Misses 54677 55220 +543
- Partials 18704 19091 +387 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…tighten reserves - Replace fixed 2_400 pool limit with dynamic peak profiling using PeakRecordingPool. - Tighten pool reserves to stay below transformer batch delta. - Assert join values and introduce typed ResourcesExhausted. - Ensure coverage of Noop and BatchSplitter.
…mment - Updated pool limit calculation to be observed peak - 1. - Added comment clarifying direct vs integration ownership. - All focused tests pass, and formatting/diff checks succeed.
…cludes shared objects
…unded‑pool regression, and cover Noop + BatchSplitter; ablation confirms count_memory is required (transformer batch exceeds by 1 byte)**
- Added real SymmetricHashJoinExec poll-lifecycle regression.
- Verifies retain/release reservation calls:
- Noop: retain → immediate release.
- BatchSplitter: retained after non-final slice → released final
slice.
- Temporary ablation removing production resize calls: both cases fail.
…emory reservation accounting, add finite‑pool regression test, pin retain event by poll lifecycle position, restore count_memory handling
…deltas test helper, skip risky cache/macro ideas
… clone - Ensures change_deltas() uses a single mutex lock for thread safety. - Removes unnecessary cloning of data within the function.
- Removed duplicate direct‑stream, transformer, nested, OOM/calibration tests. - Simplified recording pool using `Vec<isize>` and removed change structs/snapshots. - Kept a real poll‑lifecycle regression test (Noop + splitter retain/release).
… bounded stream reservation test; add retain-transition test - Restored direct stream alias tests for Noop + Splitter and shared batch + nested shared child. - Restored exact bounded stream reservation test. - Added bounded SymmetricHashJoinExec retain-transition test that validates independent retained‑batch delta; test fails when `count_memory` is removed.
kosiew
force-pushed
the
memcalc-03-23393
branch
from
August 18, 2026 05:23
c25f101 to
00f20ed
Compare
…d-stream terminal-poll regression test
kosiew
marked this pull request as ready for review
August 18, 2026 07:35
Contributor
|
Thanks @kosiew I'm checking this today |
comphead
reviewed
Aug 19, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Thanks @kosiew doing first round of LLM review
P1 — Array-object dedup silently fails for Union / RunEndEncoded
UnionArray::new_unchecked (arrow-array/src/array/union_array.rs:164) and RunArray::try_new (arrow-array/src/array/run_array.rs:118) round-trip children through Array::into_data() → make_array(),
minting new Arc<dyn Array> allocations. Arc::as_ptr identity therefore never matches across parents built from the same Arc::clone(&shared_child). The PR's own tests concede this by using
get_buffer_memory_size (not get_array_memory_size) as the deduplicated amount for these two types — while the docstring on count_batch_with_array_overhead claims "counts recursively reachable Arrow
array objects once, including objects shared by multiple batches".
Fix: narrow the docstring to state that Union / RunArray children are dedup'd at buffer granularity only; or drop the Union / RunEndEncoded branches from array_children since they add code that
does nothing.
P2 — Fragile subtract-input_buffer in size()
size += self.left.size() - self.left.input_buffer.get_array_memory_size();Cross-file arithmetic on OneSideHashJoiner::size()'s implementation detail (symmetric_hash_join.rs:1448). A future edit to how size() counts input_buffer breaks this quietly.
Fix: extract OneSideHashJoiner::size_without_input_buffer(&self) -> usize.
P2 — Metric set before try_resize succeeds
update_reservation writes stream_memory_usage.set(capacity) before reservation.try_resize(capacity). On ResourcesExhausted, the metric reflects a value that never actually reserved. Move .set(...)
after the ?.
P2 — Test suite has ~200 LoC of removable redundancy
Three overlapping patterns:
- 5 nearly identical recursive-child tests (
..._shared_list_child,..._shared_map_children,..._shared_union_child,..._shared_dictionary_child,..._shared_run_end_encoded_child) — collapse
into one table-driven test. Missing coverage forFixedSizeList,ListView,LargeListViewbecomes trivial to add. RecordingMemoryPool(~60 LoC) + two overlapping integration tests (symmetric_hash_join_updates_reservation_while_transforming_outputand
symmetric_hash_join_transformer_retention_exhausts_bounded_pool) — merge into one test; consider using existingTrackConsumersPool(datafusion/execution/src/memory_pool/pool.rs:405) or asserting on the
stream_memory_usagemetric.stream_accounts_for_transformer_batches_oncevsstream_deduplicates_nested_transformer_batches— the second subsumes the first.
…thout_input_buffer, update metrics after successful resize, add failed‑resize regression assertion - Removes Union/RunEnd object‑identity branches/tests. - Adds size_without_input_buffer(). - Updates metrics only after successful resize. - Adds failed‑resize metric regression assertion.
Contributor
Author
|
Thanks for the review.
The direct stream test covers reservation growth, release, and a direct input-buffer alias; the nested test covers recursive child accounting. Likewise, the recording-pool test verifies retain/release transitions, while the bounded-pool test verifies the externally observable failure. I think it's a good idea keep the remaining focused tests separate because they assert distinct contracts. |
Contributor
|
Looks like |
comphead
reviewed
Aug 26, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Thanks @kosiew
I made a second pass, no blockers, 2 meds to address
1. Medium — eager free() under-reports still-resident memory
BothExhausted{final_result:true} calls free() + set(0) before drop, but left.input_buffer, right.input_buffer, and both PruningJoinHashMaps stay resident until the stream drops. This is the
under-count direction (pool told bytes are free while occupied). A parent that retains the exhausted stream lets another consumer reserve those bytes while live → real over-subscription.
Fix: rely on Drop (revert the free()), or clear the retained state (empty input_buffers, clear hash maps) before free() since final results are already emitted.
2. Medium — array_children omits Union and RunEndEncoded (dedup gap, over-count)
Both fall to vec![], so child object overhead is keyed only by the parent Arc and double-counts when a child is shared across parents (e.g. UnionArray::slice() clones field Arcs unsliced, so a
BatchSplitter slice shares them). Buffers are still deduped by the buffer counter, so only struct overhead inflates → over-count only (no OOM risk, but can trip spurious ResourcesExhausted).
Contradicts the fn doc ("...objects shared by multiple batches") and is untested (tests cover Struct/List/Map/Dictionary only).
Fix: add Union/RunEndEncoded arms mirroring count_array_memory_size (Union: as_union() → fields()/child(type_id); Run: dispatch on run_ends.data_type() → values()), or document + add a
regression test.
…ms, restore Union/RunEnd child traversal, and add slice‑sharing regression tests - Ensure that terminal streams retain their reservation and associated metric information until the stream is explicitly dropped, preventing premature release or loss of accounting data. - Restore correct child traversal behavior for `Union` and `RunEnd` nodes, aligning the implementation with the documented identity guarantees and fixing the regression that caused incorrect traversal order. - Add comprehensive regression tests for slice‑sharing scenarios to catch future breakage and verify that shared slices maintain proper reservation and metric tracking.
- Terminal test now drives normal exhaustion/finalization. - Union + RunEnd use nonzero partial slices.
Contributor
Author
|
- Fixed the `iter_on_single_items` function in `datafusion/common/src/utils/memory.rs`. - Replaced the custom iterator logic with `std::iter::once(...)`.
Contributor
|
run benchmark tpch tcpds |
Contributor
|
run benchmark tpch tpcds |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tcpdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing memcalc-03-23393 (d922d3d) to 6eaca8b (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
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
4 participants
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?
sizefunctions #23393Rationale for this change
SymmetricHashJoinStreamdid not include theRecordBatchretained by itsBatchTransformerin the stream's memory reservation. As a result, the symmetric hash join could under-account memory while eitherNoopBatchTransformerorBatchSplitterheld an output batch.Transformer-held batches can also share Arrow buffers or nested array objects with other batches retained by the stream, so independently summing batch memory would double-count shared allocations.
This change makes the stream reservation reflect the memory retained at that boundary exactly once. Corrected accounting can increase reservation pressure and may cause a bounded memory pool to report resource exhaustion earlier, without changing join results.
What changes are included in this PR?
RecordBatchMemoryCounter::count_batch_with_array_overheadto count unique Arrow buffers and recursively reachable array-object allocations across a sequence of batches.BatchTransformer::count_memoryand implements it for bothNoopBatchTransformerandBatchSplitter.SymmetricHashJoinStream::size()to account for its input buffers and transformer-retained batch through one sharedRecordBatchMemoryCounter, avoiding duplicate accounting of shared allocations.OneSideHashJoineraccounting formulas otherwise unchanged.Are these changes tested?
Yes. The patch adds regression coverage for:
test_record_batch_memory_counter_array_overhead_shared_across_batchestest_record_batch_memory_counter_deduplicates_shared_nested_array_overheadtest_record_batch_memory_counter_deduplicates_shared_list_childtest_record_batch_memory_counter_deduplicates_shared_map_childrentest_record_batch_memory_counter_deduplicates_shared_union_childtest_record_batch_memory_counter_deduplicates_shared_dictionary_childtest_record_batch_memory_counter_deduplicates_shared_run_end_encoded_childstream_accounts_for_transformer_batches_oncesymmetric_hash_join_releases_reservation_when_completestream_deduplicates_nested_transformer_batchestransformer_reservation_exhausts_poolsymmetric_hash_join_updates_reservation_while_transforming_output, for bothNoopBatchTransformerandBatchSplittersymmetric_hash_join_transformer_retention_exhausts_bounded_pool, for bothNoopBatchTransformerandBatchSplitterThe patch shown does not establish that the full
cargo test -p datafusion-physical-plan, clippy, repository lint, or extended test suites have been run.Are there any user-facing changes?
There are no intended changes to join results or query semantics.
Memory accounting for symmetric hash joins is more accurate. In bounded-memory configurations, a join that previously under-accounted transformer-retained output may now reach the configured memory limit and report resource exhaustion earlier.
No public API breaking change is introduced by this patch.
LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated content has been manually reviewed.