fix: cancel the NestedLoopJoin coordinated fallback when a partition is dropped unfinished by viirya · Pull Request #25004 · apache/datafusion

viirya · 2026-09-07T00:27:56Z

Which issue does this PR close?

Closes NestedLoopJoin coordinated fallback hangs surviving partitions when one is dropped unfinished #25003.

Rationale for this change

In the coordinated memory-limited NestedLoopJoin fallback, dropping an unfinished output partition can strand the shared chunk and leave other partitions waiting indefinitely. This change cancels the coordinated execution synchronously, releases coordinator-owned resources, and wakes unfinished partitions and in-flight loaders so they report an execution error. Normally completed partitions do not cancel their peers, and chunks still held by streams remain memory-accounted until released.

The timing that hangs: chunk advancement needs every probing partition to report, since the last one to report becomes the emitter and releases the coordinator slot. A partition that vanishes never reports, so no emitter is elected and nothing releases the slot. A survivor then asks for the next chunk and finds Case 1 unmatched (the slot still holds the previous chunk) while Cases 2 and 3 both require current.is_none(), so it falls into the notified() wait in next_chunk with nothing left to wake it. Because the coordinator is owned by the exec rather than the streams, a plan that outlives them also keeps the slot's Arc<JoinLeftData>, and with it the chunk's reservation.

Two shapes reach this: a release future dropped while pending on the coordinator mutex, and — more broadly — a mid-probe cancellation where no emitter is ever elected, so release_chunk is never called at all. The second is why hardening the release future alone cannot fix it.

What changes are included in this PR?

The coordinator lock is now synchronous (parking_lot::Mutex). Every critical section already was: the one slow operation, load_one_chunk, runs after the guard is dropped. next_chunk is restructured to decide under the lock and act after releasing it, carrying a Decision out of the locked block. This is what lets cleanup finish without another poll or await — cancellation runs from Drop, which has neither.

Cancellation travels on its own broadcast with registered waiters. cancel() sets a permanent cancelled flag and signals a dedicated cancel_notify; it is the only signaller, so a wake from it always means a real cancellation. Waiters enable their Notified before reading the flag, so a cancellation landing between those steps is delivered rather than lost. Streams hold a registered cancellation_watcher across polls and poll it every iteration — a stream parked on its right input is not waiting on chunk progress, so a flag read alone would never reach it.

A loader reading outside the lock cannot publish into a cancelled coordinator. The publish path would otherwise reinstate the stream and reservation that cancel had just dropped. The load also races the cancellation signal, so a read parked on input does not have to finish before the cancellation is observed.

Dead machinery removed. With the release synchronous, chunk_release_in_flight, its poll sites, handle_releasing_final_chunk and NLJState::ReleasingFinalChunk had no remaining purpose, and the "release dropped while pending" shape disappears with them.

What is the testing strategy for this PR?

Fourteen tests in nested_loop_join.rs covering the cancellation timings: a dropped partition not hanging its peers; cancellation releasing coordinator-held memory; cancellation before watcher registration; chunk-progress traffic not resolving a cancellation watcher; cancellation during an in-flight load discarding its publish, and not waiting for a read that may never finish; wakeups for streams parked on build input, right input, and the global-right replay; an errored stream cancelling peers on drop; and normal completion of both partitions leaving the coordinator alone, returning all rows and freeing memory while the plan stays alive.

They are mutation-checked rather than merely passing: neutering Drop fails 7 of them, removing Pending from Drop's match fails the retention test, and weakening the stream-side check to a plain flag read fails the parked-stream tests.

Also run: joins (1168), memory_limit (38, including #24746's regressions), the nested_loop_join / joins / information_schema sqllogictest files (8/8), and cargo clippy -p datafusion-physical-plan --all-targets --all-features -- -D warnings.

Unrelated and pre-existing: cargo clippy -p datafusion-sqllogictest --all-targets --all-features -- -D warnings fails needless_pass_by_value at datafusion/sqllogictest/src/engines/conversion.rs:99 on clean main (3266eaa91) as well.

Are there any user-facing changes?

Dropping an unfinished partition cancels the coordinated NLJ execution, causing other unfinished partitions to report an execution error. Normally completed partitions do not cancel their peers. Chunks remain memory-accounted while any stream still references them.

Previously those partitions hung, so this replaces an indefinite wait with a reported failure. Reviewers may want to weigh in on that semantic specifically: the alternative would be letting survivors finish, which cannot produce a complete result once a partition is gone.

The existing opt-out behavior is unchanged: it disables coordinated fallback for multi-partition joins requiring final left-side emission. No API changes.

…ished The partitions of a coordinated memory-limited fallback are not independent: chunk advancement needs every one of them to report. A partition that goes away before finishing therefore never reports, so no emitter is elected, nothing releases the coordinator slot, and the survivors fall through to the `notified()` wait in `next_chunk` and hang. Because the coordinator is owned by the exec rather than the streams, it also keeps holding the chunk it published, so its reservation stays charged for as long as the plan is alive. Reproduced directly against the coordinator: one partition takes chunk 0 and disappears, and a second partition asking for chunk 1 never returns. Three changes, which only make sense together: 1. The coordinator lock becomes a synchronous `parking_lot::Mutex`. Every critical section here was already synchronous -- the one slow operation, `load_one_chunk`, runs after the guard is dropped -- so `next_chunk` is restructured to decide under the lock, release it, then act: a `Decision` enum carries `Serve` / `Finished` / `Load` / `Wait` / `Cancelled` out of the locked block. This is what lets cleanup complete synchronously instead of depending on a future that a dropped stream takes with it. 2. A `cancelled` flag plus `cancel()`, driven from `Drop for NestedLoopJoinStream`. A stream that reached `Done` finished its work and cancels nothing; any other state cancels the fallback, dropping `current`, `carryover`, the shared left stream and the coordinator's reservation, and waking the waiters so they return an error rather than blocking. The guard exists from stream construction, so it also covers a cancellation that happens before the stream ever takes a chunk. Chunks other partitions still hold stay accounted until they release them. 3. A loader that was reading outside the lock must not publish into a coordinator that `cancel` has already cleaned up; doing so would reinstate the stream and reservation it just dropped. The publish path now discards its result and reports the cancellation instead. This accepts a deliberate semantic: cancelling one unfinished partition cancels the coordinated execution. It can no longer produce a complete result, so failing the remaining partitions is more honest than leaving them hanging or letting them report success from partial input. Making the release synchronous also removes the machinery that existed only to drive it across polls: `chunk_release_in_flight`, its three poll sites, `handle_releasing_final_chunk`, the `NLJState::ReleasingFinalChunk` state, and the `cx` argument `handle_emit_left_unmatched` no longer needs. With no future to interrupt, the "release dropped while pending" shape disappears on its own. Tests cover the cancellation timings that matter: a dropped partition not hanging the survivors, cancel releasing what the coordinator held, cancellation while a load is in flight discarding its result, and a normally-completed stream not cancelling anything. Each fails if `cancel()` is neutered -- the hang test by timing out, the memory test with 505 bytes still reserved. Co-authored-by: Claude Code

Review of the previous commit found three places where cancellation was not actually observed, plus a hole in my own test coverage. Drop only handled `SpillState::Active`, so a stream cancelled while still `Pending` -- set up for the fallback but not yet holding a chunk -- reported nothing and cancelled nothing. That is exactly the "cancelled before taking a chunk" case the guard was supposed to cover, so the claim that placing it at stream construction was sufficient was wrong. Both states now cancel; only `Disabled` does not. Cancellation was only checked inside `next_chunk`, which a stream holding the final chunk never calls again. Such a survivor ran to completion and reported success built from an execution that had lost a partition. `poll_next` now checks on every iteration. `cancel()` could not interrupt a load already in flight: the loader awaits `load_one_chunk` directly, and waking `notify` does not reach a read parked on its input. The load now races the cancellation signal, and on cancellation drops its local stream and reservation, clears the leader claim and reports the cancellation rather than waiting for a read that may never finish. The publish check alone prevented reinstatement but not the stall. The test gap is the more important lesson. My four tests all called `coordinator.cancel()` directly, so neutering `cancel()` failed them while neutering the entire `Drop` body left all 68 tests passing -- they exercised the function but never the wiring that calls it. The four review tests added here drive real plans and streams, and two of them fail when `Drop` is neutered. Co-authored-by: Claude Code

…iters Follow-up review found two remaining hang paths, both traceable to using the chunk-progress `Notify` to carry cancellation. `Notify` broadcasts to waiters that already exist and stores no state, so a cancellation landing before a waiter registered was simply lost, and a task parked on something other than chunk progress had no waker registered at all. Cancellation now travels on its own `cancel_notify`, and every waiter registers with `Notified::enable()` before reading the `cancelled` flag. A cancellation arriving between those two steps is delivered rather than lost, which closes the gap at initial registration. Because `cancel_notify` is signalled only by `cancel()`, a wake from it always means a real cancellation: the re-arm loop that re-registered after unrelated wakes is gone, and with it the race that loop carried. For streams, checking a flag was not enough. A stream returning Pending from its right input is not waiting on the coordinator, so dropping a peer never woke it and the poll-loop check was not reached until some unrelated event happened to poll the task. Streams now hold a registered `cancellation_watcher` across polls and poll it each iteration, so the waker really is with the coordinator. Writing a replacement for the now-inapplicable re-arm test caught a third bug of my own: the loader's watcher was still constructed from the progress `Notify`, so any chunk-progress traffic resolved it and failed the load as if cancelled. The new test drives unrelated `notify_waiters()` past a parked loader and requires it to stay pending, then cancels for real. The obsolete re-arm test is dropped -- it drove a loop that no longer exists -- and the replacement needs no instrumentation. The hook for cancelling before watcher registration stays, since that gap still needs guarding. Verified: 8 review tests pass; joins 1162; memory_limit 38 including apache#24746's regressions; clippy clean. Neutering `Drop` fails 4 of the review tests, and weakening the stream check back to a plain flag read fails 1, so both mechanisms are covered rather than merely present. Co-authored-by: Claude Code

Round 3 found no blocking correctness issue: the reviewer walked the four notification interleavings and confirmed that enabling the `Notified` before reading `cancelled` covers each one, and that a single-signaller broadcast makes the removed re-arm loop unnecessary. What it did find was documentation and test names claiming more than the code delivers. Adds the reviewer's four tests: a watcher constructed before cancellation but first polled after it, plus one constructed after; several live watchers with waker replacement and progress traffic that must not wake them; a real `NestedLoopJoinStream` parked on build input receiving a cancellation wakeup when its peer is dropped; and both partitions running to normal completion and being dropped, which must leave the coordinator usable, return all nine rows, and free the memory while the plan stays alive. That last one closes the gap my own "completed stream" test never covered. Documentation corrections, all cases of promising more than happens: - `cancellation_watcher` said the future was already queued. Registration happens on its first poll, so callers have to poll it, not just hold it. - `cancel_notify` was described as notified "once, permanently". It is a broadcast carrying no state; `cancelled` is what persists. - `cancel`'s documentation had been split by a later insertion and was sitting above `is_cancelled`. Moved back, and it now records that `cancel` is the only signaller of `cancel_notify` -- the invariant the missing re-arm loop rests on. Two of my tests were renamed because their names overstated what they did. `test_nlj_cancel_during_load_discards_the_result` cancels *before* starting the load, so it exercises the entry check, not a publish after an in-flight read; it is now `test_nlj_cancelled_coordinator_refuses_to_serve_chunks` and points at the test that does pause a real read. `test_nlj_completed_stream_drop_does_not_cancel` never built or dropped a stream; it is now `test_nlj_uncancelled_coordinator_serves_and_stays_live` and points at the reviewer's test that drops real streams. Verified: joins 1166; memory_limit 38 including apache#24746's regressions; clippy clean. Neutering `Drop` fails 5 review tests and weakening the stream check to a plain flag read fails 2. Co-authored-by: Claude Code

Round 3 listed these as optional, sharing the already-reviewed watcher and Drop mechanisms. They are cheap and they close the two cases where I was reasoning rather than testing. A stream can park in `EmitGlobalRightUnmatched` rather than `BufferingLeft`: that state reopens the spilled right side and polls it. Since the watcher is polled at the top of every `poll_next` iteration regardless of state, it should be woken there too, and now that is asserted rather than assumed -- a FULL join parked on a pending replay records a wake when its peer is dropped. An unfinished stream that ends in an error also reaches `Drop` without passing through `Done`, so it cancels its peers. That follows from the `Done` check, but the error path had no test of its own; one now injects a failing right input and confirms the peers are cancelled. Both fail if `Drop` is neutered. Dropping `Pending` back out of `Drop`'s match -- the round-1 defect -- now fails 5 review tests rather than the 1 it did then. Verified: 14 review tests; joins 1168; memory_limit 38 including apache#24746's regressions; clippy clean. While adding these I removed a `#[tokio::test]` that belonged to `review_v3_cancel_wakes_pending_build_input` and restored it, then audited every test in the module for a missing attribute. Only the shared `review_paused_loader` helper lacks one, correctly. This is the second time an attribute has gone missing to careless editing here, and a silently unregistered test is worse than a failing one. Co-authored-by: Claude Code

Review round 4 found that neither fixture added in the previous commit exercised its stated scenario, and both were rewritten rather than patched. The error fixture built a stream with a failing right input and dropped it without polling, so it only repeated "an unstarted stream cancels its peers", which other tests already cover. Polling alone would not have helped: its `left_data` was a permanently pending future, so it could never leave `BufferingLeft` to reach the right input at all. It now resolves the build side through `LeftLoad::Spilled`, drives the stream under a bounded timeout until the injected error surfaces, and asserts that error and a non-`Done` state *before* dropping -- so a fixture that stops reaching the right input fails instead of quietly degrading. The replay fixture assigned `EmitGlobalRightUnmatched` on top of `SpillState::Pending`, a combination execution never produces. It reached the pending read only because `right_data` was already `Some`, bypassing the `Active`-only reopen branch, so it was really testing a `Pending` watcher while claiming to test the replay configuration. It now polls until the spilled build side puts it in `Active`, asserts that, and only then parks in the replay stage. The replay reader is still injected rather than reopened -- that shortcut is now stated in the doc comment -- and the test additionally requires the wake to surface the cancellation rather than merely counting a wakeup. Also corrects the rustdoc above `test_nlj_cancelled_coordinator_refuses_to_serve_chunks`, which still described cancelling mid-load and validating publication after the rename fixed the name and the body comment, and softens the `cancel_notify` wording: a delivered broadcast is itself sufficient, so it is not merely a nudge to go and read the flag. On test registration: `cargo test -- --list` reports all 14 review tests, which is the right check. The regex sweep I ran was the wrong tool -- the module has legitimate async helpers, so requiring an attribute on every `async fn` would be wrong -- and it is not committed. Verified: 14 review tests; joins 1168; memory_limit 38 including apache#24746's regressions; clippy clean. Neutering `Drop` fails both fixtures; removing the watcher poll fails the replay one only, which is the expected split since the error fixture does not depend on the watcher. Co-authored-by: Claude Code

Review round 5 reported no blocking findings and cleared the change for upstream, but pointed at three comments still describing the pre-refactor code. The coordinator's inner state was documented as guarded by an async mutex; it is a `parking_lot::Mutex`, and the reason matters enough to record: cancellation and chunk release have to finish inside one `poll_next` rather than depending on a future a dropped stream would take away, and no critical section awaits. The coordinator reservation was described as holding the current chunk's memory, registered via `initiate_fallback` -- a function that no longer exists. It only holds bytes while a load runs; `load_one_chunk` then moves them into the chunk's `JoinLeftData` with `take()`, so accounting follows the data. The comment above `buffered_left_data = None` claimed the `Arc` reaches zero once the last partition lets go. The slot holds a strong reference too, so the reservation is freed only when every holder *and* the slot release. That is precisely why the slot cannot hold the chunk weakly, which two earlier attempts here established the hard way, so the comment now says it. Also corrects an overreach in my own reporting: I had described the sqllogictest suite as blocked by the pre-existing `needless_pass_by_value` failure in the sqllogictest crate. That was wrong -- clippy escalates warnings with `-D warnings`, `cargo test` does not. Running it works: `cargo test -p datafusion-sqllogictest --test sqllogictests -- nested_loop_join joins information_schema` passes 8/8 here. The lint failure is real but separate, and reproduces on clean `upstream/main` (`3266eaa91`) with `cargo clippy -p datafusion-sqllogictest --all-targets --all-features -- -D warnings`. Verified: joins 1168; sqllogictest 8/8 on the join and information_schema files; clippy clean for datafusion-physical-plan. Co-authored-by: Claude Code

Round 6 cleared the change for upstream and corrected three pieces of wording. The important one: I had justified the slot's strong `Arc` by saying the reservation frees only after every holder and the slot release. True, but that is a consequence, not the reason. The reason is a handoff: a faster partition can drop its reference before a slower one has taken the chunk at all, and the slow one is served from the slot (`Decision::Serve`), so the slot has to keep the chunk alive across that gap. Stated that way, the comment explains why holding it weakly cannot work -- which two earlier attempts here established by failing. "Complete inside a single `poll_next`" was wrong about cancellation, which runs from `Drop` and has neither a poll nor an await; it now says "without another poll or await". "Only holds bytes while a load is in progress" was too absolute: the error path returns the reservation to the coordinator with its bytes still accounted, since `take()` happens only on success. The comment now describes the successful path without claiming the stronger invariant. Verified: joins 1168; clippy clean for datafusion-physical-plan. Comment-only, so memory_limit and sqllogictest were not rerun. Co-authored-by: Claude Code

Review-process naming and duplication that should not reach upstream. The fourteen tests carried `review_v2`/`v3`/`v4` prefixes recording which review round produced them, which means nothing to anyone reading the file later. They are renamed after the behaviour they check, under one `nlj_` prefix: what is woken, what is released, what refuses to serve. The two shared helpers are renamed the same way (`run_paused_loader_cancellation`, `cancellation_test_plan`). The coordinator carried a test-only `review_cancel_before_watch` field with no documentation at all -- a field named after a review round, sitting in production state. It is now `cancel_at_leader_claim` and says what it is for: firing a cancellation in the window between claiming a load and registering the watcher, which is the interleaving where a lost notification would strand every other partition, and which a test cannot otherwise reach. Two `eprintln!`s left over from debugging are gone; the assertions beside them already carried the meaning. The `WakeCount` waker was defined identically in four tests and is now one helper with `new`/`count`/`reset` -- deliberately just that, not a test framework. While doing the extraction a regex of mine rewrote the helper's own body into `self.reset()` calling itself, which the test run caught as a stack overflow rather than a wrong answer. Fixed, and a reminder that mechanical renames need the result read back. Verified: 44 NLJ tests, all registered per `cargo test -- --list`; joins 1168; clippy clean. Neutering `Drop` still fails 7 of them, and disabling the renamed test seam fails the test that depends on it, so the renames did not quietly detach any coverage. Co-authored-by: Claude Code

Round 7 found no blocking issue but caught five precision problems, one of them mine from the previous commit. Inserting the shared `WakeCount` put it directly beneath an existing rustdoc, so it captured documentation belonging to `multi_partition_memory_limited_join_collect_concurrent` -- the block explaining why that helper must collect partitions concurrently rather than sequentially. `WakeCount` now carries only its own docs and the explanation is back on the helper it describes. Same class of mistake as the recursive-body regex last commit: an insertion that looked local but moved something adjacent. `nlj_normal_completion_leaves_coordinator_usable` claimed more than it checks -- it establishes that normal completion does not cancel peers of the same execution, not that a plan or coordinator can be re-executed (it cannot; the coordinator's state is one-shot). Renamed to `nlj_normal_completion_does_not_cancel_peers`. `nlj_drop_while_pending_releases_chunk` named a chunk release that has not happened at that point: no chunk is loaded before the drop. It asserts the memory is not retained afterwards, so it is now `nlj_drop_pending_partition_does_not_retain_chunk_memory`. Removing `Pending` from `Drop`'s match still fails it, so the name change tracks the same coverage. The test seam's documentation said it reproduced an interleaving "a test cannot otherwise hit" and that a lost cancellation would "strand every other partition". Both overstated: it reproduces the interleaving deterministically, and what it strands is the loader itself, since other observers may already be returning errors. Verified: joins 1168; clippy clean. For the record, the nested_loop_join module holds 82 tests; the `nlj_` prefix selects 44 of them, so earlier reports of "44 NLJ tests" were describing the filter, not the module. Co-authored-by: Claude Code

codecov-commenter · 2026-09-07T00:52:56Z

Codecov Report

❌ Patch coverage is 89.35722% with 101 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.76%. Comparing base (3266eaa) to head (21817bb).
⚠️ Report is 36 commits behind head on main.

Files with missing lines	Patch %	Lines
...fusion/physical-plan/src/joins/nested_loop_join.rs	89.35%	22 Missing and 79 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25004      +/-   ##
==========================================
+ Coverage   81.64%   81.76%   +0.12%     
==========================================
  Files        1124     1128       +4     
  Lines      413173   417431    +4258     
  Branches   413173   417431    +4258     
==========================================
+ Hits       337320   341302    +3982     
+ Misses      55995    55991       -4     
- Partials    19858    20138     +280

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

kosiew

@viirya,

Thanks for working on this. The synchronous chunk release and cancellation cleanup look like good improvements for the coordinated NLJ fallback. I found one issue around SpillState::Pending that I think needs to be addressed before merging. I also left one non-blocking test coverage suggestion.

kosiew · 2026-09-08T08:57:07Z

+                        SpillState::Active(active) => {
+                            Some(active.coordinator.cancellation_watcher())
+                        }
+                        SpillState::Pending {


I think we need to be a bit more careful about cancellation while the stream is SpillState::Pending.

Pending only means spilling is possible. Every eligible execution enters this state before the shared left load determines whether the input actually needs to spill. If collect_left_input resolves to LeftLoad::InMemory, there is no coordinated fallback or shared chunk counter, so the right partitions are still independent.

With the current behavior, dropping an unfinished partition while it is still Pending can cancel the coordinator and cause a surviving partition to fail even though the left side ultimately fits in memory.

Could we defer or separately record the pending drop, and only apply the cancellation if the shared load resolves to Spilled?

I don't think limiting cancellation to Active streams is sufficient either. One partition may already have transitioned to Active while another is still Pending. If that pending partition is dropped, the active peer is already relying on coordinated fallback and needs to be cancelled.

It would also be good to add a regression test with a sufficiently large memory pool where one pending partition is dropped and another partition can still be collected successfully.

kosiew · 2026-09-08T08:57:07Z

+    }
+
+    #[tokio::test]
+    async fn nlj_stream_stops_producing_after_cancellation_error() -> Result<()> {


Small non-blocking test suggestion: this uses the shared fixture with batch_size = 1, so it can't exercise cancellation while the coalescer has a partial batch buffered.

It might be worth adding a focused test with a larger batch size, cancelling while the survivor has buffered output, and checking that the terminal behavior is still an error followed by None. This would give us a little more confidence that buffered output can't escape after cancellation.

A stream dropped while its execution was still `SpillState::Pending` cancelled the coordinator outright. `Pending` is entered by every eligible execution before the shared load decides whether the left side spills, so a join whose left side fits in memory -- one that never coordinates at all -- was failed by a partition going away. Record the drop instead and let the coordinator decide. `pending_drop` and `coordination_started` are updated and read under the same lock, so whichever of the two happens second performs the cancellation: * drop, then coordinate -- `begin_coordination` sees the recorded drop. * coordinate, then drop -- `record_pending_drop` sees coordination started and cancels immediately, because that peer is already waiting on a probe report the departing partition will never make. The earlier fix only covered the first order, since nothing remembered that coordination had begun; a plan reaching `Active` before the drop was left hanging with the chunk still reserved. An execution that resolves to `InMemory` never calls `begin_coordination`, so the recorded drop stays inert -- which is the point. Also stop a cancelled stream from emitting after its error. `handle_done` pads an empty result with one empty batch so an all-filtered join keeps its schema; a survivor cancelled before producing rows took that path and returned `Some(Ok(<empty>))` after the failure. `cancelled_terminally` ends the stream at `None` instead. Two of the watcher tests now drop a peer that has really reached `Active` rather than calling `cancel()` by hand, so they exercise the `Drop` wiring. `nlj_cancel_wakes_stream_parked_on_build_input` keeps its explicit `begin_coordination()` -- its `left_data` never resolves, so no peer there can coordinate -- and asserts the drop stays inert first. Comments and two test names that the earlier mechanism left stale are corrected in passing: the removed `pending_drops`/`apply_pending_drops` are no longer referenced, the `Drop` comment no longer claims the `Active` peer cancels itself, and neither test name now promises a premise it does not set up. Co-authored-by: Claude Code

viirya · 2026-09-08T19:08:23Z

Thanks — you were right on both counts, and the Pending one was a real regression I introduced.

The Pending drop. Pending is entered by every eligible execution before the shared load decides whether the left side actually spills, so cancelling on a drop there failed joins that never coordinate at all. I reproduced it with an ample pool: an execution whose left side fits in memory got cancelled by an unrelated partition going away.

The drop is now recorded rather than acted on, and the coordinator decides. pending_drop and coordination_started are written and read under the same lock, so whichever of the two happens second performs the cancellation:

drop, then coordinate → begin_coordination picks up the recorded drop;
coordinate, then drop → record_pending_drop sees coordination has begun and cancels immediately, since that peer is already waiting on a probe report the departing partition will never make.

My first attempt at this only handled the first order — nothing remembered that coordination had begun, so a plan that reached Active before the drop was still left hanging with the chunk reserved. A reviewer caught that the test I had written exercised the opposite order, which is why it passed. There is now a real-plan test for each direction, and I checked that removing the coordination_started read fails the one covering the order I had missed.

An execution that resolves to InMemory never calls begin_coordination, so a recorded drop stays inert there — which is the behaviour you were asking for.

One more thing that came out of it. handle_done pads an empty result with a single empty batch so an all-filtered join keeps its schema. A survivor cancelled before producing any rows took that path, so the stream returned Some(Ok(<empty>)) after its cancellation error instead of ending. cancelled_terminally now short-circuits that, so a failure termination ends at None.

Two of the watcher tests were also calling cancel() by hand instead of dropping a peer that had really reached Active; they now poll the peer into Active, assert it, and drop it, so they cover the actual Drop wiring. nlj_cancel_wakes_stream_parked_on_build_input keeps an explicit call, because its left_data never resolves and no peer there can coordinate — it asserts the drop stays inert first, and the comment says why the seam is there.

joins 1174 pass, the memory-limit suite 38, clippy clean.

kosiew

@viirya, thanks for the follow-up. The blocking Pending-drop issue looks addressed. Pending drops now stay inert when the shared left load resolves to LeftLoad::InMemory, while both drop/coordination orderings correctly cancel once coordinated fallback is involved.

The buffered-output suggestion is also covered by nlj_cancellation_after_buffered_rows_ends_without_output, which establishes a partial coalescer buffer and verifies that cancellation produces an error followed by None.

I didn't find any new correctness issues in the follow-up changes.

One small coverage-description note: the drop-then-coordinate test directly exercises FallbackCoordinator::begin_coordination, while the coordinate-then-drop case goes through a real plan. This is just a description nuance and not something that needs another change.

Validation: cargo test -p datafusion-physical-plan nlj_ --lib -- --nocapture (50 passed).

viirya · 2026-09-09T07:41:02Z

Thanks for the re-review, and you're right about the coverage description — that imprecision was in my comment, not the code. I wrote "a real-plan test for each direction," but only the coordinate-then-drop case goes through a real plan (nlj_pending_drop_cancels_active_peer_real_plan); the drop-then-coordinate case is a unit test calling FallbackCoordinator::begin_coordination directly on a bare coordinator.

For what it's worth, the test's own doc comment already draws that line — it says the opposite order "runs on a real plan in nlj_pending_drop_cancels_active_peer_real_plan" — so the source doesn't carry the overstatement. I'll keep the split as it is: the direct call is what lets that test pin the deferral without a plan that has to spill, and the ordering that actually regressed is the one covered end to end.

viirya · 2026-09-09T18:29:15Z

Note for anyone else hitting this: the merge queue is currently failing on build and run with wasm-pack, but not for any reason in this PR.

The job dies in the Install dependencies step, before anything is compiled:

E: Failed to fetch https://dl.google.com/linux/chrome-stable/deb/dists/stable/main/binary-amd64/Packages.gz
   Hash Sum mismatch
   Last modification reported: Wed, 09 Sep 2026 09:41:12 +0000
   Release file created at:    Wed, 09 Sep 2026 17:16:59 +0000

Google's apt repository is serving a Packages.gz that doesn't match its own Release file, so apt-get refuses the index and the Chrome install exits 100. The wasm tests themselves are then skipped.

It isn't specific to this branch — the same job is failing in the queued runs for #24976, #25028, #25032 and #25073 as well, and this commit passed the same check on the PR head. Retrying just re-runs into the same upstream mirror, so I'll wait for it to resolve rather than keep re-queuing.

viirya · 2026-09-10T03:06:16Z

Thanks @kosiew for reviewing!

@viirya

…d drop the cancel protocol (apache#25542) ## Which issue does this PR close? - Part of apache#15760. Follows apache#22038 and apache#25004. ## Rationale for this change Thanks @viirya for apache#22038 and apache#25004, which made `LEFT` / `FULL` / `LEFT SEMI` / `LEFT ANTI` / `LEFT MARK` nested loop joins spill with a multi-partition right side instead of failing. This PR keeps that capability, keeps the way left chunks are scheduled, and changes how matches are tracked. Two user-visible problems today: - Dropping one unfinished partition of a spilling join fails the other partitions with an execution error (the behaviour apache#25004 chose over hanging). On the in-memory path, and in `HashJoinExec`, the others simply finish. apache#25004 asked reviewers to weigh in on "letting survivors finish"; nobody did, so I am picking that question up here. - About 1,270 lines of `nested_loop_join.rs` (roughly 300 of production code, the rest tests and test seams) exist only to make that early drop safe: `cancel_notify`, registered watchers polled on every `poll_next`, `pending_drop` / `coordination_started`, a `select!` racing each chunk load against cancellation, a publish guard, and `cancelled_terminally`. Both come from one coupling. On the in-memory path the shared probe counter only decides who emits the unmatched left rows. In the fallback the same per-chunk counter also decides when the next chunk may load, and the partition it elects holds the only `release_chunk` call while it emits. A partition that goes away can therefore strand the chunk, and the loader going away takes the shared left stream with it. ## What changes are included in this PR? Three commits, meant to be reviewed in order. It will be squashed on merge; the split is only there to make the last one easier to read. ### 1. `test:` coverage that does not depend on this change These pass on `main` as well, so they show the same behaviour before and after. - `test_nested_loop_join_spill_fuzz`: all 10 join types, right side spread over 4 partitions, comparing a run under a 4 KB limit with the same join without a limit, and asserting the limited run spilled. Until now no fuzz test reached the NLJ fallback. - `nested_loop_join_spill.slt`: multi-partition `LEFT MARK` and `LIMIT` cases. - `memory_limit/nlj_spill_unmatched.rs`: assert `spill_count > 0` exactly when a memory limit is set. These three regressions would otherwise keep passing if the join stopped spilling. ### 2. `chore:` small cleanups, no behaviour change - `execute()` computed the "can this join spill" condition twice, with the same long comment; compute it once. - Remove `left_buffered_in_one_pass`, which is written in three places and never read. - `process_left_unmatched_range` copied the visited bitmap bit by bit under an `assert!` whose message is `"DBG: .."`; use `BooleanBuffer::collect_bool`. - The module docs still said each output partition re-executes the left child; since apache#24677 it is executed once and spilled during that load. ### 3. `refactor:` separate match tracking from chunk scheduling **The schedule does not change.** All partitions still share one left chunk at a time, as large as memory allows, and the next chunk loads once every partition has finished the current one. That is what keeps right-side replays at a minimum, and it is the part of the current design worth keeping (see the measurements below). **Match tracking moves to one bitmap over the whole spilled left side, with one final emission pass.** This is the mirror image of what apache#24746 did for the right side (`global_right_bitmaps` + `EmitGlobalRightUnmatched`), and it is what other engines do: DuckDB's `PhysicalNestedLoopJoin` keeps one `OuterJoinMarker` over the whole build side and emits unmatched build rows in a separate final scan; Spark's `BroadcastNestedLoopJoinExec` ORs per-task bitsets into one and emits as a final step. - `LeftSpillData`, already shared by all partitions through `OnceAsync<LeftLoad>`, gets the bitmap (one bit per spilled left row, addressed by the row's position in the spill file) and the probe-threads counter — the same two things `JoinLeftData` carries on the in-memory path. - While probing a chunk each partition uses a private bitmap, so the probe path takes no contended lock, and merges it into the global one when it finishes the chunk. - `ProbeEnd` is entered once per stream, after its last chunk. The last stream to report streams the left spill file once more and emits the final left rows batch by batch. With no per-chunk emitter, what is left of the coordinator is a count, `LeftChunkBarrier`: `finished == live` means move on to the next chunk. - The chunk load is a shared future, the same pattern as `OnceFut` for the whole left side. Whichever waiting partition is polled drives it and an error reaches all of them, so it no longer matters which partition started the load or whether it is still around. `Decision`, `loader_in_flight`, `abandon_load` and the `select!` go away. - A stream dropped before `Done` departs: `live -= 1` (and `finished -= 1` if it had already finished the current chunk). The others carry on. Because the departed stream never reports probe completion, the final left rows are not emitted — exactly what happens on the in-memory path. - The barrier's clone of the load keeps the chunk alive for partitions that reach it late and lets go of it when the barrier moves on, so the chunk's reservation follows the data, as asked for in apache#22038. `nested_loop_join.rs` goes from 7,547 to 6,226 lines: production code −262, tests −1,059 (24 cancellation tests and their helpers are replaced by three drop scenarios). ### What I tried first, and why this PR looks different from its first version The first version of this PR also gave every partition its own left chunk, so that partitions never wait for each other. It was simpler still, but it costs what the shared chunk saves: memory split P ways means chunks P times smaller and P times more right-side replays. On a wide-row workload (below) `main` makes 7 passes over the right side and that version made 98, for about 1.6× the time and 1.5× the RSS. A variant sharing chunks through futures without a barrier reached time parity only at one hand-tuned chunk size, and never `main`'s memory. A partition that runs ahead has to either wait or hold a second chunk; I kept the wait. ## What is the testing strategy for this PR? - The tests in the first commit pass on `main` and after this change. - The multi-partition unit tests now use a multi-chunk left side for all five left-emitting join types (before, only `LEFT` and `FULL` had multi-chunk coverage). - Three drop scenarios, each asserting that the remaining partition finishes within a timeout, without error, and that no memory stays reserved: dropped while holding a chunk (`..._dropped_partition_does_not_stall_peers`), dropped after finishing the current chunk while waiting for the others (`..._partition_dropped_while_waiting_for_peers`, which also checks the barrier's counts), and dropped right after starting a chunk load (`..._partition_dropped_while_loading`). - `test_nlj_memory_limited_releases_memory_{left,full}_join`: nothing stays reserved after completion while the plan is still alive. - Extended test suite (`--features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`): 11,992 passed, 0 failed. - Performance, `datafusion-cli` release builds of this branch and of the `main` commit it is based on, 16 partitions, 24 MB limit, `LEFT JOIN` with a 50 MB build side (50K rows with a 1 KB string), interleaved runs on a machine that was not idle: | right side | `main` | this PR | | --- | --- | --- | | right-side passes (`input_rows` / 60K) | 7 | 7 | | 60K rows × 1 KB, fair pool | 0.78–0.88 s, 422–466 MB RSS | 0.81–0.96 s (one 1.38 s outlier), 453–481 MB | | 60K rows × 1 KB, greedy pool | 0.86–1.12 s, 408–505 MB | 0.81–1.02 s, 462–499 MB | | 600K rows × 8 bytes (compute-bound) | 7.7–10.3 s, 345–392 MB | 8.0–11.7 s, 360–425 MB | I read that as no measurable difference in time and RSS within about 5%. ## Are there any user-facing changes? - Dropping one unfinished partition of a spilling nested loop join no longer fails the other partitions; they finish, and the final unmatched left rows are not emitted, as on the in-memory path. **This is the behaviour change I would most like reviewers to weigh in on.** - Unchanged: partitions of a spilling join still have to be polled concurrently (that is inherent to sharing one chunk, and not reachable from planned queries, since the join declares no output ordering and nothing parks one partition above it); `datafusion.execution.enable_nlj_coordinated_fallback` keeps its behaviour and rationale; only its documentation is reworded, since `FallbackCoordinator` no longer exists. - No public API changes.

viirya added 10 commits September 6, 2026 12:40

github-actions Bot added the physical-plan Changes to the physical-plan crate label Sep 7, 2026

viirya requested a review from kosiew September 8, 2026 02:57

kosiew requested changes Sep 8, 2026

View reviewed changes

kosiew approved these changes Sep 9, 2026

View reviewed changes