viirya commented on code in PR #25542:
URL: https://github.com/apache/datafusion/pull/25542#discussion_r4064064130
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -1431,649 +1410,352 @@ pub(crate) enum LeftLoad {
/// The left side fit the memory budget and is buffered as one batch.
InMemory(Arc<JoinLeftData>),
/// The budget ran out, so the left side was spilled during that same
pass. Every partition
- /// shares this handle, and each left chunk pass re-opens the file.
+ /// shares this handle, and the chunks that are read back from the file.
Spilled(Arc<LeftSpillData>),
}
+/// The spill file [`spill_left_input`] wrote, before it is wrapped in a
[`LeftSpillData`].
+struct SpilledLeftFile {
+ spill_manager: SpillManager,
+ spill_file: Arc<dyn SpillFile>,
+ /// Total number of rows written to `spill_file`
+ num_rows: usize,
+}
+
/// The spilled left side, shared by every output partition.
+///
+/// This is the memory-limited counterpart of [`JoinLeftData`]: the rows live
in
+/// a spill file instead of memory, but the visited bitmap and the
probe-threads
+/// counter cover the whole left side and are shared by all partitions in the
+/// same way.
+///
+/// The rows come back one chunk at a time (see [`LeftChunkBarrier`]), while
the
+/// bitmap spans all of them: bits are addressed by the row's position in the
+/// file. Keeping match tracking apart from the chunks is what lets a chunk be
+/// dropped as soon as it has been probed, with the final left rows emitted in
+/// one pass at the very end.
pub(crate) struct LeftSpillData {
/// SpillManager used to read the spill file (has the left schema)
spill_manager: SpillManager,
/// The spill file containing all left-side batches
spill_file: Arc<dyn SpillFile>,
/// Left-side schema
schema: SchemaRef,
+ /// Total number of rows in `spill_file`
+ num_rows: usize,
+ /// The pass over `spill_file` that chunks are read from. Each chunk load
+ /// takes it and hands it back for the load of the following chunk.
+ reader: Arc<Mutex<Option<LeftChunkReader>>>,
+ /// Visited bitmap over every row of `spill_file`. Empty when the join type
+ /// does not need it.
+ visited: SharedBitmapBuilder,
+ /// Counter of partitions that have not finished probing every chunk
+ probe_threads_counter: AtomicUsize,
+ /// Memory reservation for `visited`
+ reservation: MemoryReservation,
}
-/// 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,
-}
+impl LeftSpillData {
+ fn new(
+ spilled: SpilledLeftFile,
+ schema: SchemaRef,
+ with_visited_left_side: bool,
+ probe_threads_count: usize,
+ reservation: MemoryReservation,
+ ) -> Self {
+ let SpilledLeftFile {
+ spill_manager,
+ spill_file,
+ num_rows,
+ } = spilled;
+ let visited = if with_visited_left_side {
+ // Use infallible `grow`: one bit per row is all that stays in
+ // memory, and the fallback path has no other recourse.
+ reservation.grow(num_rows.div_ceil(8));
+ let mut buffer = BooleanBufferBuilder::new(num_rows);
+ buffer.append_n(num_rows, false);
+ buffer
+ } else {
+ BooleanBufferBuilder::new(0)
+ };
+ Self {
+ spill_manager,
+ spill_file,
+ schema,
+ num_rows,
+ reader: Arc::new(Mutex::new(None)),
+ visited: Mutex::new(visited),
+ probe_threads_counter: AtomicUsize::new(probe_threads_count),
+ reservation,
+ }
+ }
-/// Inner state of [`FallbackCoordinator`], guarded by a synchronous mutex.
-///
-/// Synchronous because cancellation and chunk release have to complete without
-/// another poll or await -- cancellation runs from `Drop`, which has neither
--
-/// rather than depending on a future a dropped stream would take with it. No
-/// critical section awaits: the one slow operation, reading a chunk, runs
after
-/// the guard is released.
-struct FallbackCoordinatorInner {
- /// Reservation the leader borrows to bound one chunk load.
- ///
- /// On a successful load `load_one_chunk` moves the accounted bytes into
the
- /// chunk's `JoinLeftData` with `take()`, so the accounting follows the
data
- /// rather than staying with this slot. Lazily registered by the first
- /// leader, once a runtime context is available.
- 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>,
- /// 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,
- /// A partition was dropped while still `Pending`, before the shared load
- /// decided whether the left side spills.
- ///
- /// `Pending` is entered by every eligible execution, including those whose
- /// left side ends up fitting in memory -- and those never build a shared
- /// chunk counter, so their right partitions stay independent and a dropped
- /// peer has nothing to coordinate. Cancelling on such a drop would fail a
- /// query that had no fallback at all, so it is only recorded here.
- ///
- /// Paired with `coordination_started`: whichever of the two happens second
- /// performs the cancellation, so the drop is honoured whether it precedes
or
- /// follows the execution becoming coordinated. A bool suffices -- one lost
- /// partition is enough to cancel, and nothing reads a count.
- pending_drop: bool,
- /// Set once any partition has entered the coordinated path.
- ///
- /// Remembered rather than checked in the moment, because a `Pending` drop
can
- /// arrive after a peer is already coordinating; that drop must cancel, and
- /// without this flag there would be nobody left to notice.
- coordination_started: bool,
- /// Set when a stream is dropped before finishing, which cancels the whole
- /// coordinated fallback.
- ///
- /// The partitions of a coordinated fallback are not independent: chunk
- /// advancement requires every one of them to report, so a partition that
- /// disappears mid-probe would otherwise leave the survivors waiting on a
- /// release nobody will ever make. Once set, chunk state is dropped,
waiters
- /// are woken with an error, and a loader that is still reading must
discard
- /// its result instead of publishing it.
- cancelled: bool,
-}
+ /// Open a new pass over the spilled left rows
+ fn open_pass(&self) -> Result<SendableRecordBatchStream> {
+ self.spill_manager
+ .read_spill_as_stream(Arc::clone(&self.spill_file), None)
+ }
-/// 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: 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,
- /// Broadcast signalled when the fallback is cancelled.
- ///
- /// This carries no state of its own -- `cancelled` is what persists. A
- /// delivered broadcast is itself sufficient to establish cancellation;
- /// observers read the flag before awaiting, so neither alone is relied on.
- /// Kept separate from `notify`
- /// because cancellation has to reach tasks that are not waiting on chunk
- /// progress at all: a stream parked on its right input, or a loader parked
- /// on a spill read. Waiters enable their `Notified` before reading
- /// `cancelled`, so a cancellation landing between those two steps is
- /// delivered rather than lost.
- cancel_notify: tokio::sync::Notify,
- /// Test seam that reproduces one cancellation interleaving
deterministically.
+ /// Load the next chunk, accounting for it in `reservation`.
///
- /// Set to `1` to make the leader cancel after claiming the load but before
- /// it registers its cancellation watcher. A cancellation lost in that
window
- /// strands the loader itself -- other observers may already be returning
- /// errors -- which is what the paired test checks. Consumed when it fires,
- /// so one store arms it once.
- #[cfg(test)]
- cancel_at_leader_claim: AtomicUsize,
-}
-
-impl FallbackCoordinator {
- fn new(right_partition_count: usize, with_visited_bitmap: bool) -> Self {
- Self {
- right_partition_count,
- with_visited_bitmap,
- inner: Mutex::new(FallbackCoordinatorInner {
- reservation: None,
- left_stream: None,
- carryover: None,
- left_exhausted: false,
- next_chunk_index: 0,
- current: None,
- loader_in_flight: false,
- pending_drop: false,
- coordination_started: false,
- cancelled: false,
- }),
- notify: tokio::sync::Notify::new(),
- cancel_notify: tokio::sync::Notify::new(),
- #[cfg(test)]
- cancel_at_leader_claim: AtomicUsize::new(0),
- }
+ /// The load is a shared future, the way the whole left side is a shared
+ /// [`OnceFut`]: every partition waiting for the chunk holds a clone, and
+ /// whichever of them is polled drives it. So it does not matter which
+ /// partition started the load, or whether that one is still around.
+ fn load_chunk(
+ &self,
+ reservation: MemoryReservation,
+ build_time: &Time,
+ ) -> LeftChunkFut {
+ load_left_chunk(
+ Arc::clone(&self.reader),
+ self.spill_manager.clone(),
+ Arc::clone(&self.spill_file),
+ Arc::clone(&self.schema),
+ reservation,
+ build_time.clone(),
+ )
+ .map(|chunk| chunk.map(Arc::new).map_err(Arc::new))
+ .boxed()
+ .shared()
}
- /// After the last partition finishes processing chunk
- /// `released_chunk_index`, drop the slot so the next leader can
- /// load chunk `released_chunk_index + 1`.
- fn release_chunk(self: &Arc<Self>, released_chunk_index: usize) {
+ /// Record the matches of a finished chunk, whose first row is the
+ /// `row_offset`-th row of the spill file.
+ fn merge_visited(&self, row_offset: usize, chunk_visited:
&BooleanBufferBuilder) {
+ let mut visited = self.visited.lock();
+ for idx in BitIndexIterator::new(chunk_visited.as_slice(), 0,
chunk_visited.len())
{
- let mut inner = self.inner.lock();
- if let Some(cur) = &inner.current
- && cur.chunk_index == released_chunk_index
- {
- inner.current = None;
- inner.next_chunk_index = released_chunk_index + 1;
- }
+ visited.set_bit(row_offset + idx, true);
}
- // Always notify: waiters may be blocked because they couldn't
- // become leader while a previous chunk was current.
- self.notify.notify_waiters();
}
- /// True once a partition was dropped unfinished, cancelling the fallback.
- ///
- /// Production code observes cancellation through `cancellation_watcher`,
so
- /// that a waker is registered; this plain read is for assertions only.
- #[cfg(test)]
- fn is_cancelled(&self) -> bool {
- self.inner.lock().cancelled
+ /// Decrements counter of running threads, and returns `true`
+ /// if caller is the last running thread
+ fn report_probe_completed(&self) -> bool {
+ self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1
}
- /// A future that resolves when the fallback is cancelled.
- ///
- /// Registration happens on the future's **first poll**, not at
construction:
- /// that poll enables the `Notified` and then reads `cancelled`, so
whichever
- /// happens first is observed. Callers must therefore poll it, not merely
hold
- /// it.
+ /// Take the visited bitmap for the final emission. Only the last running
+ /// thread may call this, after which the bitmap is complete.
///
- /// Callers that park on something other than chunk progress -- a stream
- /// waiting on its right input, for instance -- need one of these polled
- /// alongside their own work, or a peer's cancellation never reaches their
- /// waker.
- fn cancellation_watcher(self: &Arc<Self>) -> BoxFuture<'static, ()> {
- let coordinator = Arc::clone(self);
- async move {
- let notified = coordinator.cancel_notify.notified();
- let mut notified = std::pin::pin!(notified);
- notified.as_mut().enable();
- if coordinator.inner.lock().cancelled {
- return;
- }
- notified.await;
- }
- .boxed()
+ /// The bitmap's memory is no longer accounted for here afterwards, so that
+ /// a plan that outlives its execution does not keep it reserved. The
+ /// caller accounts for the returned buffer instead.
+ fn take_visited(&self) -> BooleanBuffer {
Review Comment:
Confirming this independently — the mechanism is exactly as described, and I
think the framing that makes it clearly a defect rather than a quirk is:
**the release is tied to the plan's lifetime, but the memory is owned by the
execution.**
`LeftSpillData` is reached through `build_side_data: OnceAsync<LeftLoad>`,
and `OnceAsync::try_once` caches the resolved future on the `Exec` itself. So
the reservation outlives every stream by construction. `take_visited()` is the
only `free()`, and it is unreachable when any partition departs early, because
not decrementing `probe_threads_counter` is precisely how "a departed partition
means no unmatched-left emission" is implemented. The two behaviours are the
same line of code, so they cannot currently be separated.
What convinced me this is worth blocking on, beyond the leak itself: this PR
already has `test_nlj_memory_limited_releases_memory_{left,full}_join`, which
asserts `pool.reserved() == 0` **while the plan is still alive**. The
dropped-partition tests only pass because they `drop(plan)` first. The PR is
therefore already treating "no memory reserved while the plan lives" as the
invariant worth testing, and the drop path quietly does not meet it. That
inconsistency between the two test groups is the tell.
On magnitude: it is `left_rows / 8` bytes, so ~12.5 MB for a 100M-row left
side — not large in isolation. But it is an unbounded leak against the memory
pool for a retained or cached plan, and the trigger (a `LIMIT` that finishes
early) is exactly the query shape this PR is fixing, so I would not wave it
through.
On the fix: I would gently push back on introducing a separate "incomplete
execution" concept. This PR's main achievement is decoupling two things that
were fused; adding a third piece of shared departure state trades some of that
back. A lighter shape that preserves the decoupling is to split what
`report_probe_completed` currently conflates — "am I the last stream to finish
with the left side" (which should govern `free()`, and which a departing stream
in `Drop` can legitimately participate in) from "should unmatched-left rows be
emitted" (which stays false if anyone departed). Then the last stream to leave,
however it leaves, releases the bitmap, and emission semantics are unchanged.
That said, this is an implementation call and I do not want to design it in
a review comment — the requirement I would hold to is just that the release
happen when the last stream goes away, not when the plan does. A regression
test in the shape of `assert_memory_released_after_completion` but with a
dropped partition, asserting before `drop(plan)`, would pin it.
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -7235,179 +5914,184 @@ pub(crate) mod tests {
assert_eq!(
pool.reserved(),
0,
- "{join_type}: the coordinator still holds the final chunk's memory
\
- while the plan is alive"
+ "{join_type}: memory is still reserved after every partition
finished"
);
Ok(())
}
- #[tokio::test]
- async fn test_nlj_memory_limited_releases_final_chunk_left_join() ->
Result<()> {
- assert_final_chunk_released(JoinType::Left).await
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_nlj_memory_limited_releases_memory_left_join() -> Result<()>
{
+ assert_memory_released_after_completion(JoinType::Left).await
}
- #[tokio::test]
- async fn test_nlj_memory_limited_releases_final_chunk_full_join() ->
Result<()> {
- assert_final_chunk_released(JoinType::Full).await
- }
-
- /// Run a NLJ across 4 right partitions, collecting every output
- /// partition CONCURRENTLY. This is required for the multi-chunk
- /// coordinator path: a chunk is not released until all partitions
- /// finish probing it, and a partition cannot advance to the next chunk
- /// until the current one is released. Collecting partitions
- /// sequentially would therefore deadlock; concurrent collection mirrors
- /// how partitions actually run under the runtime.
- async fn multi_partition_memory_limited_join_collect_concurrent(
- left: Arc<dyn ExecutionPlan>,
+ #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+ async fn test_nlj_memory_limited_releases_memory_full_join() -> Result<()>
{
+ assert_memory_released_after_completion(JoinType::Full).await
+ }
+
+ /// A two-partition LEFT join over a left side of several chunks, under a
+ /// memory limit that makes it spill.
+ fn dropped_partition_test_plan(
right: Arc<dyn ExecutionPlan>,
- join_type: &JoinType,
- join_filter: Option<JoinFilter>,
- context: Arc<TaskContext>,
- ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> {
- let partition_count = 4;
+ ) -> Result<(Arc<NestedLoopJoinExec>, Arc<TaskContext>)> {
+ let task_ctx = task_ctx_with_memory_limit(50, 1)?;
let right = Arc::new(RepartitionExec::try_new(
right,
- Partitioning::RoundRobinBatch(partition_count),
+ Partitioning::RoundRobinBatch(2),
)?) as Arc<dyn ExecutionPlan>;
-
- let nested_loop_join = Arc::new(NestedLoopJoinExec::try_new(
- left,
+ let plan = Arc::new(NestedLoopJoinExec::try_new(
+ build_left_table_multi_chunk(),
right,
- join_filter,
- join_type,
+ Some(prepare_join_filter()),
+ &JoinType::Left,
None,
)?);
- let columns = columns(&nested_loop_join.schema());
+ Ok((plan, task_ctx))
+ }
+
+ /// The partitions move through the left chunks together, so one that is
Review Comment:
+1 on `FULL` specifically, and I think there is a concrete reason it is more
valuable than "one more join type".
`FULL` has a long drop window that `LEFT` does not have. In memory-limited
mode a `FULL` partition finishes its last chunk (`finish_chunk()`, so
`chunk_index` is already `N`), then replays its **entire** right spill file in
`EmitGlobalRightUnmatched`, and only reaches `ProbeEnd` — and thus
`report_probe_completed` — after that. So there is an arbitrarily long interval
in which the stream has finished every chunk but has not reported, and a drop
landing there takes a different branch of `depart`'s `chunk_index >
inner.chunk_index` test than any `LEFT` drop does.
I worked through that path statically and believe it is correct, but that is
a derivation, not a test, and it is the sort of thing worth pinning. `LEFT
MARK` is worth covering for a different reason — its final output semantics
differ from the other left-emitting types (every left row appears, carrying a
bool), so "no emitter" means something slightly different there.
--
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]
