jayzhan211 commented on code in PR #25028:
URL: https://github.com/apache/datafusion/pull/25028#discussion_r3950625795
##########
datafusion/physical-plan/src/joins/hash_join/stream.rs:
##########
@@ -1054,19 +1140,69 @@ impl HashJoinStream {
)?;
let push_status = self.output_buffer.push_batch(batch)?;
- // If limit reached, finish the coalescer
+ // If limit reached, finish the coalescer and stop emitting
if push_status == PushBatchStatus::LimitReached {
self.output_buffer.finish()?;
+ self.state = HashJoinStreamState::Completed;
}
}
timer.done();
- self.state = HashJoinStreamState::Completed;
Ok(StatefulStreamResult::Continue)
}
}
+/// Returns the next chunk of final build-side indices for join types that
+/// produce build rows once the probe side is exhausted, starting at `cursor`
+/// and holding at most `batch_size` rows. Advances `cursor` past the build
+/// rows examined; `cursor == visited.len()` means every row has been examined.
+///
+/// The build indices are always valid. The probe indices are NULL for every
+/// row (`Left`, `LeftAnti`, `Full`, `LeftSemi`), except for `LeftMark`, where
+/// every build row is emitted and a NULL probe index marks an unmatched row.
+///
+/// For example, with `visited = [true, false, true, true, false]`:
+/// - `Left`: build `[1, 4]`, probe `[null, null]`
+/// - `LeftSemi`: build `[0, 2, 3]`, probe `[null, null, null]`
+/// - `LeftMark`: build `[0, 1, 2, 3, 4]`, probe `[0, null, 0, 0, null]`
+fn next_final_indices_chunk(
+ visited: &BooleanBuffer,
+ cursor: &mut usize,
+ join_type: JoinType,
+ batch_size: usize,
+) -> (UInt64Array, UInt32Array) {
+ let num_rows = visited.len();
+ let start = *cursor;
+
+ if join_type == JoinType::LeftMark {
+ // Every build row is emitted, so a chunk is a plain range of rows.
+ let end = (start + batch_size).min(num_rows);
+ let build_indices = (start as u64..end as
u64).collect::<UInt64Array>();
+ let probe_indices = (start..end)
+ .map(|idx| visited.value(idx).then_some(0))
+ .collect::<UInt32Array>();
+ *cursor = end;
+ return (build_indices, probe_indices);
+ }
+
+ // `LeftSemi` emits the matched build rows; `Left`, `LeftAnti` and `Full`
+ // emit the unmatched ones.
+ let emit_visited = join_type == JoinType::LeftSemi;
+ let mut build_indices = Vec::with_capacity(batch_size.min(num_rows -
start));
+ let mut idx = start;
+ while idx < num_rows && build_indices.len() < batch_size {
Review Comment:
Unmatched build-side rows are now emitted in batch_size chunks instead of a
single batch over the whole build side.
--
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]