jayzhan211 commented on code in PR #25542:
URL: https://github.com/apache/datafusion/pull/25542#discussion_r4082846990
##########
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:
One correction to the failure mode, though: a drop cannot land between the
two lines. Drop needs &mut self, and finish_left_chunk holds it for the whole
call, so both lines run before any drop can happen. What matters is that they
are paired, not the order they run in. depart reads chunk_index >
barrier.chunk_index as "has called finish_chunk for the current chunk", so a
path that did one without the other would break it.
--
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]
