perf: buffer the NestedLoopJoin build side as coalesced chunks instead of one concat_batches allocation - #24820
ranflarion wants to merge 5 commits into
Conversation
…ne concat_batches allocation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24820 +/- ##
========================================
Coverage 81.90% 81.91%
========================================
Files 1133 1134 +1
Lines 424993 425504 +511
Branches 424993 425504 +511
========================================
+ Hits 348083 348541 +458
- Misses 56278 56306 +28
- Partials 20632 20657 +25 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for working on this. The chunked build-side representation looks like a good way to avoid the large concat_batches allocation while preserving the existing join semantics. I walked through the global-to-local row translation in probing, bitmap updates, unmatched-left emission, and spill/replay, and I did not find any blocking issues.
I have one non-blocking testing suggestion below. The implementation looks good to me otherwise.
| @@ -3440,6 +3534,102 @@ pub(crate) mod tests { | |||
| Arc::new(TestMemoryExec::update_cache(&source)) | |||
There was a problem hiding this comment.
Could we add an execution-level regression test where the left child emits multiple sub-target batches, so JoinLeftData actually contains more than one coalesced chunk? It would be useful to exercise both a match and an unmatched-left row across a chunk boundary.
The new tests do a good job covering buffer reuse and locate(), but the existing join fixtures appear to provide a single build-side RecordBatch. With the smaller batch-size variants, that batch takes the large-batch bypass, so those tests do not exercise the new chunk transition paths during probing and final unmatched-left emission.
I am marking this as non-blocking since the global/local index translation looks consistent on inspection, but an end-to-end test would give us good regression coverage for the main correctness-sensitive part of this change.
There was a problem hiding this comment.
Added join_across_build_chunk_boundaries, which runs every join type over a 36-row build side delivered two ways, as 1-row batches that the load coalesces into chunks of exactly batch_size (4 or 12), and as 7-row batches that take the bypass so chunk edges fall off the output batch size and the probe and emission ranges have to be clamped. Rows 11|12 match on both sides of a boundary in the 4- and 12-row layouts, 27|28 in the 4- and 7-row layouts, every other boundary has unmatched rows on both sides, and the right side arrives as 1-row batches so batch size 12 exercises the range probe path and 4 the single-row one. Each case asserts the chunk count through collect_left_input, then checks the output against the same rows delivered as one batch, which is a single chunk. I mutation-tested it: swapping the local index for the global one in the range probe fails 20 of 40 cases, dropping the chunk-end clamp in unmatched-left emission fails 10, and that second one only shows up with the 7-row delivery, since exactly-batch_size chunks never need the clamp, which is why both deliveries are in the matrix.
There was a problem hiding this comment.
Added join_across_build_chunk_boundaries, which runs every join type over a 36-row build side delivered two ways, as 1-row batches that the load coalesces into chunks of exactly batch_size (4 or 12), and as 7-row batches that take the bypass so chunk edges fall off the output batch size and the probe and emission ranges have to be clamped.
…d chunk boundaries
|
run benchmark nlj |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing nlj-chunked-build (a5bbd22) to 5e168c9 (merge-base) diff Run configurationrun benchmark nljResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing nlj-chunked-build (a5bbd22) to 5e168c9 (merge-base) diff Run configurationrun benchmark nljCPU Details (lscpu)Details
Resource Usagenlj — base (merge-base)
nlj — branch
File an issue against this benchmark runner |
2010YOUY01
left a comment
2010YOUY01
left a comment
There was a problem hiding this comment.
Thank you for working on it. Left some suggestions.
| /// Build-side data as bounded chunks, in input order. Kept as chunks rather than one | ||
| /// `concat_batches` result so buffering never needs input and output to coexist, and a | ||
| /// chunk that already arrived at target size is retained without being copied at all. | ||
| chunks: Vec<RecordBatch>, | ||
| /// Row index of the first row of each chunk, i.e. prefix sums over the chunk lengths. | ||
| /// The visited-left bitmap is indexed by these global row numbers. | ||
| row_offsets: Vec<usize>, | ||
| total_rows: usize, | ||
| /// Build-side schema, kept so an empty chunk list still knows its shape | ||
| schema: SchemaRef, |
There was a problem hiding this comment.
Can we wrap it in a module so that callers can use it as if it were a single concatenated batch?
I feel the current implementation has several leaks -- operator logic has to understand the internal physical representation of the chunked build-side data. Some of these could be avoided with such a design.
There was a problem hiding this comment.
the operator now asks JoinLeftData::range(start, max_len) for the next run of rows and gets back a ChunkRange with the batch to address, the local start, the global start for the bitmap and the length, already clamped at the chunk end, so the translation and the clamp live in one place and the probe and unmatched-left paths just consume ranges.
| let mut batches: Vec<RecordBatch> = Vec::new(); | ||
| let mut chunks: Vec<RecordBatch> = Vec::new(); | ||
| // Batches at or above half the target size pass through without being copied. | ||
| let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), target_batch_size) |
There was a problem hiding this comment.
This coalescer is not necessary, there is a hidden convention: each operator should promise to output batches coalesced to batch_size, so here we can assume input doesn't contain small batches.
There was a problem hiding this comment.
The convention does not hold at this operator's build side. coalesce_batches wraps only FilterExec, HashJoinExec and RepartitionExec, so aggregates, unions, memory tables and row-group tails arrive at whatever size produced them, and MemorySourceConfig forwards stored batches as they are.
There was a problem hiding this comment.
The convention does not hold at this operator's build side.
I think it should hold, but there might be some violations here and there, we should fix them instead.
The evidence for such convention is no operator is doing this input coalescing currently. I think it's a reasonable design, since we don't have to introduce complexity/execution overhead in all operators for such pre-processing, as long as all operator promise to produce batch-size coalesced output.
| Arc::new(TestMemoryExec::update_cache(&source)) | ||
| } | ||
|
|
||
| /// A build side that already arrives in target-sized batches is retained as-is: the chunks |
There was a problem hiding this comment.
Can we move all tests to end-to-end tests, like sqllogictests, or remove them?
The issue with those UTs asserting internal properties is that they have a lot of correct assertions, but it's very hard to figure out what the end goal of the test is, and maintaining them is very hard. Today's AI likes to over-generate them.
For those two tests, I can't easily understand their test goals.
There was a problem hiding this comment.
Dropped join_left_data_drops_zero_row_chunks, the execution-level join_across_build_chunk_boundaries covers the row translation and I mutation-tested it against the same faults. Kept build_side_chunks_reuse_the_input_buffers with its goal stated in one line: it asserts by buffer pointer that a build side arriving at target size is retained without a copy, which is the property this PR exists for and which no sqllogictest can observe since the results are identical either way.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Reviewed a5bbd22 against 5e168c9 with five independent review passes. I found one P2 memory-budget regression, described inline.
Validation: all 1,886 physical-plan library tests passed. An independent expected-row oracle passed 10,500 executions on each of base and head, including 3,900 executions that spilled per version. No incorrect results were found. The full workspace suite was not run.
| coalescer.push_batch(batch)?; | ||
| while let Some(chunk) = coalescer.next_completed_batch() { | ||
| spill_file.append_batch(&chunk)?; | ||
| } |
There was a problem hiding this comment.
[P2] Bound spill coalescing by memory
After reservation.free(), this loop keeps accumulating input until the coalescer reaches its row target or EOF. For wide strings arriving in small batches, this retains substantial unreserved build data and creates oversized spill batches that replay must accept over budget.
I reproduced this on base 5e168c9d0 and head a5bbd22a0 using a lazy source that emits 512 one-row Utf8 batches containing 64 KiB strings, a 256 KiB memory pool, the default 8192-row target, and an empty right input:
- Base: 0.51 MiB peak live allocations.
- Head: 64.18 MiB peak live allocations, with 32.17 MiB still live just before the left source returns EOF while the pool reports 0 bytes reserved.
These are tracked live Rust allocations relative to the pre-execution baseline, not RSS. The source retains no previous batches; both runs spill once and return zero rows. Setting the head's batch target to 1 restores approximately base memory usage, confirming that the amplification follows the new coalescing.
Please flush partial chunks according to a byte budget, or preserve the input-batch spill granularity, so spilling does not accumulate much more memory than the configured limit. Otherwise this workload can exhaust process memory despite spilling.
There was a problem hiding this comment.
Yes, reproduced. Fixed in cc8badf by finishing the coalescer's partial chunk when the load spills (its inputs were all reserved, so that copy is inside the budget the pool already granted) and writing the rest of the stream to the spill file batch by batch as it arrives, the way the base did.
… at input granularity, coalesce each replay pass, and move chunk range clamping into JoinLeftData::range
|
This is the main part I'd suggest implementing differently. I'm happy to take it on myself — WDYT? Idea: #24820 (comment) I agree this seems optional if we only look at this PR. I feel strongly about this structure because we're likely to implement the same idea repeatedly in HJ and piecewise merge join, given #23076. A reusable struct would help there, so we don't have to handle the same complexity inside each operator. |
Understood. I'm happy to let you implement this. Let me know if I can be of any assistance. |
|
@ranflarion |
…ack coordinator's chunk loader
There was a problem hiding this comment.
Thanks for working through the earlier review feedback. The chunk-boundary coverage looks much better now, and the spill draining issue has been addressed by processing input a batch at a time.
I found one remaining memory-budget issue in the replay path that I think needs to be fixed before merging. The new replay-side coalescing can temporarily hold both the reserved decoded input batches and the newly concatenated output batch at the same time, which can push the nested-loop join beyond its configured memory limit.
Thanks again for the updates.
| let merged_batch = concat_batches(&left_schema, &pending_batches)?; | ||
| let n_rows = merged_batch.num_rows(); | ||
| // Every batch of the pass is reserved above, so compacting it here stays within budget. | ||
| let chunks = coalesce_chunks(pending_batches, &left_schema, target_batch_size)?; |
There was a problem hiding this comment.
I think this can exceed the configured memory limit during replay.
load_one_chunk keeps reserving decoded input batches until try_grow fails, and those batches are still live when they are passed to coalesce_chunks. The coalescing step then allocates the concatenated output without acquiring an additional reservation for that copy.
For a chunk made up of many small or wide batches, this means we can temporarily hold roughly both the fully reserved inputs and another full copied output. That brings back the same kind of memory amplification the earlier spill handling was trying to avoid, just on the replay side.
Could we either preserve the replay batches at their existing granularity, or explicitly limit or reserve the coalescing allocation before materializing the combined batch?
There was a problem hiding this comment.
You're right. I made some updates and added relevant testing. @2010YOUY01 feel free to use this as a reference if the decision is to do a replacement PR.
| variant | peak live allocation |
|---|---|
| PR head (both copies) | 40.05 MiB |
| head with replay kept raw (transition copy only) | 40.02 MiB |
| fixed with replay coalesced again (replay copy only) | 13.74 MiB |
| fixed | 8.14 MiB |
…ed and replayed batches as they arrived so the memory-limited path never copies a pass
|
@kosiew Hi, I'm doing a replacement PR due to #24820 (comment) Happy to hear thoughts on this direction. If we can agree on it, I'll put together a PR in 1–2 days. |
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for continuing to work through the memory accounting issues here. I checked the earlier review concerns against the current branch. The chunk-boundary regression coverage is now present, and the replay-side concern is addressed by keeping replay batches raw in load_one_chunk. The new spill/replay allocator test also passes.
I found one remaining issue on the non-spilling path, which I left inline. coalesce_chunks can allocate a near-full target-sized copied chunk while all of the fully reserved input batches are still live. That means a build side can fit within the memory pool but still exceed the configured budget during coalescing.
I think we should fix this before merging. If @2010YOUY01 can deliver the reusable chunk abstraction discussed earlier promptly, I would prefer replacing this PR with that implementation. The segmented-batch representation and range mechanics seem likely to be useful for hash join and piecewise merge join as well. I would keep NLJ-specific bitmap handling, probe coordination, cancellation, and spill policy local to NLJ.
If the replacement timing is uncertain, I think this PR should stay open long enough to address the in-memory accounting issue first. Any replacement should also retain the current chunk-boundary tests and allocator-peak coverage for spilling and replay, and add equivalent coverage for a near-limit in-memory build.
| let merged_batch = concat_batches(&schema, &batches)?; | ||
| // Compacted only once the whole side is reserved, so a load that spills never has a | ||
| // partially built chunk to materialize while the pool is exhausted. | ||
| let chunks = coalesce_chunks(batches, &schema, target_batch_size)?; |
There was a problem hiding this comment.
I think there is still a memory-budget issue here on the non-spilling path. coalesce_chunks runs while batches still owns all of the fully reserved input batches, but the allocation for the concatenated chunk is not separately reserved.
As a result, a build side that fits just below the pool limit can temporarily hold both the full input and a near-full target-sized copied chunk. Being fully reserved tells us that the input fits, but it does not account for the output allocation that exists at the same time during coalescing.
I reproduced this by varying the allocator test to use 126 one-row 64-KiB strings, a 9-MiB pool, and an empty right batch so the build stays in memory and there is no probe cross-product contribution. Peak live allocation reached 15.84 MiB, above the test's 13.50-MiB bound.
The committed spill/replay test passes, but it does not exercise this in-memory case. Could we either preserve the in-memory batches without copying here, or reserve and limit the output materialization before doing the copy? We should also make sure the ownership accounting remains exact after the input batches are dropped.
The replacement PR is now open: |
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Re-reviewed 1659313 against c5257f0 with five independent passes. The earlier spill-coalescing memory regression is fixed: its reproducer now peaks at approximately 0.51 MiB on both base and head.
I found two P2 issues, detailed inline: an empty-build spill panic and an allocator test that includes backtrace initialization in its measurement.
Validation: 2,251 physical-plan library tests and the allocator test passed under default local settings. An independent expected-row oracle passed 24,000 executions per revision, including 15,000 confirmed spilling executions per revision. The targeted empty-build reproducer fails only on head, and controlled backtrace runs reproduce the allocator-test failure. The full workspace suite was not run.
| for batch in batches { | ||
| spill_file.append_batch(&batch)?; | ||
| } |
There was a problem hiding this comment.
[P2] Preserve empty-batch filtering during spill
Removing the num_rows() > 0 guard allows an entirely empty build side to create a spill file. append_batch initializes the writer even for a zero-row batch, so finish() returns a file; replay then skips every batch and returns no chunk. This bypasses the assignment to active.left_schema, and final right-side emission panics with left_schema must be set.
I reproduced this against base c5257f054 and head 16593136d using three empty Int32 build batches, right-side row 7, and a 97-byte memory pool (each empty batch accounts for 96 bytes). For a Right join, base returns (NULL, 7) while head panics. Full, RightSemi, RightAnti, and RightMark also panic on head while base succeeds. The same Right-join failure reproduces with empty slices retaining 4 MiB buffers and a roughly 4 MiB pool, so it also occurs beyond tiny test budgets.
Please retain the previous empty-batch guard here and add coverage for an empty build side that exhausts its reservation after buffering an empty batch.
| let baseline = LIVE_BYTES.load(Ordering::Relaxed); | ||
| PEAK_BYTES.store(baseline, Ordering::Relaxed); | ||
| let output = collect(join.execute(0, ctx)?).await?; | ||
| let peak = PEAK_BYTES.load(Ordering::Relaxed) - baseline; |
There was a problem hiding this comment.
[P2] Initialize backtrace diagnostics before the allocation baseline
The baseline is taken before the first reservation failure. With DataFusion backtraces enabled and debug information retained, that failure initializes process-wide symbolization caches, which this test counts against the join's 12 MiB threshold. The new test therefore fails even when build-side spilling stays bounded; the AMD64 CI job reports 14.21 MiB against 12 MiB.
In a controlled copy of this test linked to the current head, the measured peak is 8.14 MiB with RUST_BACKTRACE=0 and 13.75 MiB with RUST_BACKTRACE=1. Calling and dropping DataFusionError::get_back_trace() before recording the baseline initializes about 5.61 MiB of retained diagnostic caches; the measured peak then returns to 8.14 MiB with backtraces still enabled and the same join execution. Retaining line-table debug information gives the same failure/pass distinction.
Please warm the backtrace machinery before taking the baseline, or isolate the measurement in a process with backtraces disabled. The existing 12 MiB threshold can remain unchanged.
6 participants
Which issue does this PR close?
Rationale for this change
NestedLoopJoinExecbuffered the whole build side into one batch viaconcat_batches, which doubles peak memory while the copy runs (inputs and output coexist), and the concat output is never reserved in the memory pool, so the doubled peak is invisible to it. On an ~880 MB build side (repro in #24819, runnable with stockdatafusion-cli), main peaks at 1738 MB RSS; with--memory-limit 1git completes while peaking at 1739 MB, 1.7x its own limit. A single allocation also caps any one string column ati32::MAXbytes (the overflow reported in #23032) and can only be spilled or released wholesale.This PR keeps the build side as target-batch-size chunks instead. Same repro after the change: 899 MB peak (1.02x the build side), 900 MB under
-m 1g. Wall time on the repro improves ~20% since the concat copy is gone.What changes are included in this PR?
JoinLeftDataholdsVec<RecordBatch>chunks with prefix-sumrow_offsets, a binary-searchlocate(global_row) -> (chunk, local_row), and the build schema (so an empty build side keeps its shape). Zero-row chunks are dropped at construction.collect_left_inputfeeds arrow'sBatchCoalescerwithwith_biggest_coalesce_batch_size(target/2): small batches are compacted to target size, batches at or above half the target pass through zero-copy. Reservation still charges each input batch'sget_array_memory_sizeas before.concat_batchesin the memory-limited rebuild is gone too.take/slice use chunk-local indices, and the visited-left bitmap keeps global row numbers, so bitmap semantics (including the multi-partition rules from fix: refuse memory-limited NestedLoopJoin fallback for left-emission joins with a multi-partition probe side #24675 and the deferred emission from fix: emit deferred unmatched rows when memory-limited NestedLoopJoin exhausts its left side #24746) are unchanged. The outputBatchCoalescerre-coalesces the occasionally smaller batch emitted at a chunk tail.One behavior note: a sliced batch at or above half the target size is now retained as-is, keeping its parent allocation alive, where the old concat incidentally un-pinned it by copying. The reservation charges the full parent buffers, so the pool over-counts rather than under-counts in that case; slice-aware accounting (dedup by allocation) is a planned follow-up.
Are these changes tested?
Existing coverage: the full NLJ suite including the memory-limited matrix and the one-shot re-execution test,
physical-planlib tests, the join fuzz suite (--features extended_tests), thememory_limitintegration tests, and the join sqllogictests all pass. New unit tests: chunks retain the input buffers by pointer identity (the zero-copy bypass), and zero-row chunks are dropped withlocate()boundary checks.Probe-throughput parity, medians over interleaved runs of ~1.6e10 pair evaluations, release builds, same machine:
Are there any user-facing changes?
No. Plans, results, and metrics are unchanged; only the build side's in-memory layout and its peak memory differ.