kosiew commented on code in PR #22038:
URL: https://github.com/apache/datafusion/pull/22038#discussion_r3901779483
##########
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:
I think there is still a memory-accounting race here. `release_chunk` calls
`reservation.resize(0)` as soon as the probe-counter emitter releases the final
chunk, but that emitter is only the last stream to finish probing. It is not
necessarily the last stream to drop its `buffered_left_data` reference.
For example, another partition can already be in `EmitLeftUnmatched`, flush
a completed matched-output batch from `maybe_flush_ready_batch`, and return
while still holding its `Arc<JoinLeftData>`. The emitter can then reach this
cleanup and release the reservation even though that chunk's `RecordBatch` and
bitmap are still live in the other partition. At that point the pool
under-accounts the actual live memory and could allow the configured limit to
be exceeded.
Could we keep the reservation charged until every partition has relinquished
its reference to the chunk? One option would be a separate per-chunk release
acknowledgement. Another would be to tie reservation ownership directly to the
shared chunk data so its lifetime follows the data naturally.
It would also be good to add a scheduling-sensitive regression that holds a
non-emitter after it flushes output while allowing the emitter to release the
final chunk. That should catch this lifetime gap.
##########
datafusion/common/src/config.rs:
##########
@@ -1031,6 +1031,25 @@ config_namespace! {
/// Default: 128 MB
pub max_spill_file_size_bytes: ConfigNonZeroUsize, default =
non_zero_usize_default(128 * 1024 * 1024)
+ /// Enables the memory-limited fallback for `NestedLoopJoinExec` join
+ /// types that emit unmatched left rows in the final output (LEFT, LEFT
+ /// SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple
+ /// partitions.
+ ///
+ /// This fallback coordinates per-chunk left state (visited bitmap and
+ /// probe-thread counter) across all right-side partitions, which
+ /// assumes every partition runs in the same process. Distributed
+ /// engines that execute each output partition as an independent task
+ /// (e.g. Ballista, datafusion-distributed) build a separate
coordinator
+ /// per task and poll only one partition, so the cross-partition
+ /// counter never reaches zero and the fallback would stall. Such
+ /// engines should set this to `false`: the coordinated fallback is
then
+ /// disabled for left-emitting multi-partition joins, which instead
fail
+ /// with a resource-exhaustion error under memory pressure rather than
+ /// deadlocking. Single-partition and non-left-emitting joins are
+ /// unaffected and always keep the fallback.
+ pub enable_nlj_coordinated_fallback: bool, default = true
Review Comment:
I think the SemVer concern raised by the bot still needs to be addressed.
`ExecutionOptions` is public and can be constructed exhaustively, so adding
this public field breaks downstream struct literals. `cargo-semver-checks`
reports `constructible_struct_adds_field` for this field.
Could we use a configuration mechanism that does not add a field to the
publicly constructible struct? Otherwise, if this API break is intentional, I
think it needs to be explicitly approved and targeted for the appropriate
major-version 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]
