viirya commented on code in PR #22038:
URL: https://github.com/apache/datafusion/pull/22038#discussion_r3887045015
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -2252,29 +2636,49 @@ impl NestedLoopJoinStream {
return ControlFlow::Break(poll);
}
- if !self.left_exhausted && self.is_memory_limited() {
- // More left data to process — free current chunk and
- // go back to BufferingLeft for the next chunk
- if let SpillState::Active(ref active) =
self.spill_state {
- active.reservation.resize(0);
+ // Drop our reference to the current chunk's
+ // `JoinLeftData` before releasing the slot. Once the
+ // last partition does this, the `Arc` reaches zero
+ // refcount and the per-chunk reservation is freed.
+ self.buffered_left_data = None;
+
+ if self.is_memory_limited() {
+ let is_emitter = self.is_unmatched_left_emitter;
+ if let SpillState::Active(active) = &mut
self.spill_state {
+ // The last partition for this chunk (the
+ // unmatched-left emitter elected in `ProbeEnd`)
+ // releases the coordinator slot so the next
+ // leader can load the following chunk.
+ if is_emitter {
+ let coordinator =
Arc::clone(&active.coordinator);
+ let released_index = active.next_chunk_index;
+ active.chunk_release_in_flight = Some(
+ async move {
+
coordinator.release_chunk(released_index).await
+ }
+ .boxed(),
+ );
+ }
+ active.next_chunk_index += 1;
}
- self.buffered_left_data = None;
+ // `is_unmatched_left_emitter` is recomputed when
+ // `ProbeEnd` is re-entered for the next chunk, so it
+ // does not need to be reset here.
+ }
+
+ if !self.left_exhausted && self.is_memory_limited() {
+ // More left data to process — go back to
+ // BufferingLeft for the next chunk.
self.left_probe_idx = 0;
self.left_emit_idx = 0;
- // Each memory-limited chunk gets a fresh per-chunk
- // `JoinLeftData`/counter; `is_unmatched_left_emitter`
is
- // recomputed when `ProbeEnd` is re-entered for the
next
- // chunk, so it does not need to be reset here.
self.state = NLJState::BufferingLeft;
} else if self.is_memory_limited()
Review Comment:
Good catch, and confirmed — this was a real leak, not just a theoretical one.
I reproduced it before fixing: holding the plan alive after collecting every
partition, the pool still had **1921 bytes** reserved for both LEFT and FULL.
Your reading of the control flow was exactly right: the release future is
created, then the final chunk goes straight to `Done` /
`EmitGlobalRightUnmatched`, neither of which polls `chunk_release_in_flight`.
Fixing it turned out to need two parts, the second of which I only found
because the test still failed:
1. A new `ReleasingFinalChunk` state that polls the release future before
finishing, then continues to whichever state would have followed. (My first
attempt just re-entered `EmitLeftUnmatched` so its existing drain would run,
but that re-runs `process_left_unmatched()` and panics with `LeftData should be
available`, since `buffered_left_data` has already been cleared by then.)
2. `release_chunk` also has to `resize(0)` the coordinator reservation once
the left side is exhausted. With only fix 1 the retained figure stayed at
exactly 1921 — the per-chunk `JoinLeftData` carries just an empty RAII
placeholder (`reservation.new_empty()`), so dropping the `Arc` frees nothing by
itself; the real bytes live in the coordinator's reservation, which otherwise
waits for a `load_one_chunk` that never comes after the final chunk.
Added `test_nlj_memory_limited_releases_final_chunk_{left,full}_join`, which
keep the plan alive after collection and assert `pool.reserved() == 0`. Both
fail without the change.
--
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]
