diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..2d8883e9c29c9 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -68,6 +68,18 @@ impl LimitedBatchCoalescer { } } + pub fn with_biggest_coalesce_batch_size( + self, + biggest_coalesce_batch_size: Option, + ) -> Self { + Self { + inner: self + .inner + .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size), + ..self + } + } + /// Return the schema of the output batches pub fn schema(&self) -> SchemaRef { self.inner.schema() diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 667a8a3ce0697..f9b214c9b41a5 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -200,16 +200,12 @@ impl ExecutionPlan for CoalesceBatchesExec { partition: usize, context: Arc, ) -> Result { - Ok(Box::pin(CoalesceBatchesStream { - input: self.input.execute(partition, context)?, - coalescer: LimitedBatchCoalescer::new( - self.input.schema(), - self.target_batch_size, - self.fetch, - ), - baseline_metrics: BaselineMetrics::new(&self.metrics, partition), - completed: false, - })) + Ok(Box::pin(CoalesceBatchesStream::new( + self.input.execute(partition, context)?, + self.target_batch_size, + self.fetch, + BaselineMetrics::new(&self.metrics, partition), + ))) } fn metrics(&self) -> Option { @@ -293,7 +289,7 @@ impl ExecutionPlan for CoalesceBatchesExec { } /// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details. -struct CoalesceBatchesStream { +pub(crate) struct CoalesceBatchesStream { /// The input plan input: SendableRecordBatchStream, /// Buffer for combining batches @@ -322,6 +318,36 @@ impl Stream for CoalesceBatchesStream { } impl CoalesceBatchesStream { + pub(crate) fn new( + input: SendableRecordBatchStream, + target_batch_size: usize, + fetch: Option, + baseline_metrics: BaselineMetrics, + ) -> Self { + CoalesceBatchesStream { + coalescer: LimitedBatchCoalescer::new( + input.schema(), + target_batch_size, + fetch, + ), + input, + baseline_metrics, + completed: false, + } + } + + pub(crate) fn with_biggest_coalesce_batch_size( + self, + biggest_coalesce_batch_size: Option, + ) -> Self { + Self { + coalescer: self + .coalescer + .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size), + ..self + } + } + fn poll_next_inner( self: &mut Pin<&mut Self>, cx: &mut Context<'_>, diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..1f0f40e3d582b 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -114,6 +114,16 @@ impl BatchBuilder { self.indices.push((cursor.batch_idx, row_idx)); } + /// Append the next `n` rows from `stream_idx` + pub fn push_n_rows(&mut self, stream_idx: usize, n: usize) { + let cursor = &mut self.cursors[stream_idx]; + let row_idx = cursor.row_idx; + + cursor.row_idx += n; + self.indices + .extend((0..n).map(|i| (cursor.batch_idx, row_idx + i))); + } + /// Returns the number of in-progress rows in this [`BatchBuilder`] pub fn len(&self) -> usize { self.indices.len() @@ -197,20 +207,22 @@ impl BatchBuilder { RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) } - /// Drains the in_progress row indexes, and builds a new RecordBatch from them + /// Drains up to `n` in_progress row indexes, and builds a new RecordBatch from them /// /// Will then drop any batches for which all rows have been yielded to the output. /// If an offset overflow occurs (e.g. string/list offsets exceed i32::MAX), /// retries with progressively fewer rows until it succeeds. /// /// Returns `None` if no pending rows - pub fn build_record_batch(&mut self) -> Result> { - if self.is_empty() { + pub fn build_record_batch(&mut self, n: usize) -> Result> { + if self.is_empty() || n == 0 { return Ok(None); } + let batch_size = self.indices.len().min(n); + let (rows_to_emit, columns) = - retry_interleave(self.indices.len(), self.indices.len(), |rows_to_emit| { + retry_interleave(batch_size, batch_size, |rows_to_emit| { self.try_interleave_columns(&self.indices[..rows_to_emit]) })?; diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index d71eaad663410..aabbe1293fe01 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -44,8 +44,18 @@ pub trait CursorValues: Debug + Sync + Send { fn eq_to_previous(cursor: &Self, idx: usize) -> bool; /// Returns comparison of `l[l_idx]` and `r[r_idx]` + /// + /// Only called with the cursors' *current* offsets, so caching + /// implementations may serve it from their cache. fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; + /// Returns comparison of `l[l_idx]` and `r[r_idx]` for *arbitrary* + /// indices, unlike [`Self::compare`]. Cold path: caching implementations + /// must override this to index directly. + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + Self::compare(l, l_idx, r, r_idx) + } + /// Notifies the values that the owning [`Cursor`] moved to `offset` (always /// `< len()`), so caching implementations can refresh the value(s) read by /// the hot comparisons. Default no-op (e.g. byte/row cursors don't benefit). @@ -96,14 +106,19 @@ impl Cursor { /// Returns true if there are no more rows in this cursor #[inline] pub fn is_finished(&self) -> bool { - self.offset == self.values.len() + self.len() == 0 } /// Advance the cursor, returning the previous row index #[inline] pub fn advance(&mut self) -> usize { + self.advance_n(1) + } + + #[inline] + pub(crate) fn advance_n(&mut self, n: usize) -> usize { let t = self.offset; - self.offset += 1; + self.offset += n; // Refresh the cache for the new position. The guard keeps `set_offset` // in bounds; a finished cursor's stale cache is never read (it is taken // before the next comparison). @@ -113,6 +128,18 @@ impl Cursor { t } + /// Compare this cursor's *last* row to `other`'s current row. Since the + /// values are sorted, `Less` means every remaining row of this cursor + /// sorts before all of `other`'s remaining rows. + pub fn compare_last_to(&self, other: &Self) -> Ordering { + T::compare_at( + &self.values, + self.values.len() - 1, + &other.values, + other.offset, + ) + } + pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor>) -> bool { if self.offset > 0 { self.is_eq_to_prev_row() @@ -122,6 +149,10 @@ impl Cursor { false } } + + pub(crate) fn len(&self) -> usize { + self.values.len() - self.offset + } } impl PartialEq for Cursor { @@ -287,6 +318,11 @@ impl CursorValues for PrimitiveValues { l.current.compare(r.current) } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Arbitrary indices (cold path), so index directly instead of using the cache. + l.values[l_idx].compare(r.values[r_idx]) + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that @@ -516,6 +552,26 @@ impl CursorValues for ArrayValues { } } + fn compare_at(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // Same null handling as `compare`, but delegates to the inner + // `compare_at` so arbitrary indices work with caching values. + match (l.is_null(l_idx), r.is_null(r_idx)) { + (true, true) => Ordering::Equal, + (true, false) => match l.options.nulls_first { + true => Ordering::Less, + false => Ordering::Greater, + }, + (false, true) => match l.options.nulls_first { + true => Ordering::Greater, + false => Ordering::Less, + }, + (false, false) => match l.options.descending { + true => T::compare_at(&r.values, r_idx, &l.values, l_idx), + false => T::compare_at(&l.values, l_idx, &r.values, r_idx), + }, + } + } + #[inline(always)] fn set_offset(&mut self, offset: usize) { // Forward to the wrapped values (e.g. caching `PrimitiveValues`). diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 986da549f75c8..c4c4338fba6f1 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,23 +18,26 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::fmt::Debug; -use std::future::poll_fn; -use std::sync::Arc; -use std::task::{Context, Poll}; - use crate::SendableRecordBatchStream; use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; +use std::cmp::Ordering; +use std::fmt::Debug; +use std::future::poll_fn; +use std::sync::Arc; +use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; -use datafusion_execution::async_try_stream; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{TryEmitter, async_try_stream}; +use futures::StreamExt; + +use crate::coalesce_batches::CoalesceBatchesStream; use futures::Stream; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] @@ -134,6 +137,9 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, + + number_of_exhausted_streams: usize, + is_exhausted: Vec, } impl SortPreservingMergeStream { @@ -163,6 +169,8 @@ impl SortPreservingMergeStream { fetch, produced: 0, enable_round_robin_tie_breaker, + number_of_exhausted_streams: 0, + is_exhausted: vec![false; stream_count], } } @@ -196,7 +204,13 @@ impl SortPreservingMergeStream { } match futures::ready!(self.streams.poll_next(cx, idx)) { - None => Poll::Ready(Ok(())), + None => { + if !self.is_exhausted[idx] { + self.is_exhausted[idx] = true; + self.number_of_exhausted_streams += 1; + } + Poll::Ready(Ok(())) + } Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok((cursor, batch))) => { self.cursors[idx] = Some(Cursor::new(cursor)); @@ -207,7 +221,11 @@ impl SortPreservingMergeStream { fn emit_in_progress_batch(&mut self) -> Result> { let rows_before = self.in_progress.len(); - let result = self.in_progress.build_record_batch(); + + // Only emit within limits + let rows_to_emit = + (self.fetch.unwrap_or(usize::MAX) - self.produced).min(self.batch_size); + let result = self.in_progress.build_record_batch(rows_to_emit); self.produced += rows_before - self.in_progress.len(); result } @@ -223,6 +241,10 @@ impl SortPreservingMergeStream { assert_eq!(uninitiated_partitions.len(), 0); + + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + // If there are no more uninitiated partitions, set up the loser tree and continue // to the next phase. @@ -230,12 +252,10 @@ impl SortPreservingMergeStream { drop(uninitiated_partitions); self.init_loser_tree(); - // NB timer records time taken on drop, so there are no - // calls to `timer.done()` below. - let elapsed_compute = self.metrics.elapsed_compute().clone(); - let mut timer = elapsed_compute.timer(); + // TODO - on first batch that is larger than the rest add to in progress - loop { + // Continue while we have more than 1 stream left + while self.number_of_exhausted_streams + 1 < self.streams.partitions() { let stream_idx = self.loser_tree[0]; if !self.advance_cursors(stream_idx) { break; @@ -262,15 +282,93 @@ impl SortPreservingMergeStream { if self.cursors[winner].is_none() { drop(timer); poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; + timer = elapsed_compute.timer(); + + // Adjusting the loser tree if necessary + self.update_loser_tree(); + + let new_winner = self.loser_tree[0]; + + // Fast path: skip comparing if the new winner batch + // we do this if: + // 1. The new winner have more than 1 value (so we won't do a needless job of comparing the entire tree if not enough values) + // 2. the last row in the new winner beat all other streams + if self.cursors[new_winner].as_ref().is_some_and(|c| c.len() > 1) + && self.winner_batch_beats_all(new_winner) + { + let cursor = self.cursors[new_winner] + .as_mut() + .expect("already validated that has cursor"); + + // TODO - avoid pushing n rows where we are above the limit or something + // Skip less than the entire number of rows so we can keep the same code flow + let number_of_rows_to_skip = cursor.len() - 1; + // If so, add the indices and + self.in_progress + .push_n_rows(new_winner, number_of_rows_to_skip); + + // Move the cursor to the end + cursor.advance_n(number_of_rows_to_skip); + } + } else { + // Adjusting the loser tree if necessary + self.update_loser_tree(); } + } + + let last_stream_idx = self.loser_tree[0]; - // Adjusting the loser tree if necessary - self.update_loser_tree(); + // Push the last stream's buffered rows that were not added to in progress + if let Some(cursor) = self.cursors[last_stream_idx].as_mut() { + let mut remaining = cursor.len(); + if let Some(fetch) = self.fetch { + remaining = remaining.min( + fetch.saturating_sub(self.produced + self.in_progress.len()), + ); + } + if remaining > 0 { + self.in_progress.push_n_rows(last_stream_idx, remaining); + cursor.advance_n(remaining); + } } drop(timer); + if self.fetch.is_none_or(|fetch| fetch > self.produced) { + self.passthrough_last_stream(emitter, last_stream_idx) + .await?; + } + + Ok(()) + }) + } + + async fn passthrough_last_stream( + &mut self, + mut emitter: TryEmitter, + last_stream_index: usize, + ) -> Result<()> { + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + let last_stream = self.streams.take_partition(last_stream_index); + + let mut last_batch = None; + + // Continue while we still have rows in the in progress builder and not reached fetch limit + while !self.in_progress.is_empty() + && self.fetch.is_none_or(|fetch| fetch > self.produced) + { + // If still not empty and we have last_batch this mean that we were unable to emit a batch with the existing indices + // and we fall back to emitting smaller one + // in that case we emit that without coalescing so we won't have error coalescing + if let Some(last_batch) = last_batch.take() { + drop(timer); + emitter.emit(last_batch).await; + timer = elapsed_compute.timer(); + } + // When `build_record_batch()` hits an i32 offset overflow (e.g. // combined string offsets exceed 2 GB), it emits a partial batch // and keeps the remaining rows in `self.in_progress.indices`. @@ -278,11 +376,58 @@ impl SortPreservingMergeStream { // otherwise they would be silently dropped. // Repeated overflows are fine — each poll emits another partial // batch until `in_progress` is fully drained. - while let Some(batch) = self.emit_in_progress_batch()? { - emitter.emit(batch).await; + last_batch = self.emit_in_progress_batch()?; + } + + // we + + // If the stream is done, we stop since nothing to emit anymore + // or if the stream is not done but we reached the limit, we stop as well (we reached the limit with the last batch) + if last_stream.is_done() || self.fetch.is_some_and(|fetch| fetch <= self.produced) + { + if let Some(last_batch) = last_batch.take() { + drop(timer); + emitter.emit(last_batch).await; + + // Not creating a timer since we are returning right away } - Ok(()) - }) + + return Ok(()); + } else { + // Undo the added produced batch since we did not emit it yes, so it will mess with coalesce limit + // TODO - this is ugly + self.produced -= last_batch.as_ref().map(|x| x.num_rows()).unwrap_or(0); + } + + let last_stream = last_stream.into_inner(); + + let last_stream = if let Some(last_batch) = last_batch { + let schema = last_stream.schema(); + let stream = futures::stream::iter(vec![Ok(last_batch)]); + + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream.chain(last_stream).boxed(), + )) as SendableRecordBatchStream + } else { + last_stream + }; + + let mut coalescer = CoalesceBatchesStream::new( + last_stream, + self.batch_size, + self.fetch.map(|x| x - self.produced), + self.metrics.intermediate(), + ) + // Don't allow for passthrough of batches with sizes other than the provided batch size + .with_biggest_coalesce_batch_size(None); + + drop(timer); + while let Some(batch) = coalescer.next().await { + emitter.emit(batch?).await; + } + + Ok(()) } /// Initialize all partitions, return `Poll::Pending` if any partition returns `Poll::Pending` @@ -338,6 +483,46 @@ impl SortPreservingMergeStream { } } + /// Returns `true` if every remaining row of `winner`'s current batch + /// sorts before every other stream's current row, using the same index + /// tie-breaking as [`Self::is_gt`]. Exhausted streams cannot compete. + fn winner_batch_beats_all(&self, winner: usize) -> bool { + let Some(winner_cursor) = &self.cursors[winner] else { + return false; + }; + + // Only the streams that lost *directly* to the winner can hold the + // overall runner-up, and those are exactly the losers stored on the + // winner's leaf-to-root path (every other stream lost to one of them + // transitively), so `O(log k)` comparisons suffice instead of `O(k)`. + let mut node = self.lt_leaf_node_index(winner); + while node != 0 { + let challenger = self.loser_tree[node]; + if let Some(other) = &self.cursors[challenger] { + let beats = match winner_cursor.compare_last_to(other) { + Ordering::Less => true, + Ordering::Equal => { + if self.enable_round_robin_tie_breaker { + // true if in round-robin because the winner already decided in a + // round-robin fashion + true + } else { + // Keep sort stable + winner < challenger + } + } + Ordering::Greater => false, + }; + if !beats { + return false; + } + } + node = self.lt_parent_node_index(node); + } + + true + } + /// For the given partition, updates the poll count. If the current value is the same /// of the previous value, it increases the count by 1; otherwise, it is reset as 0. fn update_poll_count_on_the_same_value(&mut self, partition_idx: usize) { @@ -544,15 +729,15 @@ impl SortPreservingMergeStream { if self.enable_round_robin_tie_breaker { match (&self.cursors[winner], &self.cursors[challenger]) { (Some(ac), Some(bc)) => match ac.cmp(bc) { - std::cmp::Ordering::Equal => { + Ordering::Equal => { self.handle_tie(cmp_node, &mut winner, challenger); } - std::cmp::Ordering::Greater => { + Ordering::Greater => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; self.update_winner(cmp_node, &mut winner, challenger); } - std::cmp::Ordering::Less => { + Ordering::Less => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; } @@ -581,6 +766,7 @@ impl SortPreservingMergeStream { #[cfg(test)] mod tests { use super::*; + use crate::EmptyRecordBatchStream; use crate::metrics::ExecutionPlanMetricsSet; use crate::sorts::stream::PartitionedStream; use arrow::array::Int32Array; @@ -589,6 +775,7 @@ mod tests { MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; use futures::TryStreamExt; + use futures::stream::Fuse; use std::cmp::Ordering; #[derive(Debug)] @@ -601,6 +788,16 @@ mod tests { 1 } + fn take_partition( + &mut self, + _partition_idx: usize, + ) -> Fuse { + // TODO - should not be empty schema + (Box::pin(EmptyRecordBatchStream::new(Arc::new(Schema::empty()))) + as SendableRecordBatchStream) + .fuse() + } + fn poll_next( &mut self, _cx: &mut Context<'_>, diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index 107631074ed3d..ac4ca1aee1d5c 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -47,6 +47,9 @@ pub trait PartitionedStream: std::fmt::Debug + Send { /// Returns the number of partitions fn partitions(&self) -> usize; + fn take_partition(&mut self, partition_idx: usize) + -> Fuse; + fn poll_next( &mut self, cx: &mut Context<'_>, @@ -67,6 +70,17 @@ impl std::fmt::Debug for FusedStreams { } impl FusedStreams { + fn take_stream(&mut self, stream_idx: usize) -> Fuse { + let stream_schema = self.0[stream_idx].get_ref().schema(); + + // Replace the stream with an empty stream, so we can drop memory usage + let empty_stream: Fuse = + (Box::pin(EmptyRecordBatchStream::new(stream_schema)) + as SendableRecordBatchStream) + .fuse(); + mem::replace(&mut self.0[stream_idx], empty_stream) + } + fn poll_next( &mut self, cx: &mut Context<'_>, @@ -209,6 +223,10 @@ impl PartitionedStream for RowCursorStream { self.streams.0.len() } + fn take_partition(&mut self, stream_idx: usize) -> Fuse { + self.streams.take_stream(stream_idx) + } + fn poll_next( &mut self, cx: &mut Context<'_>, @@ -279,6 +297,10 @@ impl PartitionedStream for FieldCursorStream { self.streams.0.len() } + fn take_partition(&mut self, stream_idx: usize) -> Fuse { + self.streams.take_stream(stream_idx) + } + fn poll_next( &mut self, cx: &mut Context<'_>, diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index e96138ef1306c..7a1dc35ecce22 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -234,6 +234,11 @@ impl<'a> StreamingMergeBuilder<'a> { let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); + // TODO - on only 1 stream use coalesce (so we keep the batch size contract) with observed (so the end time will be set) + // if streams.len() == 1 { + // return streams[0].clone(); + // } + // Special case single column comparisons with optimized cursor implementations if expressions.len() == 1 { let sort = expressions[0].clone();