diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..8f07a1bcd5569 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -68,6 +68,16 @@ impl LimitedBatchCoalescer { } } + /// Override the default Biggest Coalesce Batch Size + pub fn with_biggest_coalesce_batch_size( + mut 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 c5b91767777f2..b755abe79e332 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -48,6 +48,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::ready; use futures::stream::{Stream, StreamExt}; +use datafusion_proto_models::protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches; /// `CoalesceBatchesExec` combines small batches into larger batches for more /// efficient vectorized processing by later operators. @@ -200,16 +201,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 { @@ -348,7 +345,7 @@ impl CoalesceBatchesExec { } /// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details. -struct CoalesceBatchesStream { +pub(crate) struct CoalesceBatchesStream { /// The input plan input: SendableRecordBatchStream, /// Buffer for combining batches @@ -377,6 +374,48 @@ 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 + } + } + + pub(crate) fn with_memory_reservation( + self, + memory_reservation: datafusion_execution::memory_pool::MemoryReservation, + ) -> Result { + Ok(Self { + coalescer: self + .coalescer + .with_memory_reservation(memory_reservation)?, + ..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..1edcf25061a74 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -114,6 +114,14 @@ impl BatchBuilder { self.indices.push((cursor.batch_idx, row_idx)); } + /// Append the next `n` row 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() @@ -204,13 +212,15 @@ impl BatchBuilder { /// 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 initial_rows_to_emit = n.min(self.indices.len()); + let (rows_to_emit, columns) = - retry_interleave(self.indices.len(), self.indices.len(), |rows_to_emit| { + retry_interleave(initial_rows_to_emit, initial_rows_to_emit, |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..19c43ba933ea2 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -102,8 +102,14 @@ impl Cursor { /// Advance the cursor, returning the previous row index #[inline] pub fn advance(&mut self) -> usize { + self.advance_n(1) + } + + /// Advance the cursor by `n`, returning the previous row index + #[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). @@ -122,6 +128,10 @@ impl Cursor { false } } + + pub(crate) fn len(&self) -> usize { + self.values.len() - self.offset + } } impl PartialEq for Cursor { diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 310416c22d982..64c24815db02c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -35,7 +35,8 @@ use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, assert_or_internal_err, internal_err}; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::{TryEmitter, async_try_stream}; -use futures::Stream; +use futures::{Stream, StreamExt}; +use crate::coalesce_batches::CoalesceBatchesStream; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -134,6 +135,14 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, + + /// How many streams have been exhausted + number_of_exhausted_streams: usize, + + /// Tracks which streams have been exhausted + exhausted_streams: Vec, + + reservation: MemoryReservation, } impl SortPreservingMergeStream { @@ -152,6 +161,7 @@ impl SortPreservingMergeStream { let stream_count = streams.partitions(); Self { + reservation: reservation.new_empty(), in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation), streams, metrics, @@ -166,6 +176,8 @@ impl SortPreservingMergeStream { fetch, produced: 0, enable_round_robin_tie_breaker, + number_of_exhausted_streams: 0, + exhausted_streams: vec![false; stream_count], } } @@ -198,7 +210,13 @@ impl SortPreservingMergeStream { } match futures::ready!(self.streams.poll_next(cx, idx)) { - None => Poll::Ready(Ok(())), + None => { + if !self.exhausted_streams[idx] { + self.exhausted_streams[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)); @@ -209,11 +227,19 @@ 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(); + let result = self.raw_emit_in_progress_batch(); self.produced += rows_before - self.in_progress.len(); result } + fn raw_emit_in_progress_batch(&mut self) -> Result> { + // 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); + result + } + async fn flush_in_progress( &mut self, mut emitter: TryEmitter, @@ -263,8 +289,8 @@ impl SortPreservingMergeStream { // 2. Init loser tree self.init_loser_tree(); - // 3. loop until all streams have been exhausted - while !self.is_exhausted() { + // 3. loop until having only 1 non-exhuasted stream + while self.number_of_exhausted_streams + 1 < self.streams.partitions() { // 3.1. add loser_tree[0] (minimum) stream to pending record batch let winner_stream = self.loser_tree[0]; self.in_progress.push_row(winner_stream); @@ -311,8 +337,25 @@ impl SortPreservingMergeStream { self.update_loser_tree(); } - // 4. Flush any remaining rows in `self.in_progress` - self.flush_in_progress(emitter).await?; + let last_stream_idx = self.loser_tree[0]; + + // 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); + + self.passthrough_last_stream(emitter, last_stream_idx).await?; Ok(()) }) @@ -383,6 +426,96 @@ impl SortPreservingMergeStream { } } + /// When only 1 stream is left passthrough all the remaining data + 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: Option = 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 + last_batch.as_ref().map(|batch| batch.num_rows()).unwrap_or(0)) + { + // 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() { + self.produced += last_batch.num_rows(); + 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.raw_emit_in_progress_batch()?; + } + + // 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 + last_batch.as_ref().map(|batch| batch.num_rows()).unwrap_or(0)) + { + if let Some(last_batch) = last_batch.take() { + self.produced += last_batch.num_rows(); + drop(timer); + emitter.emit(last_batch).await; + + // Not creating a timer since we are returning right away + } + + return Ok(()); + } + + 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 + // To keep the contract of batch size + // TODO - add a memory test that fail if we have Multi Level Merge Sort + // and we emit smaller batches but huge in memory but right below the largest one + // that have batch size rows - then the whole Multi Level Merge Sort falls apart since + // it assumes that the batches are all batch size except the last one + // But because we are using bypass memory pool there it wont be caught + .with_biggest_coalesce_batch_size(None) + .with_memory_reservation(self.reservation.take())?; + + drop(timer); + while let Some(batch) = coalescer.next().await { + emitter.emit(batch?).await; + } + + Ok(()) + } + /// 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) {