SubhamSinghal commented on code in PR #26287:
URL: https://github.com/apache/datafusion/pull/26287#discussion_r4439931843
##########
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(
+ &mut self,
+ buffered_batch: &RecordBatch,
+ streamed_schema: &Schema,
+ buffered_range: Range<usize>,
+ ) -> Result<()> {
+ for (start, len) in chunks(buffered_range, self.batch_size) {
+ let columns = copy_runs(buffered_batch, &[(start, len)], len)?
+ .into_iter()
+ .chain(null_columns(streamed_schema, len))
+ .collect();
+ self.push_columns(columns, len)?;
+ }
+ Ok(())
+ }
+
+ /// Builds the pending rows into a completed batch.
+ fn flush(
+ &mut self,
+ buffered_batch: &RecordBatch,
+ stream_batch: &RecordBatch,
+ ) -> Result<()> {
+ let num_rows = self.pending_rows();
+ if num_rows == 0 {
+ return Ok(());
+ }
+ let mut columns = copy_runs(buffered_batch, &self.buffered_runs,
num_rows)?;
+ self.buffered_runs.clear();
+
+ let streamed_indices = UInt32Array::from(std::mem::replace(
+ &mut self.streamed_indices,
+ Vec::with_capacity(self.batch_size),
+ ));
+ for column in stream_batch.columns() {
+ columns.push(take(column.as_ref(), &streamed_indices, None)?);
+ }
+ self.push_columns(columns, num_rows)
+ }
+
+ fn push_columns(&mut self, columns: Vec<ArrayRef>, num_rows: usize) ->
Result<()> {
+ let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
+ self.completed.push_back(RecordBatch::try_new_with_options(
+ Arc::clone(&self.schema),
+ columns,
+ &options,
+ )?);
+ Ok(())
+ }
+}
+
+// Splits `range` into `(start, len)` chunks of at most `size` rows.
+fn chunks(range: Range<usize>, size: usize) -> impl Iterator<Item = (usize,
usize)> {
+ let end = range.end;
+ range
+ .step_by(size)
+ .map(move |start| (start, size.min(end - start)))
+}
+
+// All-null columns of `len` rows for the fields of `schema`.
+fn null_columns(schema: &Schema, len: usize) -> impl Iterator<Item = ArrayRef>
+ '_ {
+ schema
+ .fields()
+ .iter()
+ .map(move |f| new_null_array(f.data_type(), len))
+}
+
+// Copies `runs` of `batch`'s rows into new columns of `num_rows` rows.
+fn copy_runs(
+ batch: &RecordBatch,
+ runs: &[(usize, usize)],
+ num_rows: usize,
+) -> Result<Vec<ArrayRef>> {
+ batch
+ .columns()
+ .iter()
+ .map(|column| {
+ let data = column.to_data();
Review Comment:
addressed in
[b0fcb65](https://github.com/apache/datafusion/pull/26287/commits/b0fcb6518e33a642e6f838c072043036020d2ce5)
--
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]
