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 fc0fae6cc34c2..e35572c8b9f5f 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 { @@ -287,7 +283,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 @@ -316,6 +312,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 8991922779d4a..a3d23fa252e04 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -16,6 +16,7 @@ // under the License. use std::cmp::Ordering; +use std::fmt::Debug; use std::sync::Arc; use arrow::array::{ @@ -32,7 +33,7 @@ use datafusion_execution::memory_pool::MemoryReservation; /// /// This is a trait as there are several specialized implementations, such as for /// single columns or for normalized multi column keys ([`Rows`]) -pub trait CursorValues { +pub trait CursorValues: Debug + Sync + Send { fn len(&self) -> usize; /// Returns true if `l[l_idx] == r[r_idx]` @@ -43,8 +44,18 @@ pub trait CursorValues { 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). @@ -99,14 +110,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). @@ -116,6 +132,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() @@ -125,6 +153,10 @@ impl Cursor { false } } + + pub(crate) fn len(&self) -> usize { + self.values.len() - self.offset + } } impl PartialEq for Cursor { @@ -290,6 +322,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 @@ -302,6 +339,7 @@ impl CursorValues for PrimitiveValues { } } +#[derive(Debug)] pub struct ByteArrayValues { offsets: OffsetBuffer, values: Buffer, @@ -518,6 +556,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 4117789777fe8..6bbb538d0c71c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,21 +18,26 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::pin::Pin; -use std::sync::Arc; -use std::task::{Context, Poll, ready}; - -use crate::RecordBatchStream; +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_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`] @@ -49,18 +54,6 @@ pub(crate) struct SortPreservingMergeStream { /// used to record execution metrics metrics: BaselineMetrics, - /// If the stream has encountered an error or reaches the - /// `fetch` limit. - done: bool, - - /// Whether buffered rows should be drained after `done` is set. - /// - /// This is enabled when we stop because the `fetch` limit has been - /// reached, allowing partial batches left over after overflow handling to - /// be emitted on subsequent polls. It remains disabled for terminal - /// errors so the stream does not yield data after returning `Err`. - drain_in_progress_on_done: bool, - /// A loser tree that always produces the minimum cursor /// /// Node 0 stores the top winner, Nodes 1..num_streams store @@ -93,12 +86,6 @@ pub(crate) struct SortPreservingMergeStream { /// reference: loser_tree: Vec, - /// If the most recently yielded overall winner has been replaced - /// within the loser tree. A value of `false` indicates that the - /// overall winner has been yielded but the loser tree has not - /// been updated - loser_tree_adjusted: bool, - /// Target batch size batch_size: usize, @@ -151,8 +138,8 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, - /// This vector contains the indices of the partitions that have not started emitting yet. - uninitiated_partitions: Vec, + number_of_exhausted_streams: usize, + is_exhausted: Vec, } impl SortPreservingMergeStream { @@ -171,8 +158,6 @@ impl SortPreservingMergeStream { in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation), streams, metrics, - done: false, - drain_in_progress_on_done: false, cursors: (0..stream_count).map(|_| None).collect(), prev_cursors: (0..stream_count).map(|_| None).collect(), round_robin_tie_breaker_mode: false, @@ -180,15 +165,31 @@ impl SortPreservingMergeStream { current_reset_epoch: 0, poll_reset_epochs: vec![0; stream_count], loser_tree: vec![], - loser_tree_adjusted: false, batch_size, fetch, produced: 0, - uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, + number_of_exhausted_streams: 0, + is_exhausted: vec![false; stream_count], } } + pub(crate) fn into_stream(self) -> SendableRecordBatchStream + where + C: 'static, + { + let schema_clone = Arc::clone(self.in_progress.schema()); + + let cloned_metrics = self.metrics.clone(); + + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + /// If the stream at the given index is not exhausted, and the last cursor for the /// stream is finished, poll the stream for the next RecordBatch and create a new /// cursor for the stream from the returned result @@ -203,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)); @@ -214,95 +221,222 @@ 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 } - fn poll_next_inner( - &mut self, - cx: &mut Context<'_>, - ) -> Poll>> { - if self.done { - // 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`. - // Drain those leftover rows before terminating the stream, - // otherwise they would be silently dropped. - // Repeated overflows are fine — each poll emits another partial - // batch until `in_progress` is fully drained. - if self.drain_in_progress_on_done && !self.in_progress.is_empty() { - return Poll::Ready(self.emit_in_progress_batch().transpose()); - } - return Poll::Ready(None); - } + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + // This vector contains the indices of the partitions that have not started emitting yet. + let mut uninitiated_partitions = + (0..self.streams.partitions()).collect::>(); + + poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) + .await?; - // Once all partitions have set their corresponding cursors for the loser tree, - // we skip the following block. Until then, this function may be called multiple - // times and can return Poll::Pending if any partition returns Poll::Pending. - if self.loser_tree.is_empty() { - ready!(self.initialize_all_partitions(cx))?; - assert_eq!( - self.uninitiated_partitions.len(), - 0, - "all partitions should be initialized" - ); + assert_eq!(uninitiated_partitions.len(), 0); // If there are no more uninitiated partitions, set up the loser tree and continue // to the next phase. // Claim the memory for the uninitiated partitions - self.uninitiated_partitions.shrink_to_fit(); + 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 _timer = elapsed_compute.timer(); + // 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(); + + // 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; + } + self.in_progress.push_row(stream_idx); + + // stop sorting if fetch has been reached + if self.fetch_reached() { + break; + } + + if self.in_progress.len() >= self.batch_size + && let Some(batch) = self.emit_in_progress_batch()? + { + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } - loop { - // Adjust the loser tree if necessary, returning control if needed - if !self.loser_tree_adjusted { let winner = self.loser_tree[0]; // Fast path: skip the `maybe_poll_stream` call (and its `Poll` // plumbing) unless the winner's cursor is exhausted and needs a // fresh batch — it is live for almost every row. if self.cursors[winner].is_none() { - match ready!(self.maybe_poll_stream(cx, winner)) { - Ok(()) => {} - Err(e) => { - self.done = true; - return Poll::Ready(Some(Err(e))); - } + 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 full batch if the last batch value is still the winner + // we do this if: + // 1. the winner did not change - so we can skip the full new batch rather than other partial batch + // 2. The new winner have more than 1 value + // 3. the last row in the new batch beat all other streams + if winner == new_winner + && 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(); } - self.update_loser_tree(); } - let stream_idx = self.loser_tree[0]; - if self.advance_cursors(stream_idx) { - self.loser_tree_adjusted = false; - self.in_progress.push_row(stream_idx); + let last_stream_idx = self.loser_tree[0]; - // stop sorting if fetch has been reached - if self.fetch_reached() { - self.done = true; - self.drain_in_progress_on_done = true; - } else if self.in_progress.len() < self.batch_size { - continue; + // 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`. + // Drain those leftover rows before terminating the stream, + // otherwise they would be silently dropped. + // Repeated overflows are fine — each poll emits another partial + // batch until `in_progress` is fully drained. + last_batch = self.emit_in_progress_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 } - return Poll::Ready(self.emit_in_progress_batch().transpose()); + return Ok(()); + } else { + // Undo the added produced batch since we did not emit it yes, so it will mess with coalesce limit + 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` /// /// This DOES NOT return `Poll::Pending` as soon as the first uninitiated partition returns `Poll::Pending` /// so we can continue to initialize the remaining partitions - fn initialize_all_partitions(&mut self, cx: &mut Context) -> Poll> { + fn initialize_all_partitions( + &mut self, + uninitiated_partitions: &mut Vec, + cx: &mut Context, + ) -> Poll> { assert_eq!( self.loser_tree.len(), 0, @@ -311,11 +445,10 @@ impl SortPreservingMergeStream { // Manual indexing since we're iterating over the vector and shrinking it in the loop let mut idx = 0; - while idx < self.uninitiated_partitions.len() { - let partition_idx = self.uninitiated_partitions[idx]; + while idx < uninitiated_partitions.len() { + let partition_idx = uninitiated_partitions[idx]; match self.maybe_poll_stream(cx, partition_idx) { Poll::Ready(Err(e)) => { - self.done = true; return Poll::Ready(Err(e)); } Poll::Pending => { @@ -331,12 +464,12 @@ impl SortPreservingMergeStream { // place which we'll try in the next loop iteration // swap_remove will change the partition poll order, but that shouldn't // make a difference since we're waiting for all streams to be ready. - self.uninitiated_partitions.swap_remove(idx); + uninitiated_partitions.swap_remove(idx); } } } - if self.uninitiated_partitions.is_empty() { + if uninitiated_partitions.is_empty() { Poll::Ready(Ok(())) } else { // There are still uninitiated partitions so return pending. @@ -348,6 +481,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) { @@ -474,7 +647,6 @@ impl SortPreservingMergeStream { } self.loser_tree[cmp_node] = winner; } - self.loser_tree_adjusted = true; } /// Resets the poll count by incrementing the reset epoch. @@ -555,15 +727,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; } @@ -586,31 +758,13 @@ impl SortPreservingMergeStream { } self.loser_tree[0] = winner; - self.loser_tree_adjusted = true; - } -} - -impl Stream for SortPreservingMergeStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - let poll = self.poll_next_inner(cx); - self.metrics.record_poll(poll) - } -} - -impl RecordBatchStream for SortPreservingMergeStream { - fn schema(&self) -> SchemaRef { - Arc::clone(self.in_progress.schema()) } } #[cfg(test)] mod tests { use super::*; + use crate::EmptyRecordBatchStream; use crate::metrics::ExecutionPlanMetricsSet; use crate::sorts::stream::PartitionedStream; use arrow::array::Int32Array; @@ -618,7 +772,8 @@ mod tests { use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; - use futures::task::noop_waker_ref; + use futures::TryStreamExt; + use futures::stream::Fuse; use std::cmp::Ordering; #[derive(Debug)] @@ -631,6 +786,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<'_>, @@ -661,8 +826,8 @@ mod tests { } } - #[test] - fn test_done_drains_buffered_rows() { + #[tokio::test] + async fn test_done_drains_buffered_rows() { let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); let pool: Arc = Arc::new(UnboundedMemoryPool::default()); let reservation = MemoryConsumer::new("test").register(&pool); @@ -678,24 +843,20 @@ mod tests { true, ); + // Simulate rows left buffered in `in_progress` (as happens when + // `build_record_batch` emits a partial batch on offset overflow). With + // an empty input stream the merge loop breaks immediately, so the only + // way these rows reach the consumer is the generator's final drain loop. let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]) .unwrap(); stream.in_progress.push_batch(0, batch).unwrap(); stream.in_progress.push_row(0); - stream.done = true; - stream.drain_in_progress_on_done = true; - let waker = noop_waker_ref(); - let mut cx = Context::from_waker(waker); + // Drive the actual stream and confirm the buffered row is drained. + let batches: Vec = stream.into_stream().try_collect().await.unwrap(); - match stream.poll_next_inner(&mut cx) { - Poll::Ready(Some(Ok(batch))) => assert_eq!(batch.num_rows(), 1), - other => { - panic!("expected buffered rows to be drained after done, got {other:?}") - } - } - assert!(stream.in_progress.is_empty()); - assert!(matches!(stream.poll_next_inner(&mut cx), Poll::Ready(None))); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); } } 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 ade24ff0534ff..7a1dc35ecce22 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -46,7 +46,7 @@ macro_rules! merge_helper { ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ let streams = FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); - return Ok(Box::pin(SortPreservingMergeStream::new( + return Ok(SortPreservingMergeStream::new( Box::new(streams), $schema, $tracking_metrics, @@ -54,7 +54,8 @@ macro_rules! merge_helper { $fetch, $reservation, $enable_round_robin_tie_breaker, - ))); + ) + .into_stream()); }}; } @@ -233,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(); @@ -254,7 +260,7 @@ impl<'a> StreamingMergeBuilder<'a> { streams, reservation.new_empty(), )?; - Ok(Box::pin(SortPreservingMergeStream::new( + Ok(SortPreservingMergeStream::new( Box::new(streams), schema, metrics, @@ -262,6 +268,7 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, reservation, enable_round_robin_tie_breaker, - ))) + ) + .into_stream()) } }