jayzhan211 commented on code in PR #26287:
URL: https://github.com/apache/datafusion/pull/26287#discussion_r4443718338
##########
datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs:
##########
@@ -520,78 +696,62 @@ fn resolve_classic_join(
"PiecewiseMergeJoin: streamed row {row_idx} reached the scan
without a match"
);
}
- let count = buffered_len - buffer_idx;
- let batch = build_matched_indices_and_mark_buffered(
- (buffer_idx, count),
- (row_idx, count),
+ mark_buffered(
+ buffer_idx,
buffered_side,
- stream_batch,
join_type,
- join_schema,
&mut batch_process_state.batch_min_marked,
+ );
+ batch_process_state.output_batches.push_matched(
+ buffered_batch,
+ &stream_batch.batch,
+ buffer_idx..buffered_len,
+ row_idx,
)?;
- batch_process_state.output_batches.push_batch(batch)?;
-
- // Flush batch and update pointers if we have a completed batch
- if let Some(batch) =
batch_process_state.output_batches.next_completed_batch() {
+ // Stop once a batch completes, and update pointers to resume from the
next row
+ if batch_process_state.output_batches.has_completed_batch() {
batch_process_state.start_buffer_idx = buffer_idx;
batch_process_state.start_stream_idx = row_idx + 1;
- return Ok(batch);
+ return Ok(());
}
}
+ // Flush the partial last batch: the pending rows index into this stream
batch, so they
+ // can't wait for the next one
+ batch_process_state
+ .output_batches
+ .flush(buffered_batch, &stream_batch.batch)?;
Review Comment:
@comphead I checked this against main before merging. Main also ends every
stream batch with a partial batch (`next_drained_batch` calls
`finish_buffered_batch` before the next stream batch). With 200 stream batches
that each match one row, both emit 200 output batches for Inner and 400 for
Right. So the layout is unchanged by this PR. If you think merging tails across
stream batches is worth doing, I can open a follow-up issue.
##########
datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs:
##########
@@ -445,13 +458,176 @@ impl BatchProcessState {
self.continue_process = true;
self.batch_min_marked = usize::MAX;
}
+}
+
+/// Collects the join's output as runs of buffered rows paired with streamed
rows, and
+/// builds each `batch_size` batch with one copy per column.
+///
+/// A matched streamed row pairs with a contiguous buffered suffix, often only
a few rows
+/// long. Building a batch per streamed row and coalescing those costs a
`take` and a
+/// `RecordBatch` per streamed row, and copies the streamed columns twice.
Here the buffered
+/// side is copied run by run (a `memcpy` for fixed-width types, not a per-row
gather), and
+/// the streamed side is gathered once.
+struct OutputBuffer {
+ schema: SchemaRef,
+ batch_size: usize,
+ // `(start, len)` pairs the buffered rows `[start, start + len)` with the
next `len`
+ // streamed indices.
+ buffered_runs: Vec<(usize, usize)>,
+ streamed_indices: Vec<u32>,
+ completed: VecDeque<RecordBatch>,
+}
- // `None` guarantees the coalescer holds no pending rows, so the caller
- // may safely transition state without losing output.
- fn next_drained_batch(&mut self) -> Result<Option<RecordBatch>> {
- self.output_batches.finish_buffered_batch()?;
- Ok(self.output_batches.next_completed_batch())
+impl OutputBuffer {
+ fn new(schema: SchemaRef, batch_size: usize) -> Self {
+ Self {
+ schema,
+ batch_size,
+ buffered_runs: vec![],
+ streamed_indices: Vec::with_capacity(batch_size),
+ completed: VecDeque::new(),
+ }
}
+
+ fn pending_rows(&self) -> usize {
+ self.streamed_indices.len()
+ }
+
+ fn has_completed_batch(&self) -> bool {
+ !self.completed.is_empty()
+ }
+
+ fn next_completed_batch(&mut self) -> Option<RecordBatch> {
+ self.completed.pop_front()
+ }
+
+ /// Pairs streamed row `streamed_idx` with every buffered row in
`buffered_range`.
+ fn push_matched(
+ &mut self,
+ buffered_batch: &RecordBatch,
+ stream_batch: &RecordBatch,
+ buffered_range: Range<usize>,
+ streamed_idx: usize,
+ ) -> Result<()> {
+ let mut start = buffered_range.start;
+ while start < buffered_range.end {
+ let n =
+ (self.batch_size - self.pending_rows()).min(buffered_range.end
- start);
+ self.buffered_runs.push((start, n));
+ self.streamed_indices
+ .extend(std::iter::repeat_n(streamed_idx as u32, n));
+ start += n;
+ if self.pending_rows() == self.batch_size {
+ self.flush(buffered_batch, stream_batch)?;
+ }
+ }
+ Ok(())
+ }
+
+ /// Emits `unmatched`, streamed rows that match no buffered row, with null
buffered
+ /// columns.
+ fn push_unmatched_streamed(
+ &mut self,
+ buffered_schema: &Schema,
+ unmatched: &RecordBatch,
+ ) -> Result<()> {
+ for (start, len) in chunks(0..unmatched.num_rows(), self.batch_size) {
+ let columns = null_columns(buffered_schema, len)
+ .chain(unmatched.slice(start, len).columns().iter().cloned())
+ .collect();
+ self.push_columns(columns, len)?;
+ }
+ Ok(())
+ }
+
+ /// Emits the buffered rows in `buffered_range` with null streamed columns.
+ ///
+ /// The rows are copied, not sliced: a slice would keep all of
`buffered_batch` alive
+ /// and report its full buffers to memory accounting downstream, once per
batch.
+ fn push_unmatched_buffered(
Review Comment:
@comphead I measured this against main before merging. The setup is a
20K-row buffered side with 42-byte string payloads and 200 output batches, and
the numbers are the sum of `get_record_batch_memory_size` over the output
batches:
| payload, case | main | this PR |
|---|---|---|
| Utf8View, each match 1000 rows | 217.8 MB | 217.7 MB |
| Utf8View, each match 100 rows | 419 MB | 237 MB |
| Utf8View, unmatched buffered rows | 207 MB | 207 MB |
| Utf8, each match 1000 rows | 236 MB | 20 MB |
| Utf8, each match 100 rows | 20 MB | 25 MB |
Main doesn't compact the views at this scale either, so this isn't a
regression. Compacting view columns in the join output would help both. If
you'd like that tracked, I can open a follow-up issue.
--
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]
