viirya commented on code in PR #25542:
URL: https://github.com/apache/datafusion/pull/25542#discussion_r4097185805
##########
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:
You're right, thanks — `Drop` can't interleave with `finish_left_chunk`
since both need `&mut self`, so my failure scenario was wrong. The pairing is
what matters, and the new comment says exactly that.
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -3130,6 +2864,103 @@ impl NestedLoopJoinStream {
}
}
+ /// Memory-limited path for handle_emit_left_unmatched.
+ ///
+ /// The global visited bitmap is complete once the last stream has reported
+ /// probe completion, and that stream is the emitter. It streams the left
+ /// spill file one more time and emits the final left rows batch by batch,
+ /// so no left chunk has to be held in memory for it. This mirrors
+ /// `EmitGlobalRightUnmatched` on the right side.
+ fn handle_emit_left_unmatched_memory_limited(
+ &mut self,
+ cx: &mut std::task::Context<'_>,
+ ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> {
+ // Return any completed batches first
+ if let Some(poll) = self.maybe_flush_ready_batch() {
+ return ControlFlow::Break(poll);
+ }
+
+ let is_emitter = need_produce_result_in_final(self.join_type)
+ && self.is_unmatched_left_emitter;
+ let SpillState::Active(active) = &mut self.spill_state else {
+ unreachable!("memory-limited EmitLeftUnmatched without Active
spill state");
+ };
+
+ // On first entry, the emitter opens its pass over the left spill file
+ if is_emitter && active.left_unmatched_pass.is_none() {
+ let join_metric = self.metrics.join_metrics.join_time.clone();
+ let _join_timer = join_metric.timer();
+ match active.left_spill.open_pass() {
+ Ok(stream) => {
+ // Every chunk is done, which leaves `chunk_reservation`
free
+ // to account for the bitmap this stream now owns.
+ let visited = active.left_spill.take_visited();
+ active.chunk_reservation.grow(visited.len().div_ceil(8));
+ active.left_unmatched_pass = Some(LeftUnmatchedPass {
+ stream,
+ visited,
+ row_offset: 0,
+ });
+ }
+ Err(e) => return ControlFlow::Break(Poll::Ready(Some(Err(e)))),
Review Comment:
+1. A shape that would cover this and make it structural: take the bitmap
inside `report_probe_completed` when the stream is elected, and account it in
that stream's `chunk_reservation` right away. Then it is stream-owned from
`ProbeEnd` on and freed on drop, so neither an `open_pass()` error nor an early
return from `maybe_flush_ready_batch()` in `EmitLeftUnmatched` (not reachable
today, but only because every state drains the coalescer before transitioning)
can strand it. Fine as a follow-up.
--
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]
