viirya opened a new pull request, #25004: URL: https://github.com/apache/datafusion/pull/25004
## Which issue does this PR close? - Closes #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. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
