viirya commented on code in PR #22038:
URL: https://github.com/apache/datafusion/pull/22038#discussion_r3911337518
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -1279,7 +1359,377 @@ pub(crate) struct LeftSpillData {
schema: SchemaRef,
}
-/// Tracks the state of the memory-limited spill mode for NLJ.
+/// Per-chunk shared state in the memory-limited fallback path.
+///
+/// Each chunk's `JoinLeftData` is loaded once by a "leader" partition and
+/// shared (via `Arc`) with every right-side output partition. The
+/// `probe_threads_counter` inside the `JoinLeftData` is initialized to
+/// `right_partition_count`, so `report_probe_completed` returns `true`
+/// only when the *last* partition has finished probing the chunk. That
+/// last partition is then responsible for emitting unmatched left rows
+/// for the chunk, mirroring the single-pass path's coordination via
+/// `collect_left_input(..., probe_threads_count)`.
+struct CurrentChunk {
+ /// 0-based monotonically increasing chunk index.
+ chunk_index: usize,
+ /// Shared per-chunk left data. Cloned by every partition that probes
+ /// this chunk; the last to call `report_probe_completed` emits
+ /// unmatched left rows.
+ data: Arc<JoinLeftData>,
+ /// True if the left stream was exhausted while loading this chunk —
+ /// no further chunks will be produced after it.
+ is_last: bool,
+}
+
+/// Inner state of [`FallbackCoordinator`], guarded by an async mutex.
+struct FallbackCoordinatorInner {
+ /// Reservation owned by the coordinator. Holds the memory for the
+ /// currently-loaded chunk. Reset (`resize(0)`) between chunks.
+ /// Lazily registered by the first leader, after the runtime context
+ /// becomes available via `initiate_fallback`.
+ reservation: Option<MemoryReservation>,
+ /// The shared left spill stream from which chunks are read. Owned by
+ /// the coordinator so only one partition reads it at a time.
+ left_stream: Option<SendableRecordBatchStream>,
+ /// Left schema. Set after the first leader resolves the spill future.
+ left_schema: Option<SchemaRef>,
+ /// One batch carried over from the previous chunk's load: when
+ /// reservation `try_grow` failed for chunk N, the offending batch is
+ /// recorded here and becomes the first batch of chunk N+1.
+ carryover: Option<RecordBatch>,
+ /// True once the left spill stream has produced `None`.
+ left_exhausted: bool,
+ /// Index of the next chunk to be loaded.
+ next_chunk_index: usize,
+ /// The currently-loaded chunk, or `None` if no chunk is currently
+ /// loaded (initial state, or the last partition has just released
+ /// chunk `next_chunk_index - 1` and the next leader hasn't taken
+ /// over yet).
+ current: Option<CurrentChunk>,
+ /// True while a partition has claimed leader role for the next
+ /// chunk and is loading it; prevents two partitions from racing.
+ loader_in_flight: bool,
+}
+
+/// Plan-level shared coordinator for the memory-limited fallback path.
+///
+/// All right-side output partitions share one of these. It serializes
+/// access to the left spill stream (so each chunk is read exactly once),
+/// publishes the loaded chunk as an `Arc<JoinLeftData>` for every
+/// partition to clone, and uses a `Notify` so partitions waiting for the
+/// next chunk can sleep without busy-looping.
+pub(crate) struct FallbackCoordinator {
+ /// Number of right-side partitions; equals the
+ /// `probe_threads_counter` initial value for each chunk.
+ right_partition_count: usize,
+ /// Whether `JoinLeftData` should carry a left visited bitmap (for
+ /// join types that emit unmatched left rows in the final output).
+ with_visited_bitmap: bool,
+ inner: tokio::sync::Mutex<FallbackCoordinatorInner>,
+ /// Notified when a new chunk becomes available, when the left stream
+ /// is exhausted, or when a chunk is released.
+ notify: tokio::sync::Notify,
+}
+
+impl FallbackCoordinator {
+ fn new(right_partition_count: usize, with_visited_bitmap: bool) -> Self {
+ Self {
+ right_partition_count,
+ with_visited_bitmap,
+ inner: tokio::sync::Mutex::new(FallbackCoordinatorInner {
+ reservation: None,
+ left_stream: None,
+ left_schema: None,
+ carryover: None,
+ left_exhausted: false,
+ next_chunk_index: 0,
+ current: None,
+ loader_in_flight: false,
+ }),
+ notify: tokio::sync::Notify::new(),
+ }
+ }
+
+ /// After the last partition finishes processing chunk
+ /// `released_chunk_index`, drop the slot so the next leader can
+ /// load chunk `released_chunk_index + 1`.
+ async fn release_chunk(self: &Arc<Self>, released_chunk_index: usize) {
+ let mut inner = self.inner.lock().await;
+ if let Some(cur) = &inner.current
+ && cur.chunk_index == released_chunk_index
+ {
+ inner.current = None;
+ inner.next_chunk_index = released_chunk_index + 1;
+ // Give the chunk's bytes back now rather than waiting for the next
+ // `load_one_chunk` to `resize(0)`: after the final chunk there is
no
+ // next load, and the coordinator outlives the streams because it
+ // hangs off the exec, so anything still reserved here would stay
+ // accounted against the pool for the life of the plan.
+ if inner.left_exhausted
+ && let Some(reservation) = inner.reservation.as_mut()
+ {
+ reservation.resize(0);
Review Comment:
Good catch — you're right, and my previous fix was releasing too early.
I confirmed the mechanism you describe: `maybe_flush_ready_batch` returns
*before* `buffered_left_data = None`, so a non-emitter that flushes a completed
output batch exits still holding its `Arc<JoinLeftData>`, while the emitter
goes on to release the slot. Since the chunk's bytes were accounted only in the
coordinator's reservation (the `JoinLeftData` carried an empty placeholder),
releasing on slot release under-accounted memory that was still live.
I took your second suggestion — tying reservation ownership to the shared
chunk data — because it removes the failure mode instead of coordinating around
it. `load_one_chunk` now moves the bytes it accounted into the chunk's
`JoinLeftData` via `MemoryReservation::take`, which is released on drop.
Accounting follows the data, so the bytes stay charged until the last reference
goes away, whichever stream holds it. `take` keeps the same `MemoryConsumer`,
so it transfers ownership rather than registering a new consumer, and it leaves
the coordinator's reservation at zero for the next chunk — which made the
`resize(0)` I added in the previous commit unnecessary, so that's gone too.
On the regression test: my first attempt asserted the pool was back to zero
after every partition finished, and I want to flag that it was **useless** — it
passed against the buggy version as well, because the old code also ended at
zero (the early `resize(0)` is what zeroed it). The observable difference is
only in the window while a reference is still held, so an end-state assertion
can't see it.
`test_nlj_chunk_memory_is_owned_by_the_chunk_data` drives the coordinator
directly instead, which also makes it independent of scheduling: load a chunk,
release the slot while still holding the chunk, assert the pool has *not* given
the bytes back, then drop the reference and assert it has. Reverting either
half of the fix makes it fail with `left: 0, right: 505` on the post-release
assertion, so it does pin the invariant.
Re-ran after the change: 59 NLJ unit tests, 1132 `joins`, 36 `memory_limit`
(including the three #24746 regressions), and the `nested_loop_join*` /
`information_schema` SLT files. Also rebased onto current `main` to clear the
conflict from #24566, and dropped a `continue` that its new `needless_continue`
lint flags in the `ReleasingFinalChunk` arm.
--
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]
