comphead commented on code in PR #26287:
URL: https://github.com/apache/datafusion/pull/26287#discussion_r4317150303
##########
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:
This flush ends every stream batch with a partial output batch.
`push_unmatched_streamed` and `push_unmatched_buffered` also emit their own
partial tails. `BatchCoalescer` used to merge these across stream batches, and
`HashJoinExec` and `NestedLoopJoinExec` still do. With `batch_size`-row stream
batches and a selective predicate, the join now emits roughly one batch of the
matched rows per stream batch instead of full `batch_size` batches. The new
tests show the layouts: `[1, 4, 2, 1]`, `[2, 2, 2, 1]` and `[1, 2, 1, 2]`.
The `pwmj` suite cannot show the cost because the query is one aggregate on
top. Could you add numbers for an operator that pays per batch above the join
(a `GROUP BY` or a `RepartitionExec`) on `selective` and `half_match`? If the
cost is real, the pending rows would have to survive across stream batches.
##########
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:
The doc says copying lets each batch report only its own memory. That holds
for fixed width and `Utf8`/`Binary`. For `Utf8View`/`BinaryView`,
`MutableArrayData` appends every data buffer of the source to the output
(`variadic_data_buffers` in arrow-data `transform/mod.rs`), so each output
batch still references the whole buffered batch. `take` on the streamed side
does the same. `BatchCoalescer` compacts sparse view sources (`need_gc` in
arrow-select `coalesce/byte_view.rs`), so the old path did not retain them.
Parquet strings are views by default, so downstream reservations
(`SortExec`, `RepartitionExec`) could see up to the full buffered string data
per output batch. This is from reading the arrow sources and is not measured.
The memory test only uses `Int32`. Could you run it with a `Utf8View` column
against `main`? If it regresses, view columns would need a compaction step.
##########
datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs:
##########
@@ -106,6 +109,8 @@ pub(super) struct ClassicPWMJStream {
pub streamed: SendableRecordBatchStream,
// Streamed schema
streamed_schema: SchemaRef,
+ // Buffered schema
+ buffered_schema: SchemaRef,
Review Comment:
The buffered batch is built with `concat_batches(&buffered.schema(), ..)` in
`exec.rs`, so `buffered_data.batch().schema_ref()` already gives this schema
inside `fetch_stream_batch`. The new field, the constructor parameter and the
`exec.rs` change can be dropped.
##########
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:
`to_data()` and `MutableArrayData::try_new` run for every buffered column on
every flush, but the buffered batch does not change during the join. Holding
the `ArrayData` per column once would leave only `try_new` per flush. It
matters most for the small end-of-stream-batch flushes. Not measured.
--
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]
