viirya commented on code in PR #25542:
URL: https://github.com/apache/datafusion/pull/25542#discussion_r4064059682
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -2770,83 +2419,142 @@ impl NestedLoopJoinStream {
/// Memory-limited path for handle_buffering_left.
///
- /// Drives an in-flight `next_chunk` future on the coordinator, which
- /// loads (or re-uses) the next per-chunk shared `JoinLeftData`.
+ /// Gets the next left chunk, which every partition shares (see
+ /// [`LeftChunkBarrier`]).
fn handle_buffering_left_memory_limited(
&mut self,
cx: &mut std::task::Context<'_>,
) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
- let build_metric_for_chunk =
self.metrics.join_metrics.build_time.clone();
+ let build_time = self.metrics.join_metrics.build_time.clone();
let SpillState::Active(active) = &mut self.spill_state else {
unreachable!(
"handle_buffering_left_memory_limited called without Active
spill state"
);
};
- // Lazily start a chunk-fetch future for `active.next_chunk_index`.
- if active.chunk_fetch_in_flight.is_none() {
- let coordinator = Arc::clone(&active.coordinator);
- let spill_data = Arc::clone(&active.left_spill);
- let task_context = Arc::clone(&active.task_context);
- let expected = active.next_chunk_index;
- let build_metric = build_metric_for_chunk.clone();
- active.chunk_fetch_in_flight = Some(
- coordinator
- .next_chunk(expected, spill_data, task_context,
build_metric)
+ let row_offset = active.chunk_row_offset;
+ if row_offset >= active.left_spill.num_rows {
+ // Only an empty spill file ends up here, and the load does not
+ // write one. Handled anyway: there is nothing left to probe.
+ self.left_exhausted = true;
+ self.enter_state_after_last_left_chunk();
+ return ControlFlow::Continue(());
+ }
+
+ if active.chunk_fetch.is_none() {
+ active.chunk_fetch = Some(
+ Arc::clone(&active.left_chunk_barrier)
+ .chunk(
+ active.chunk_index,
+ Arc::clone(&active.left_spill),
+ Arc::clone(&active.memory_pool),
+ build_time.clone(),
+ )
.boxed(),
);
}
-
- let fut = active
- .chunk_fetch_in_flight
+ let chunk_fetch = active
+ .chunk_fetch
.as_mut()
- .expect("chunk_fetch_in_flight installed above");
- let result = match fut.poll_unpin(cx) {
- Poll::Ready(r) => r,
+ .expect("chunk_fetch installed above");
+ let chunk = match chunk_fetch.poll_unpin(cx) {
+ Poll::Ready(Ok(chunk)) => chunk,
+ Poll::Ready(Err(e)) => return
ControlFlow::Break(Poll::Ready(Some(Err(e)))),
Poll::Pending => return ControlFlow::Break(Poll::Pending),
};
- active.chunk_fetch_in_flight = None;
+ active.chunk_fetch = None;
- match result {
- Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))),
- Ok(None) => {
- // No chunk to deliver: left side fully consumed.
- self.left_exhausted = true;
- if self.is_memory_limited() &&
self.should_track_unmatched_right {
- self.right_data = None;
- self.state = NLJState::EmitGlobalRightUnmatched;
- } else {
- self.state = NLJState::Done;
- }
- ControlFlow::Continue(())
- }
- Ok(Some((data, is_last))) => {
- // The operator's own work on the delivered chunk: recording
- // metrics and opening the right-side pass.
- // `load_one_chunk` times the reading it does, but a chunk can
- // also be served straight from the coordinator's slot, in
which
- // case this is the only build work there is.
- let _build_timer = build_metric_for_chunk.timer();
- let n_rows = data.batch().num_rows();
- self.metrics.join_metrics.build_input_batches.add(1);
- self.metrics.join_metrics.build_input_rows.add(n_rows);
- self.buffered_left_data = Some(data);
- self.left_exhausted = is_last;
- self.left_buffered_in_one_pass = is_last &&
active.next_chunk_index == 0;
-
- active.right_batch_index = 0;
- match active.right_input.open_pass() {
- Ok(stream) => {
- self.right_data = Some(stream);
- }
- Err(e) => {
- return ControlFlow::Break(Poll::Ready(Some(Err(e))));
- }
- }
+ let _build_timer = build_time.timer();
+ let batch = chunk.batch.clone();
+ let n_rows = batch.num_rows();
+ self.left_exhausted = row_offset + n_rows >=
active.left_spill.num_rows;
+
+ // Matches are tracked per partition while probing, so the probe path
+ // never contends with other partitions, and merged into the global
+ // bitmap once in `finish_left_chunk`.
+ let visited_left_side = if
need_produce_result_in_final(self.join_type) {
+ // Use infallible `grow` for the bitmap -- it's small
+ active.chunk_reservation.grow(n_rows.div_ceil(8));
+ let mut buffer = BooleanBufferBuilder::new(n_rows);
+ buffer.append_n(n_rows, false);
+ buffer
+ } else {
+ BooleanBufferBuilder::new(0)
+ };
- self.state = NLJState::FetchingRight;
- ControlFlow::Continue(())
+ // This `JoinLeftData` is private to the partition: it shares the
+ // chunk's rows but has its own bitmap, so its probe-threads counter is
+ // not used. Probe completion is reported once for the whole left side,
+ // on `LeftSpillData`.
+ self.buffered_left_data = Some(Arc::new(JoinLeftData::new(
+ batch,
+ Mutex::new(visited_left_side),
+ AtomicUsize::new(1),
+ active.chunk_reservation.take(),
+ )));
+ active.current_chunk = Some(chunk);
+
+ active.right_batch_index = 0;
+ match active.right_input.open_pass() {
+ Ok(stream) => {
+ self.right_data = Some(stream);
}
+ Err(e) => {
+ return ControlFlow::Break(Poll::Ready(Some(Err(e))));
+ }
+ }
+
+ self.state = NLJState::FetchingRight;
+ ControlFlow::Continue(())
+ }
+
+ /// Record the matches of the chunk that was just probed and move on to the
+ /// next one. Memory-limited mode only.
+ ///
+ /// Unmatched-left rows are not emitted here, which would mean holding on
to
+ /// the chunk until every partition had probed it and one of them had gone
+ /// through its rows again. Emission is deferred until every partition has
+ /// probed every chunk, see
+ /// [`Self::handle_emit_left_unmatched_memory_limited`].
+ fn finish_left_chunk(&mut self) -> Result<()> {
+ let Some(left_data) = self.buffered_left_data.take() else {
+ return internal_err!("LeftData should be available");
+ };
+ let SpillState::Active(active) = &mut self.spill_state else {
+ return internal_err!("finish_left_chunk called without Active
spill state");
+ };
+
+ if need_produce_result_in_final(self.join_type) {
+ active
+ .left_spill
+ .merge_visited(active.chunk_row_offset,
&left_data.bitmap().lock());
+ }
+ active.chunk_row_offset += left_data.batch().num_rows();
+ // Let go of the chunk before reporting, so that its memory is free by
+ // the time the last partition has reported and the next one is loaded.
+ drop(left_data);
+ active.current_chunk = None;
+ active.chunk_index += 1;
+ active.left_chunk_barrier.finish_chunk();
Review Comment:
Non-blocking, but I would like a comment here: the ordering of these two
lines is load-bearing and currently implicit.
`active.chunk_index += 1` **must** happen before `finish_chunk()`, because
that is what makes `depart`'s `chunk_index > inner.chunk_index` test able to
distinguish "already finished the current chunk and waiting for peers" from
"still working on the current chunk". If these were swapped, a stream dropped
in that window would fail to decrement `finished` and the barrier could advance
while a peer was still probing.
I traced this and it is correct as written — it just reads like incidental
sequencing rather than an invariant, and the next person to touch this function
has no signal that reordering is unsafe.
##########
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 {
+ self.reservation.free();
+ self.visited.lock().finish()
}
+}
- /// Records a partition dropped before the shared load decided whether the
- /// left side spills.
- ///
- /// Whether this cancels depends on what the surviving partitions are
doing,
- /// which is why the drop is recorded rather than acted on unconditionally:
- ///
- /// * No peer has coordinated yet -- only `pending_drop` is set. If the
load
- /// resolves to `InMemory` nobody ever coordinates and this stays inert,
so
- /// an execution that never needed the coordinator is not failed by it.
- /// * A peer is already coordinating -- cancel now. That peer is waiting
on a
- /// probe report this partition will never make.
- ///
- /// The second case is the mirror of [`Self::begin_coordination`]: the two
- /// share `pending_drop` and `coordination_started` under one lock, so
- /// whichever runs second performs the cancellation and neither order is
lost.
- fn record_pending_drop(self: &Arc<Self>) {
- let cancel_now = {
- let mut inner = self.inner.lock();
- if inner.cancelled {
- return;
+/// One chunk of the spilled left side, shared by the partitions probing it.
+struct LeftChunk {
+ batch: LogicalBatch,
+ /// Memory reservation for `batch`, cleared on drop
+ #[expect(dead_code)]
+ reservation: MemoryReservation,
+}
+
+/// A load of a [`LeftChunk`] that several partitions can wait on
+type LeftChunkFut = Shared<BoxFuture<'static, SharedResult<Arc<LeftChunk>>>>;
+
+/// The pass over the left spill file that chunks are read from
+struct LeftChunkReader {
+ stream: SendableRecordBatchStream,
+ /// The batch that did not fit the previous chunk, which starts the next
+ carryover: Option<RecordBatch>,
+}
+
+/// Load the next chunk: as many batches as `reservation` accepts, and always
at
+/// least one so that the join makes progress.
+async fn load_left_chunk(
+ reader_slot: Arc<Mutex<Option<LeftChunkReader>>>,
+ spill_manager: SpillManager,
+ spill_file: Arc<dyn SpillFile>,
+ schema: SchemaRef,
+ reservation: MemoryReservation,
+ build_time: Time,
+) -> Result<LeftChunk> {
+ // Chunks are loaded one after another, so the reader the previous load
+ // handed back is where this chunk starts. The first load opens it.
+ let reader = reader_slot.lock().take();
+ let mut reader = match reader {
+ Some(reader) => reader,
+ None => LeftChunkReader {
+ stream: spill_manager.read_spill_as_stream(spill_file, None)?,
+ carryover: None,
+ },
+ };
+
+ let mut batches = vec![];
+ // The batch that did not fit the previous chunk is already in memory, so
it
+ // is accounted for infallibly.
+ if let Some(batch) = reader.carryover.take() {
+ reservation.grow(batch.get_array_memory_size());
+ batches.push(batch);
+ }
+ while let Some(batch) = reader.stream.next().await {
+ let batch = batch?;
+ // Times only the work this operator does on the batch, not the wait
+ // for the spill stream to produce it.
+ let _build_timer = build_time.timer();
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let batch_size = batch.get_array_memory_size();
+ if reservation.try_grow(batch_size).is_err() {
+ if !batches.is_empty() {
+ // Chunk is full, defer this batch to the next chunk.
+ reader.carryover = Some(batch);
+ break;
}
- inner.pending_drop = true;
- // A peer may already be coordinating, in which case this drop is
not
- // hypothetical: that peer is waiting on a report this partition
will
- // never make.
- inner.coordination_started
- };
- if cancel_now {
- self.cancel();
+ // No batches yet -- accept the batch even over budget so we make
+ // progress.
+ reservation.grow(batch_size);
}
+ batches.push(batch);
}
- /// Records that this execution is now coordinating, and honours any drop
that
- /// happened while it was not.
- ///
- /// Called when a stream enters memory-limited mode. Setting the flag
matters
- /// as much as the check: a `Pending` partition dropped *after* this point
has
- /// to cancel too, and `record_pending_drop` reads this flag to decide
that.
- fn begin_coordination(self: &Arc<Self>) {
- let cancel_now = {
- let mut inner = self.inner.lock();
- inner.coordination_started = true;
- !inner.cancelled && inner.pending_drop
- };
- if cancel_now {
- self.cancel();
- }
+ let _build_timer = build_time.timer();
+ let batch = LogicalBatch::new(schema, batches)?;
+ *reader_slot.lock() = Some(reader);
+ Ok(LeftChunk { batch, reservation })
+}
+
+/// Lets the partitions of a memory-limited join share one left chunk at a
time.
+///
+/// Re-reading the right side once per left chunk is what the fallback costs,
so
+/// chunks should be as large as memory allows. One chunk shared by every
+/// partition is as large as it gets, but it means the next chunk can only be
+/// loaded once every partition is done with the current one. That is all this
+/// type arranges: it counts the partitions that have finished the current
+/// chunk, and moves on when all of them have.
+///
+/// It lives on the plan, next to the shared left load, because partitions take
+/// part before it is known whether the left side spills at all. A partition
+/// that goes away unfinished is taken out of the count, so the others carry on
+/// without it, as they do when the left side fits in memory.
+#[derive(Debug)]
+pub(crate) struct LeftChunkBarrier {
+ inner: Mutex<LeftChunkBarrierInner>,
+ /// Signalled when the barrier moves on to the next chunk
+ notify: tokio::sync::Notify,
+}
+
+struct LeftChunkBarrierInner {
+ /// Index of the chunk the partitions are on
+ chunk_index: usize,
+ /// The load of that chunk, once a partition has asked for it. Holding it
+ /// keeps the chunk in memory for partitions that get to it later.
+ chunk: Option<LeftChunkFut>,
+ /// Partitions that have not gone away
+ live: usize,
+ /// How many of them have finished the current chunk
+ finished: usize,
+ /// Accounts for the chunks. Registered by the first load, because
partitions
+ /// take part before it is known that there will be one.
+ reservation: Option<MemoryReservation>,
+}
+
+impl std::fmt::Debug for LeftChunkBarrierInner {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ f.debug_struct("LeftChunkBarrierInner")
+ .field("chunk_index", &self.chunk_index)
+ .field("live", &self.live)
+ .field("finished", &self.finished)
+ .finish()
}
+}
- /// Cancels the coordinated fallback and drops everything the coordinator
- /// holds, synchronously.
- ///
- /// Called from a stream's drop guard when it goes away without finishing.
- /// Chunks other partitions still hold stay accounted until they release
- /// them; what this drops is the coordinator's own state, which nothing
will
- /// come back for.
- ///
- /// This is the only place that signals `cancel_notify`. Watchers rely on
- /// that: a wake from it always means a real cancellation, which is why
they
- /// need no re-arm loop. Keep it that way if you add call sites.
- fn cancel(self: &Arc<Self>) {
- {
- let mut inner = self.inner.lock();
- if inner.cancelled {
- return;
- }
- inner.cancelled = true;
- inner.current = None;
- inner.carryover = None;
- inner.left_stream = None;
- inner.reservation = None;
+impl LeftChunkBarrierInner {
+ /// Move on if every remaining partition has finished the current chunk.
+ /// Returns whether it did.
+ fn advance_if_all_finished(&mut self) -> bool {
+ if self.live > 0 && self.finished < self.live {
+ return false;
}
- self.notify.notify_waiters();
- self.cancel_notify.notify_waiters();
+ // Letting go of the load is what frees the chunk: by now no partition
+ // holds it either.
+ self.chunk = None;
+ self.chunk_index += 1;
+ self.finished = 0;
+ true
}
+}
- /// Gives up the leader claim without publishing a chunk, and wakes
waiters so
- /// they do not block on a release the failed leader never makes.
- fn abandon_load(&self) {
- self.inner.lock().loader_in_flight = false;
- self.notify.notify_waiters();
+impl LeftChunkBarrier {
+ fn new(right_partition_count: usize) -> Self {
+ Self {
+ inner: Mutex::new(LeftChunkBarrierInner {
+ chunk_index: 0,
+ chunk: None,
+ live: right_partition_count,
+ finished: 0,
+ reservation: None,
+ }),
+ notify: tokio::sync::Notify::new(),
+ }
}
- /// Fetch `expected_chunk_index`, becoming leader to load it from the
- /// left spill stream if no other partition has done so. Returns
- /// `Ok(None)` when the left stream is exhausted and no chunk with
- /// the requested index exists.
- async fn next_chunk(
+ /// The chunk with index `chunk_index`, once every partition has finished
+ /// the one before it.
+ async fn chunk(
self: Arc<Self>,
- expected_chunk_index: usize,
- spill_data: Arc<LeftSpillData>,
- task_context: Arc<TaskContext>,
+ chunk_index: usize,
+ left_spill: Arc<LeftSpillData>,
+ memory_pool: Arc<dyn MemoryPool>,
build_time: Time,
- ) -> Result<Option<(Arc<JoinLeftData>, bool)>> {
- // `spill_data` is already resolved: every partition receives the
- // same `Arc<LeftSpillData>` from the shared `OnceAsync<LeftLoad>`,
- // so the left child is executed and spilled exactly once.
+ ) -> Result<Arc<LeftChunk>> {
loop {
- // Decide what to do with the lock held, then act on that decision
- // after releasing it. The guard must not survive into the `.await`
- // below -- it is not `Send`, and holding it across the load would
- // serialize every partition behind the leader's disk reads.
- let decision = {
+ // Created under the lock, so that it cannot miss the signal of an
+ // advance that happens right after the lock is released.
+ let advanced = {
let mut inner = self.inner.lock();
-
- if inner.cancelled {
- Decision::Cancelled
- } else if let Some(cur) = &inner.current
- && cur.chunk_index == expected_chunk_index
- {
- // Case 1: requested chunk is already loaded.
- Decision::Serve(Arc::clone(&cur.data), cur.is_last)
- } else if inner.left_exhausted
- && inner.current.is_none()
- && inner.carryover.is_none()
- {
- // Case 2: left side finished and nothing left to deliver.
- Decision::Finished
- } else if inner.current.is_none() && !inner.loader_in_flight {
- // Case 3: claim the leader role and take the shared
- // resources out so the load can run without the lock.
- inner.loader_in_flight = true;
- let stream = inner.left_stream.take();
- let reservation = inner.reservation.take();
- let carryover = inner.carryover.take();
- let chunk_index_to_load = inner.next_chunk_index;
- debug_assert_eq!(chunk_index_to_load,
expected_chunk_index);
- Decision::Load {
- stream,
- reservation,
- carryover,
- chunk_index: chunk_index_to_load,
- }
+ if chunk_index == inner.chunk_index {
+ let inner = &mut *inner;
+ let reservation = inner.reservation.get_or_insert_with(|| {
+ MemoryConsumer::new("NestedLoopJoinFallbackChunk")
+ .with_can_spill(true)
+ .register(&memory_pool)
+ });
+ let chunk = inner.chunk.get_or_insert_with(|| {
+ left_spill.load_chunk(reservation.new_empty(),
&build_time)
+ });
+ Err(chunk.clone())
} else {
- // Case 4: someone else is loading, or this chunk index has
- // already been passed -- wait to be notified.
- Decision::Wait(self.notify.notified())
+ Ok(self.notify.notified())
}
};
-
- match decision {
- Decision::Cancelled => {
- return exec_err!(
- "NestedLoopJoin coordinated fallback was cancelled
because a \
- partition was dropped before finishing"
- );
- }
- Decision::Serve(data, is_last) => return Ok(Some((data,
is_last))),
- Decision::Finished => return Ok(None),
- Decision::Wait(notified) => {
- notified.await;
- }
- Decision::Load {
- stream,
- reservation,
- carryover,
- chunk_index,
- } => {
- // Build whatever the slot did not already have.
- let mut left_stream = match stream {
- Some(stream) => stream,
- None => {
- match
spill_data.spill_manager.read_spill_as_stream(
- Arc::clone(&spill_data.spill_file),
- None,
- ) {
- Ok(stream) => stream,
- Err(e) => {
- self.abandon_load();
- return Err(e);
- }
- }
- }
- };
- let mut reservation = match reservation {
- Some(reservation) => reservation,
- None => {
-
MemoryConsumer::new("NestedLoopJoinFallbackChunk".to_string())
- .with_can_spill(true)
- .register(task_context.memory_pool())
- }
- };
-
- // Race the read against cancellation. A loader parked on
- // its input is not waiting on `notify`, so without this a
- // `cancel` would not be observed until the read finished
on
- // its own -- which may be never if the input is gone.
- // Test seam: fire a cancellation in the window between
- // claiming the load and registering the watcher below. See
- // `cancel_at_leader_claim`.
- #[cfg(test)]
- if self.cancel_at_leader_claim.swap(0, Ordering::SeqCst)
== 1 {
- self.cancel();
- }
- let cancelled = self.cancel_notify.notified();
- let load = self.load_one_chunk(
- &mut left_stream,
- &mut reservation,
- carryover,
- Arc::clone(&spill_data.schema),
- build_time.clone(),
- );
- let load_result = {
- let mut load = std::pin::pin!(load);
- let mut cancelled = std::pin::pin!(cancelled);
- // Queue the waiter, then read the flag. If
cancellation
- // already happened we bail without awaiting; if it
lands
- // just after, the enabled future receives the
broadcast
- // rather than losing it.
- cancelled.as_mut().enable();
- if self.inner.lock().cancelled {
- None
- } else {
- tokio::select! {
- biased;
- result = &mut load => Some(result),
- () = &mut cancelled => None,
- }
- }
- };
- let Some(load_result) = load_result else {
- // Cancelled mid-read. Drop the local stream and
- // reservation instead of putting them back, and
- // clear the leader claim so nothing waits on us.
- drop(left_stream);
- drop(reservation);
- self.abandon_load();
- return exec_err!(
- "NestedLoopJoin coordinated fallback was cancelled
\
- while a chunk was being loaded"
- );
- };
-
- // Publish, unless the fallback was cancelled while this
- // load was running -- putting the stream and reservation
- // back would undo the cleanup `cancel` just did.
- let published = {
- let mut inner = self.inner.lock();
- inner.loader_in_flight = false;
- if inner.cancelled {
- None
- } else {
- inner.left_stream = Some(left_stream);
- inner.reservation = Some(reservation);
- match load_result {
- Ok(LoadOutcome::Chunk {
- data,
- is_last,
- carryover,
- }) => {
- inner.carryover = carryover;
- if is_last {
- inner.left_exhausted = true;
- }
- inner.current = Some(CurrentChunk {
- chunk_index,
- data: Arc::clone(&data),
- is_last,
- });
- Some(Ok(Some((data, is_last))))
- }
- Ok(LoadOutcome::Empty) => {
- inner.left_exhausted = true;
- Some(Ok(None))
- }
- Err(e) => Some(Err(e)),
- }
- }
- };
- self.notify.notify_waiters();
- match published {
- Some(result) => return result,
- None => {
- return exec_err!(
- "NestedLoopJoin coordinated fallback was
cancelled \
- while a chunk was being loaded"
- );
- }
- }
- }
+ match advanced {
+ Ok(advanced) => advanced.await,
+ Err(chunk) => return
chunk.await.map_err(DataFusionError::Shared),
}
}
}
- /// Read one chunk worth of left batches into a `JoinLeftData`,
- /// honoring the coordinator's reservation as the memory budget.
- async fn load_one_chunk(
- &self,
- left_stream: &mut SendableRecordBatchStream,
- reservation: &mut MemoryReservation,
- carryover: Option<RecordBatch>,
- left_schema: SchemaRef,
- build_time: Time,
- ) -> Result<LoadOutcome> {
- // The previous chunk's bytes were moved into its `JoinLeftData`, so
- // this reservation is already back to zero; resize defensively in case
- // a load bailed out after growing it (an error path, or `Empty`).
- reservation.resize(0);
-
- let mut pending_batches: Vec<RecordBatch> = Vec::new();
- let mut left_stream_exhausted = false;
- let mut next_carryover: Option<RecordBatch> = None;
-
- // First, account for any carryover batch from the previous
- // chunk's load attempt. Its memory is already in-flight, so we
- // grow the reservation infallibly.
- if let Some(batch) = carryover {
- let bytes = batch.get_array_memory_size();
- reservation.grow(bytes);
- pending_batches.push(batch);
+ /// A partition has finished the current chunk
+ fn finish_chunk(&self) {
+ let mut inner = self.inner.lock();
+ inner.finished += 1;
+ if inner.advance_if_all_finished() {
+ drop(inner);
+ self.notify.notify_waiters();
}
+ }
- loop {
- match left_stream.next().await {
- Some(Ok(batch)) => {
- // Times only the work this operator does on the batch, not
- // the wait for the child stream to produce it.
- let _build_timer = build_time.timer();
- if batch.num_rows() == 0 {
- continue;
- }
- let bytes = batch.get_array_memory_size();
- let can_grow = reservation.try_grow(bytes).is_ok();
- if !can_grow && !pending_batches.is_empty() {
- // Defer this batch to the next chunk.
- next_carryover = Some(batch);
- break;
- } else if !can_grow {
- // No pending batches — accept the batch even
- // over budget so we make progress.
- reservation.grow(bytes);
- }
- pending_batches.push(batch);
- }
- Some(Err(e)) => return Err(e),
- None => {
- left_stream_exhausted = true;
- break;
- }
- }
+ /// A partition that was going to ask for the chunk with index
`chunk_index`
+ /// next has gone away before finishing.
Review Comment:
Minor doc accuracy: "was going to ask for the chunk with index `chunk_index`
next" is not true for every caller.
A stream dropped after finishing the *last* chunk — in `ProbeEnd`,
`EmitLeftUnmatched`, or (for FULL) during the `EmitGlobalRightUnmatched` replay
— arrives here with `chunk_index == N`, an index that does not exist and that
it never intended to request. The arithmetic is still right (`live` and
`finished` both drop by one, which is what keeps the barrier consistent), so
this is purely about the description sending a reader looking for a "next
chunk" that isn't there.
Something closer to "the index this partition had advanced to, which may be
one past the last chunk" would describe all the callers.
--
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]
