From 4e64c6505e4dc26edcbf6bb9831ab0bc67e1d2e8 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 09:44:05 +0300 Subject: [PATCH 01/29] use yield --- Cargo.lock | 32 +++- Cargo.toml | 1 + datafusion/physical-plan/Cargo.toml | 1 + .../physical-plan/src/recursive_query.rs | 177 +++++++----------- 4 files changed, 91 insertions(+), 120 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5b43435ec0a7b..06f241366a04c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -112,7 +112,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -123,7 +123,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -476,6 +476,17 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50" +[[package]] +name = "async-fn-stream" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4ba0c4baf81a0d8ab31618ffa3ae29ceeb970a6d0d82f76130753462e39d0ea" +dependencies = [ + "futures-util", + "pin-project-lite", + "smallvec", +] + [[package]] name = "async-recursion" version = "1.1.1" @@ -2451,6 +2462,7 @@ dependencies = [ "arrow-ipc", "arrow-ord", "arrow-schema", + "async-fn-stream", "async-trait", "bytes", "criterion", @@ -2761,7 +2773,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2900,7 +2912,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4172,7 +4184,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -5330,7 +5342,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5792,7 +5804,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -5892,7 +5904,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -6061,7 +6073,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6980,7 +6992,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 0bfaad9a68b3e..a682f0967d9d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,6 +113,7 @@ arrow-ipc = { version = "59.1.0", default-features = false, features = [ ] } arrow-ord = { version = "59.1.0", default-features = false } arrow-schema = { version = "59.1.0", default-features = false } +async-fn-stream = "0.3" async-trait = "0.1.89" bigdecimal = "0.4.8" bytes = "1.11" diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index c43ae81003ccc..21b597eff41b2 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -64,6 +64,7 @@ arrow-data = { workspace = true } arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } +async-fn-stream = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 7289ac43e510c..5481f8938e380 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -19,7 +19,6 @@ use std::any::Any; use std::sync::Arc; -use std::task::{Context, Poll}; use super::work_table::{ReservedBatches, WorkTable}; use crate::aggregates::group_values::{GroupValues, new_group_values}; @@ -29,8 +28,9 @@ use crate::execution_plan::{Boundedness, EmissionType, reset_plan_states}; use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, }; +use crate::stream::RecordBatchStreamAdapter; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, }; use arrow::array::{BooleanArray, BooleanBuilder}; @@ -45,7 +45,8 @@ use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use futures::{Stream, StreamExt, ready}; +use async_fn_stream::try_fn_stream; +use futures::StreamExt; /// Recursive query execution plan. /// @@ -198,14 +199,14 @@ impl ExecutionPlan for RecursiveQueryExec { let static_stream = self.static_term.execute(partition, Arc::clone(&context))?; let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); - Ok(Box::pin(RecursiveQueryStream::new( + RecursiveQueryStream::build( context, Arc::clone(&self.work_table), Arc::clone(&self.recursive_term), static_stream, self.is_distinct, baseline_metrics, - )?)) + ) } fn metrics(&self) -> Option { @@ -259,14 +260,6 @@ struct RecursiveQueryStream { work_table: Arc, /// The dynamic part (recursive term) as is (without being executed) recursive_term: Arc, - /// The static part (static term) as a stream. If the processing of this - /// part is completed, then it will be None. - static_stream: Option, - /// The dynamic part (recursive term) as a stream. If the processing of this - /// part has not started yet, or has been completed, then it will be None. - recursive_stream: Option, - /// The schema of the output. - schema: SchemaRef, /// In-memory buffer for storing a copy of the current results. Will be /// cleared after each iteration. buffer: Vec, @@ -279,87 +272,93 @@ struct RecursiveQueryStream { } impl RecursiveQueryStream { - /// Create a new recursive query stream - fn new( + /// Build the recursive-query output as a linear async generator. + /// + /// A recursive CTE is inherently iterative: emit the static term, then run + /// the recursive term repeatedly — each iteration feeding the previous + /// iteration's rows back through the work table — until an iteration + /// produces no new rows. Expressed with `try_fn_stream`, that description is + /// the code (see the algorithm doc above): a `while` over the static stream + /// followed by a `loop` of recursive iterations, `emitter.emit`-ing every + /// batch. State that must survive across iterations (buffer, reservation, + /// dedup, metrics) lives on `self`, captured by the async body — the two + /// `Option` sub-stream fields and the hand-rolled poll state machine are + /// gone. + fn build( task_context: Arc, work_table: Arc, recursive_term: Arc, static_stream: SendableRecordBatchStream, is_distinct: bool, baseline_metrics: BaselineMetrics, - ) -> Result { + ) -> Result { let schema = static_stream.schema(); let reservation = MemoryConsumer::new("RecursiveQuery").register(task_context.memory_pool()); let distinct_deduplicator = is_distinct .then(|| DistinctDeduplicator::new(Arc::clone(&schema), &task_context)) .transpose()?; - Ok(Self { + + let state = Self { task_context, work_table, recursive_term, - static_stream: Some(static_stream), - recursive_stream: None, - schema, buffer: vec![], reservation, distinct_deduplicator, baseline_metrics, - }) - } + }; - /// Push a clone of the given batch to the in memory buffer, and then return - /// a poll with it. - fn push_batch( - mut self: std::pin::Pin<&mut Self>, - mut batch: RecordBatch, - ) -> Poll>> { - let baseline_metrics = self.baseline_metrics.clone(); + let stream = try_fn_stream(|emitter| async move { + let mut state = state; + let mut static_stream = static_stream; - if let Some(deduplicator) = &mut self.distinct_deduplicator { - let _timer_guard = baseline_metrics.elapsed_compute().timer(); - batch = deduplicator.deduplicate(&batch)?; - } + // Static term: forward and buffer every batch. + while let Some(batch) = static_stream.next().await { + emitter.emit(state.push_batch(batch?)?).await; + } - if let Err(e) = self.reservation.try_grow(batch.get_array_memory_size()) { - return Poll::Ready(Some(Err(e))); - } - self.buffer.push(batch.clone()); - (&batch).record_output(&baseline_metrics); - Poll::Ready(Some(Ok(batch))) + // Recursive term: iterate until an iteration produces no new rows. + loop { + let buffered_rows: usize = + state.buffer.iter().map(|b| b.num_rows()).sum(); + if buffered_rows == 0 { + return Ok(()); + } + + // Feed the previous iteration's rows into the work table. + state.work_table.update(ReservedBatches::new( + std::mem::take(&mut state.buffer), + state.reservation.take(), + )); + + // We always (re-)execute partition 0; downstream plans should + // not expect any partitioning. + let recursive_plan = + reset_plan_states(Arc::clone(&state.recursive_term))?; + let mut recursive_stream = + recursive_plan.execute(0, Arc::clone(&state.task_context))?; + while let Some(batch) = recursive_stream.next().await { + emitter.emit(state.push_batch(batch?)?).await; + } + } + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } - /// Start polling for the next iteration, will be called either after the static term - /// is completed or another term is completed. It will follow the algorithm above on - /// to check whether the recursion has ended. - fn poll_next_iteration( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll>> { - let total_length = self - .buffer - .iter() - .fold(0, |acc, batch| acc + batch.num_rows()); - - if total_length == 0 { - return Poll::Ready(None); + /// Deduplicate (when `DISTINCT`), account for memory, buffer a copy for the + /// next iteration's work table, record output metrics, and return the batch. + fn push_batch(&mut self, mut batch: RecordBatch) -> Result { + if let Some(deduplicator) = &mut self.distinct_deduplicator { + let _timer_guard = self.baseline_metrics.elapsed_compute().timer(); + batch = deduplicator.deduplicate(&batch)?; } - // Update the work table with the current buffer - let reserved_batches = ReservedBatches::new( - std::mem::take(&mut self.buffer), - self.reservation.take(), - ); - self.work_table.update(reserved_batches); - - // We always execute (and re-execute iteratively) the first partition. - // Downstream plans should not expect any partitioning. - let partition = 0; - - let recursive_plan = reset_plan_states(Arc::clone(&self.recursive_term))?; - self.recursive_stream = - Some(recursive_plan.execute(partition, Arc::clone(&self.task_context))?); - self.poll_next(cx) + self.reservation.try_grow(batch.get_array_memory_size())?; + self.buffer.push(batch.clone()); + (&batch).record_output(&self.baseline_metrics); + Ok(batch) } } @@ -387,48 +386,6 @@ fn assign_work_table( .data() } -impl Stream for RecursiveQueryStream { - type Item = Result; - - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - if let Some(static_stream) = &mut self.static_stream { - // While the static term's stream is available, we'll be forwarding the batches from it (also - // saving them for the initial iteration of the recursive term). - let batch_result = ready!(static_stream.poll_next_unpin(cx)); - match &batch_result { - None => { - // Once this is done, we can start running the setup for the recursive term. - self.static_stream = None; - self.poll_next_iteration(cx) - } - Some(Ok(batch)) => self.push_batch(batch.clone()), - _ => Poll::Ready(batch_result), - } - } else if let Some(recursive_stream) = &mut self.recursive_stream { - let batch_result = ready!(recursive_stream.poll_next_unpin(cx)); - match batch_result { - None => { - self.recursive_stream = None; - self.poll_next_iteration(cx) - } - Some(Ok(batch)) => self.push_batch(batch), - _ => Poll::Ready(batch_result), - } - } else { - Poll::Ready(None) - } - } -} - -impl RecordBatchStream for RecursiveQueryStream { - /// Get the schema - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} /// Deduplicator based on a hash table. struct DistinctDeduplicator { From 4432b46f040022bb1518a86c7e99335240915bf9 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 09:53:45 +0300 Subject: [PATCH 02/29] prototype --- datafusion/physical-plan/src/sorts/merge.rs | 329 ++++++++---------- .../src/sorts/streaming_merge.rs | 8 +- 2 files changed, 152 insertions(+), 185 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4583d19e91061..bade13e39858c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,28 +18,30 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::pin::Pin; +use std::fmt::{Debug, Formatter}; +use std::future::poll_fn; use std::sync::Arc; -use std::task::{Context, Poll, ready}; +use std::task::{Context, Poll}; -use crate::RecordBatchStream; -use crate::metrics::BaselineMetrics; +use crate::SendableRecordBatchStream; +use crate::metrics::{BaselineMetrics, RecordOutput}; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; +use crate::stream::RecordBatchStreamAdapter; 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 futures::Stream; +use async_fn_stream::{try_fn_stream, TryStreamEmitter}; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; /// Merges a stream of sorted cursors and record batches into a single sorted stream -#[derive(Debug)] +// #[derive(Debug)] pub(crate) struct SortPreservingMergeStream { in_progress: BatchBuilder, @@ -49,18 +51,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 @@ -153,6 +143,21 @@ pub(crate) struct SortPreservingMergeStream { /// This vector contains the indices of the partitions that have not started emitting yet. uninitiated_partitions: Vec, + + emitter: TryStreamEmitter, +} + +// TODO - fix the debug not showing all the properties now +impl Debug for SortPreservingMergeStream { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SortPreservingMergeStream") + .field("current_reset_epoch", &self.current_reset_epoch) + .field("prev_cursors", &self.prev_cursors) + .field("fetch", &self.fetch) + .field("produced", &self.produced) + .field("uninitiated_partitions", &self.uninitiated_partitions) + .finish() + } } impl SortPreservingMergeStream { @@ -164,6 +169,7 @@ impl SortPreservingMergeStream { fetch: Option, reservation: MemoryReservation, enable_round_robin_tie_breaker: bool, + emitter: TryStreamEmitter ) -> Self { let stream_count = streams.partitions(); @@ -171,8 +177,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, @@ -186,28 +190,53 @@ impl SortPreservingMergeStream { produced: 0, uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, + emitter, } } + pub(crate) fn create(streams: CursorStream, + schema: SchemaRef, + metrics: BaselineMetrics, + batch_size: usize, + fetch: Option, + reservation: MemoryReservation, + enable_round_robin_tie_breaker: bool,) -> SendableRecordBatchStream { + + let schema_clone = Arc::clone(&schema); + + let stream = try_fn_stream(|emitter| async move { + let s = SortPreservingMergeStream::new( + streams, + schema, + metrics, + batch_size, + fetch, + reservation, + enable_round_robin_tie_breaker, + emitter + ); + + s.run().await + }); + + Box::pin(RecordBatchStreamAdapter::new(schema_clone, stream)) + } + /// 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 - fn maybe_poll_stream( - &mut self, - cx: &mut Context<'_>, - idx: usize, - ) -> Poll> { + async fn maybe_poll_stream(&mut self, idx: usize) -> Result<()> { if self.cursors[idx].is_some() { // Cursor is not finished - don't need a new RecordBatch yet - return Poll::Ready(Ok(())); + return Ok(()); } - match futures::ready!(self.streams.poll_next(cx, idx)) { - None => Poll::Ready(Ok(())), - Some(Err(e)) => Poll::Ready(Err(e)), + match poll_fn(|cx| self.streams.poll_next(cx, idx)).await { + None => Ok(()), + Some(Err(e)) => Err(e), Some(Ok((cursor, batch))) => { self.cursors[idx] = Some(Cursor::new(cursor)); - Poll::Ready(self.in_progress.push_batch(idx, batch)) + self.in_progress.push_batch(idx, batch) } } } @@ -219,112 +248,70 @@ impl SortPreservingMergeStream { 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); + async fn run(mut self) -> Result<()> { + // Prime the first batch of every partition, then build the loser tree. + // With yield support this is a plain sequential loop — no need to track + // `uninitiated_partitions` across re-entrant polls, and a `Pending` + // upstream simply suspends here instead of unwinding the whole call. + // An error aborts the stream immediately; buffered rows are dropped, so + // no data is emitted after an `Err`. + for partition_idx in 0..self.cursors.len() { + self.maybe_poll_stream(partition_idx).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() { - // 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]; - match self.maybe_poll_stream(cx, partition_idx) { - Poll::Ready(Err(e)) => { - self.done = true; - return Poll::Ready(Some(Err(e))); - } - Poll::Pending => { - // The polled stream is pending which means we're already set up to - // be woken when necessary - // Try the next stream - idx += 1; - } - _ => { - // The polled stream is ready - // Remove it from uninitiated_partitions - // Don't bump idx here, since a new element will have taken its - // 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); - } - } - } + self.init_loser_tree(); - if self.uninitiated_partitions.is_empty() { - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. + // NB timer records time taken on drop; we drop it around `.await`s so + // upstream/consumer wait time is not counted as compute. + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); - // Claim the memory for the uninitiated partitions - self.uninitiated_partitions.shrink_to_fit(); - self.init_loser_tree(); - } else { - // There are still uninitiated partitions so return pending. - // We only get here if we've polled all uninitiated streams and at least one of them - // returned pending itself. That means we will be woken as soon as one of the - // streams would like to be polled again. - // There is no need to reschedule ourselves eagerly. - return Poll::Pending; + loop { + // `loser_tree[0]` is the current, already-adjusted winner. + let stream_idx = self.loser_tree[0]; + if !self.advance_cursors(stream_idx) { + // The overall minimum is exhausted, which only happens once + // every input is exhausted. We're done. + break; } - } + self.in_progress.push_row(stream_idx); - // 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(); + // Stop sorting if the fetch limit has been reached. Remaining + // buffered rows (e.g. a partial batch left after overflow handling) + // are drained below. + if self.fetch_reached() { + break; + } - 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))); - } - } + if self.in_progress.len() >= self.batch_size { + if let Some(batch) = self.emit_in_progress_batch()? { + drop(timer); + self.emitter.emit(batch).await; + timer = elapsed_compute.timer(); } - 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); - - // 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; - } + // Refill the just-consumed winner (only if its cursor is exhausted — + // live for almost every row) and re-establish the tree top. + let winner = self.loser_tree[0]; + if self.cursors[winner].is_none() { + drop(timer); + self.maybe_poll_stream(winner).await?; + timer = elapsed_compute.timer(); } + self.update_loser_tree(); + } - return Poll::Ready(self.emit_in_progress_batch().transpose()); + drop(timer); + + // Drain any buffered rows: the normal final partial batch, plus leftovers + // kept after a `fetch` stop or after `build_record_batch()` hit an i32 + // offset overflow (e.g. combined string offsets exceed 2 GB) and emitted + // only a partial batch. Each iteration emits another partial batch until + // `in_progress` is fully drained, so nothing is silently dropped. + while let Some(batch) = self.emit_in_progress_batch()? { + self.emitter.emit(batch).await; } + Ok(()) } /// For the given partition, updates the poll count. If the current value is the same @@ -569,24 +556,6 @@ impl SortPreservingMergeStream { } } -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::*; @@ -597,7 +566,6 @@ mod tests { use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; - use futures::task::noop_waker_ref; use std::cmp::Ordering; #[derive(Debug)] @@ -640,41 +608,40 @@ mod tests { } } - #[test] - 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); - let metrics = ExecutionPlanMetricsSet::new(); - - let mut stream = SortPreservingMergeStream::::new( - Box::new(EmptyPartitionedStream), - Arc::clone(&schema), - BaselineMetrics::new(&metrics, 0), - 16, - Some(1), - reservation, - true, - ); - - 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); - - 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))); - } + // TODO - uncomment this test + // + // /// The merge loop's terminal drain (`while let Some(b) = emit()? { send }`) + // /// relies on `emit_in_progress_batch` returning buffered rows and then + // /// `None` once empty. This exercises that primitive directly. + // #[test] + // fn test_emit_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); + // let metrics = ExecutionPlanMetricsSet::new(); + // + // let mut stream = SortPreservingMergeStream::::create( + // Box::new(EmptyPartitionedStream), + // Arc::clone(&schema), + // BaselineMetrics::new(&metrics, 0), + // 16, + // Some(1), + // reservation, + // true, + // ); + // + // 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); + // + // let batch = stream + // .emit_in_progress_batch() + // .unwrap() + // .expect("buffered row should be emitted"); + // assert_eq!(batch.num_rows(), 1); + // assert!(stream.in_progress.is_empty()); + // assert!(stream.emit_in_progress_batch().unwrap().is_none()); + // } } diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index ade24ff0534ff..2e83662994f53 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::create( Box::new(streams), $schema, $tracking_metrics, @@ -54,7 +54,7 @@ macro_rules! merge_helper { $fetch, $reservation, $enable_round_robin_tie_breaker, - ))); + )); }}; } @@ -254,7 +254,7 @@ impl<'a> StreamingMergeBuilder<'a> { streams, reservation.new_empty(), )?; - Ok(Box::pin(SortPreservingMergeStream::new( + Ok(SortPreservingMergeStream::create( Box::new(streams), schema, metrics, @@ -262,6 +262,6 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, reservation, enable_round_robin_tie_breaker, - ))) + )) } } From 079fc8a5d223e3ebb34b17487cc78c0024ca0cd6 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:04:42 +0300 Subject: [PATCH 03/29] fix compilation --- datafusion/physical-plan/src/sorts/cursor.rs | 4 +++- datafusion/physical-plan/src/sorts/merge.rs | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 8991922779d4a..ff2753c19bc18 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]` @@ -302,6 +303,7 @@ impl CursorValues for PrimitiveValues { } } +#[derive(Debug)] pub struct ByteArrayValues { offsets: OffsetBuffer, values: Buffer, diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bade13e39858c..84ac9b58dfc6e 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -160,7 +160,7 @@ impl Debug for SortPreservingMergeStream { } } -impl SortPreservingMergeStream { +impl SortPreservingMergeStream { pub(crate) fn new( streams: CursorStream, schema: SchemaRef, @@ -200,7 +200,7 @@ impl SortPreservingMergeStream { batch_size: usize, fetch: Option, reservation: MemoryReservation, - enable_round_robin_tie_breaker: bool,) -> SendableRecordBatchStream { + enable_round_robin_tie_breaker: bool) -> SendableRecordBatchStream where C: 'static { let schema_clone = Arc::clone(&schema); From 18262df1a333bced126ac3dcc9444a158f989800 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:38:30 +0300 Subject: [PATCH 04/29] try another way with different crate --- Cargo.lock | 29 +-- Cargo.toml | 2 +- datafusion/physical-plan/Cargo.toml | 2 +- .../physical-plan/src/recursive_query.rs | 177 +++++++++++------- datafusion/physical-plan/src/sorts/builder.rs | 5 - datafusion/physical-plan/src/sorts/merge.rs | 71 +++---- 6 files changed, 157 insertions(+), 129 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 06f241366a04c..6179841f3c9e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -476,17 +476,6 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50" -[[package]] -name = "async-fn-stream" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4ba0c4baf81a0d8ab31618ffa3ae29ceeb970a6d0d82f76130753462e39d0ea" -dependencies = [ - "futures-util", - "pin-project-lite", - "smallvec", -] - [[package]] name = "async-recursion" version = "1.1.1" @@ -2462,7 +2451,6 @@ dependencies = [ "arrow-ipc", "arrow-ord", "arrow-schema", - "async-fn-stream", "async-trait", "bytes", "criterion", @@ -2480,6 +2468,7 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "futures", + "genawaiter", "half", "hashbrown 0.17.1", "indexmap 2.14.0", @@ -3184,6 +3173,22 @@ dependencies = [ "prost-build", ] +[[package]] +name = "genawaiter" +version = "0.99.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c86bd0361bcbde39b13475e6e36cb24c329964aa2611be285289d1e4b751c1a0" +dependencies = [ + "futures-core", + "genawaiter-macro", +] + +[[package]] +name = "genawaiter-macro" +version = "0.99.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b32dfe1fdfc0bbde1f22a5da25355514b5e450c33a6af6770884c8750aedfbc" + [[package]] name = "generic-array" version = "0.14.7" diff --git a/Cargo.toml b/Cargo.toml index a682f0967d9d9..c847b76df2a7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,7 +113,7 @@ arrow-ipc = { version = "59.1.0", default-features = false, features = [ ] } arrow-ord = { version = "59.1.0", default-features = false } arrow-schema = { version = "59.1.0", default-features = false } -async-fn-stream = "0.3" +genawaiter = {version = "0.99.1", default-features = false} async-trait = "0.1.89" bigdecimal = "0.4.8" bytes = "1.11" diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 21b597eff41b2..343c15feedcb1 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -64,7 +64,7 @@ arrow-data = { workspace = true } arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } -async-fn-stream = { workspace = true } +genawaiter = { workspace = true, features = ["futures03"] } async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true } diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 5481f8938e380..7289ac43e510c 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -19,6 +19,7 @@ use std::any::Any; use std::sync::Arc; +use std::task::{Context, Poll}; use super::work_table::{ReservedBatches, WorkTable}; use crate::aggregates::group_values::{GroupValues, new_group_values}; @@ -28,9 +29,8 @@ use crate::execution_plan::{Boundedness, EmissionType, reset_plan_states}; use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, }; -use crate::stream::RecordBatchStreamAdapter; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use arrow::array::{BooleanArray, BooleanBuilder}; @@ -45,8 +45,7 @@ use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use async_fn_stream::try_fn_stream; -use futures::StreamExt; +use futures::{Stream, StreamExt, ready}; /// Recursive query execution plan. /// @@ -199,14 +198,14 @@ impl ExecutionPlan for RecursiveQueryExec { let static_stream = self.static_term.execute(partition, Arc::clone(&context))?; let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); - RecursiveQueryStream::build( + Ok(Box::pin(RecursiveQueryStream::new( context, Arc::clone(&self.work_table), Arc::clone(&self.recursive_term), static_stream, self.is_distinct, baseline_metrics, - ) + )?)) } fn metrics(&self) -> Option { @@ -260,6 +259,14 @@ struct RecursiveQueryStream { work_table: Arc, /// The dynamic part (recursive term) as is (without being executed) recursive_term: Arc, + /// The static part (static term) as a stream. If the processing of this + /// part is completed, then it will be None. + static_stream: Option, + /// The dynamic part (recursive term) as a stream. If the processing of this + /// part has not started yet, or has been completed, then it will be None. + recursive_stream: Option, + /// The schema of the output. + schema: SchemaRef, /// In-memory buffer for storing a copy of the current results. Will be /// cleared after each iteration. buffer: Vec, @@ -272,93 +279,87 @@ struct RecursiveQueryStream { } impl RecursiveQueryStream { - /// Build the recursive-query output as a linear async generator. - /// - /// A recursive CTE is inherently iterative: emit the static term, then run - /// the recursive term repeatedly — each iteration feeding the previous - /// iteration's rows back through the work table — until an iteration - /// produces no new rows. Expressed with `try_fn_stream`, that description is - /// the code (see the algorithm doc above): a `while` over the static stream - /// followed by a `loop` of recursive iterations, `emitter.emit`-ing every - /// batch. State that must survive across iterations (buffer, reservation, - /// dedup, metrics) lives on `self`, captured by the async body — the two - /// `Option` sub-stream fields and the hand-rolled poll state machine are - /// gone. - fn build( + /// Create a new recursive query stream + fn new( task_context: Arc, work_table: Arc, recursive_term: Arc, static_stream: SendableRecordBatchStream, is_distinct: bool, baseline_metrics: BaselineMetrics, - ) -> Result { + ) -> Result { let schema = static_stream.schema(); let reservation = MemoryConsumer::new("RecursiveQuery").register(task_context.memory_pool()); let distinct_deduplicator = is_distinct .then(|| DistinctDeduplicator::new(Arc::clone(&schema), &task_context)) .transpose()?; - - let state = Self { + Ok(Self { task_context, work_table, recursive_term, + static_stream: Some(static_stream), + recursive_stream: None, + schema, buffer: vec![], reservation, distinct_deduplicator, baseline_metrics, - }; - - let stream = try_fn_stream(|emitter| async move { - let mut state = state; - let mut static_stream = static_stream; - - // Static term: forward and buffer every batch. - while let Some(batch) = static_stream.next().await { - emitter.emit(state.push_batch(batch?)?).await; - } - - // Recursive term: iterate until an iteration produces no new rows. - loop { - let buffered_rows: usize = - state.buffer.iter().map(|b| b.num_rows()).sum(); - if buffered_rows == 0 { - return Ok(()); - } - - // Feed the previous iteration's rows into the work table. - state.work_table.update(ReservedBatches::new( - std::mem::take(&mut state.buffer), - state.reservation.take(), - )); - - // We always (re-)execute partition 0; downstream plans should - // not expect any partitioning. - let recursive_plan = - reset_plan_states(Arc::clone(&state.recursive_term))?; - let mut recursive_stream = - recursive_plan.execute(0, Arc::clone(&state.task_context))?; - while let Some(batch) = recursive_stream.next().await { - emitter.emit(state.push_batch(batch?)?).await; - } - } - }); - - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + }) } - /// Deduplicate (when `DISTINCT`), account for memory, buffer a copy for the - /// next iteration's work table, record output metrics, and return the batch. - fn push_batch(&mut self, mut batch: RecordBatch) -> Result { + /// Push a clone of the given batch to the in memory buffer, and then return + /// a poll with it. + fn push_batch( + mut self: std::pin::Pin<&mut Self>, + mut batch: RecordBatch, + ) -> Poll>> { + let baseline_metrics = self.baseline_metrics.clone(); + if let Some(deduplicator) = &mut self.distinct_deduplicator { - let _timer_guard = self.baseline_metrics.elapsed_compute().timer(); + let _timer_guard = baseline_metrics.elapsed_compute().timer(); batch = deduplicator.deduplicate(&batch)?; } - self.reservation.try_grow(batch.get_array_memory_size())?; + if let Err(e) = self.reservation.try_grow(batch.get_array_memory_size()) { + return Poll::Ready(Some(Err(e))); + } self.buffer.push(batch.clone()); - (&batch).record_output(&self.baseline_metrics); - Ok(batch) + (&batch).record_output(&baseline_metrics); + Poll::Ready(Some(Ok(batch))) + } + + /// Start polling for the next iteration, will be called either after the static term + /// is completed or another term is completed. It will follow the algorithm above on + /// to check whether the recursion has ended. + fn poll_next_iteration( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll>> { + let total_length = self + .buffer + .iter() + .fold(0, |acc, batch| acc + batch.num_rows()); + + if total_length == 0 { + return Poll::Ready(None); + } + + // Update the work table with the current buffer + let reserved_batches = ReservedBatches::new( + std::mem::take(&mut self.buffer), + self.reservation.take(), + ); + self.work_table.update(reserved_batches); + + // We always execute (and re-execute iteratively) the first partition. + // Downstream plans should not expect any partitioning. + let partition = 0; + + let recursive_plan = reset_plan_states(Arc::clone(&self.recursive_term))?; + self.recursive_stream = + Some(recursive_plan.execute(partition, Arc::clone(&self.task_context))?); + self.poll_next(cx) } } @@ -386,6 +387,48 @@ fn assign_work_table( .data() } +impl Stream for RecursiveQueryStream { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + if let Some(static_stream) = &mut self.static_stream { + // While the static term's stream is available, we'll be forwarding the batches from it (also + // saving them for the initial iteration of the recursive term). + let batch_result = ready!(static_stream.poll_next_unpin(cx)); + match &batch_result { + None => { + // Once this is done, we can start running the setup for the recursive term. + self.static_stream = None; + self.poll_next_iteration(cx) + } + Some(Ok(batch)) => self.push_batch(batch.clone()), + _ => Poll::Ready(batch_result), + } + } else if let Some(recursive_stream) = &mut self.recursive_stream { + let batch_result = ready!(recursive_stream.poll_next_unpin(cx)); + match batch_result { + None => { + self.recursive_stream = None; + self.poll_next_iteration(cx) + } + Some(Ok(batch)) => self.push_batch(batch), + _ => Poll::Ready(batch_result), + } + } else { + Poll::Ready(None) + } + } +} + +impl RecordBatchStream for RecursiveQueryStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} /// Deduplicator based on a hash table. struct DistinctDeduplicator { diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..ec767be30e079 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -124,11 +124,6 @@ impl BatchBuilder { self.indices.is_empty() } - /// Returns the schema of this [`BatchBuilder`] - pub fn schema(&self) -> &SchemaRef { - &self.schema - } - /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 84ac9b58dfc6e..5779f75a43e9a 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,13 +18,12 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::fmt::{Debug, Formatter}; +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, RecordOutput}; +use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; @@ -32,16 +31,16 @@ use crate::stream::RecordBatchStreamAdapter; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::Result; use datafusion_execution::memory_pool::MemoryReservation; -use async_fn_stream::{try_fn_stream, TryStreamEmitter}; +use genawaiter::sync::{Co, Gen}; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; /// Merges a stream of sorted cursors and record batches into a single sorted stream -// #[derive(Debug)] +#[derive(Debug)] pub(crate) struct SortPreservingMergeStream { in_progress: BatchBuilder, @@ -140,27 +139,9 @@ 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, - - emitter: TryStreamEmitter, } -// TODO - fix the debug not showing all the properties now -impl Debug for SortPreservingMergeStream { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_struct("SortPreservingMergeStream") - .field("current_reset_epoch", &self.current_reset_epoch) - .field("prev_cursors", &self.prev_cursors) - .field("fetch", &self.fetch) - .field("produced", &self.produced) - .field("uninitiated_partitions", &self.uninitiated_partitions) - .finish() - } -} - -impl SortPreservingMergeStream { +impl SortPreservingMergeStream { pub(crate) fn new( streams: CursorStream, schema: SchemaRef, @@ -169,7 +150,6 @@ impl SortPreservingMergeStream { fetch: Option, reservation: MemoryReservation, enable_round_robin_tie_breaker: bool, - emitter: TryStreamEmitter ) -> Self { let stream_count = streams.partitions(); @@ -188,24 +168,26 @@ impl SortPreservingMergeStream { batch_size, fetch, produced: 0, - uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, - emitter, } } - pub(crate) fn create(streams: CursorStream, - schema: SchemaRef, - metrics: BaselineMetrics, - batch_size: usize, - fetch: Option, - reservation: MemoryReservation, - enable_round_robin_tie_breaker: bool) -> SendableRecordBatchStream where C: 'static { - + pub(crate) fn create( + streams: CursorStream, + schema: SchemaRef, + metrics: BaselineMetrics, + batch_size: usize, + fetch: Option, + reservation: MemoryReservation, + enable_round_robin_tie_breaker: bool, + ) -> SendableRecordBatchStream + where + C: 'static, + { let schema_clone = Arc::clone(&schema); - let stream = try_fn_stream(|emitter| async move { - let s = SortPreservingMergeStream::new( + let stream = Gen::new(|co| async move { + let mut s = SortPreservingMergeStream::new( streams, schema, metrics, @@ -213,10 +195,12 @@ impl SortPreservingMergeStream { fetch, reservation, enable_round_robin_tie_breaker, - emitter ); - s.run().await + // An error aborts the stream: yield it as the final item, then stop. + if let Err(e) = s.run(&co).await { + co.yield_(Err(e)).await; + } }); Box::pin(RecordBatchStreamAdapter::new(schema_clone, stream)) @@ -248,7 +232,7 @@ impl SortPreservingMergeStream { result } - async fn run(mut self) -> Result<()> { + async fn run(&mut self, co: &Co>) -> Result<()> { // Prime the first batch of every partition, then build the loser tree. // With yield support this is a plain sequential loop — no need to track // `uninitiated_partitions` across re-entrant polls, and a `Pending` @@ -285,7 +269,7 @@ impl SortPreservingMergeStream { if self.in_progress.len() >= self.batch_size { if let Some(batch) = self.emit_in_progress_batch()? { drop(timer); - self.emitter.emit(batch).await; + co.yield_(Ok(batch)).await; timer = elapsed_compute.timer(); } } @@ -309,7 +293,7 @@ impl SortPreservingMergeStream { // only a partial batch. Each iteration emits another partial batch until // `in_progress` is fully drained, so nothing is silently dropped. while let Some(batch) = self.emit_in_progress_batch()? { - self.emitter.emit(batch).await; + co.yield_(Ok(batch)).await; } Ok(()) } @@ -567,6 +551,7 @@ mod tests { MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; use std::cmp::Ordering; + use std::task::{Context, Poll}; #[derive(Debug)] struct EmptyPartitionedStream; From d01f8413ed826e39250b431f82bab416042e29ab Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:10:32 +0300 Subject: [PATCH 05/29] fix - emit all and not wait before emitting next --- datafusion/physical-plan/src/sorts/merge.rs | 102 +++++++++++++++++--- 1 file changed, 89 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 5779f75a43e9a..935d90ef831c2 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -21,6 +21,7 @@ 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; @@ -209,22 +210,33 @@ impl SortPreservingMergeStream { /// 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 - async fn maybe_poll_stream(&mut self, idx: usize) -> Result<()> { + fn maybe_poll_stream( + &mut self, + cx: &mut Context<'_>, + idx: usize, + ) -> Poll> { if self.cursors[idx].is_some() { // Cursor is not finished - don't need a new RecordBatch yet - return Ok(()); + return Poll::Ready(Ok(())); } - match poll_fn(|cx| self.streams.poll_next(cx, idx)).await { - None => Ok(()), - Some(Err(e)) => Err(e), + match futures::ready!(self.streams.poll_next(cx, idx)) { + None => Poll::Ready(Ok(())), + Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok((cursor, batch))) => { self.cursors[idx] = Some(Cursor::new(cursor)); - self.in_progress.push_batch(idx, batch) + Poll::Ready(self.in_progress.push_batch(idx, batch)) } } } + /// 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 + async fn poll_stream_and_wait(&mut self, idx: usize) -> Result<()> { + poll_fn(|cx| self.maybe_poll_stream(cx, idx)).await + } + fn emit_in_progress_batch(&mut self) -> Result> { let rows_before = self.in_progress.len(); let result = self.in_progress.build_record_batch(); @@ -234,14 +246,25 @@ impl SortPreservingMergeStream { async fn run(&mut self, co: &Co>) -> Result<()> { // Prime the first batch of every partition, then build the loser tree. - // With yield support this is a plain sequential loop — no need to track - // `uninitiated_partitions` across re-entrant polls, and a `Pending` - // upstream simply suspends here instead of unwinding the whole call. + // Poll *all* not-yet-ready inputs on each wakeup (rather than awaiting + // them one at a time) so their upstream pipelines start concurrently — + // a single sequential `await` would serialize startup, since many + // streams only begin their background work on first poll. // An error aborts the stream immediately; buffered rows are dropped, so // no data is emitted after an `Err`. - for partition_idx in 0..self.cursors.len() { - self.maybe_poll_stream(partition_idx).await?; - } + let mut uninitiated_partitions = + (0..self.streams.partitions()).collect::>(); + + poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) + .await?; + + 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 + drop(uninitiated_partitions); self.init_loser_tree(); // NB timer records time taken on drop; we drop it around `.await`s so @@ -279,7 +302,7 @@ impl SortPreservingMergeStream { let winner = self.loser_tree[0]; if self.cursors[winner].is_none() { drop(timer); - self.maybe_poll_stream(winner).await?; + self.poll_stream_and_wait(winner).await?; timer = elapsed_compute.timer(); } self.update_loser_tree(); @@ -298,6 +321,59 @@ impl SortPreservingMergeStream { Ok(()) } + fn initialize_all_partitions( + &mut self, + uninitiated_partitions: &mut Vec, + cx: &mut Context, + ) -> Poll> { + // 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. + + assert_eq!( + self.loser_tree.len(), + 0, + "loser tree must be empty when initializing" + ); + + // Manual indexing since we're iterating over the vector and shrinking it in the loop + let mut idx = 0; + while idx < uninitiated_partitions.len() { + let partition_idx = uninitiated_partitions[idx]; + match self.maybe_poll_stream(cx, partition_idx) { + Poll::Ready(Err(e)) => { + return Poll::Ready(Err(e)); + } + Poll::Pending => { + // The polled stream is pending which means we're already set up to + // be woken when necessary + // Try the next stream + idx += 1; + } + _ => { + // The polled stream is ready + // Remove it from uninitiated_partitions + // Don't bump idx here, since a new element will have taken its + // 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. + uninitiated_partitions.swap_remove(idx); + } + } + } + + if uninitiated_partitions.is_empty() { + Poll::Ready(Ok(())) + } else { + // There are still uninitiated partitions so return pending. + // We only get here if we've polled all uninitiated streams and at least one of them + // returned pending itself. That means we will be woken as soon as one of the + // streams would like to be polled again. + // There is no need to reschedule ourselves eagerly. + Poll::Pending + } + } + /// 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) { From ffd603e3bdcfcbe2bbeebde8d4cc17c1d66bb769 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:14:56 +0300 Subject: [PATCH 06/29] c --- datafusion/physical-plan/src/sorts/merge.rs | 40 ++++++++------------- 1 file changed, 14 insertions(+), 26 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 935d90ef831c2..0d2cd6d706fb7 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -198,7 +198,6 @@ impl SortPreservingMergeStream { enable_round_robin_tie_breaker, ); - // An error aborts the stream: yield it as the final item, then stop. if let Err(e) = s.run(&co).await { co.yield_(Err(e)).await; } @@ -230,9 +229,6 @@ impl SortPreservingMergeStream { } } - /// 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 async fn poll_stream_and_wait(&mut self, idx: usize) -> Result<()> { poll_fn(|cx| self.maybe_poll_stream(cx, idx)).await } @@ -245,13 +241,6 @@ impl SortPreservingMergeStream { } async fn run(&mut self, co: &Co>) -> Result<()> { - // Prime the first batch of every partition, then build the loser tree. - // Poll *all* not-yet-ready inputs on each wakeup (rather than awaiting - // them one at a time) so their upstream pipelines start concurrently — - // a single sequential `await` would serialize startup, since many - // streams only begin their background work on first poll. - // An error aborts the stream immediately; buffered rows are dropped, so - // no data is emitted after an `Err`. let mut uninitiated_partitions = (0..self.streams.partitions()).collect::>(); @@ -267,24 +256,19 @@ impl SortPreservingMergeStream { drop(uninitiated_partitions); self.init_loser_tree(); - // NB timer records time taken on drop; we drop it around `.await`s so - // upstream/consumer wait time is not counted as compute. + // 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(); loop { - // `loser_tree[0]` is the current, already-adjusted winner. let stream_idx = self.loser_tree[0]; if !self.advance_cursors(stream_idx) { - // The overall minimum is exhausted, which only happens once - // every input is exhausted. We're done. break; } self.in_progress.push_row(stream_idx); - // Stop sorting if the fetch limit has been reached. Remaining - // buffered rows (e.g. a partial batch left after overflow handling) - // are drained below. + // stop sorting if fetch has been reached if self.fetch_reached() { break; } @@ -297,9 +281,11 @@ impl SortPreservingMergeStream { } } - // Refill the just-consumed winner (only if its cursor is exhausted — - // live for almost every row) and re-establish the tree top. + // Adjust the loser tree if necessary, returning control if needed 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() { drop(timer); self.poll_stream_and_wait(winner).await?; @@ -310,11 +296,13 @@ impl SortPreservingMergeStream { drop(timer); - // Drain any buffered rows: the normal final partial batch, plus leftovers - // kept after a `fetch` stop or after `build_record_batch()` hit an i32 - // offset overflow (e.g. combined string offsets exceed 2 GB) and emitted - // only a partial batch. Each iteration emits another partial batch until - // `in_progress` is fully drained, so nothing is silently dropped. + // 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. while let Some(batch) = self.emit_in_progress_batch()? { co.yield_(Ok(batch)).await; } From e6df5f57e9a7360483dae36e15efed562636c5e5 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:18:32 +0300 Subject: [PATCH 07/29] add comment --- datafusion/physical-plan/src/sorts/merge.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 0d2cd6d706fb7..0982bfb1b2787 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -241,6 +241,7 @@ impl SortPreservingMergeStream { } async fn run(&mut self, co: &Co>) -> Result<()> { + // This vector contains the indices of the partitions that have not started emitting yet. let mut uninitiated_partitions = (0..self.streams.partitions()).collect::>(); From 847172e71ca76e8475603dfe09c6997a5426a607 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:40:41 +0300 Subject: [PATCH 08/29] add metrics support for adapter --- datafusion/physical-plan/src/sorts/merge.rs | 13 +++-------- datafusion/physical-plan/src/stream.rs | 25 +++++++++++++++++++-- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 0982bfb1b2787..699e5b49ada96 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -83,12 +83,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, @@ -165,7 +159,6 @@ 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, @@ -187,6 +180,8 @@ impl SortPreservingMergeStream { { let schema_clone = Arc::clone(&schema); + let cloned_metrics = metrics.clone(); + let stream = Gen::new(|co| async move { let mut s = SortPreservingMergeStream::new( streams, @@ -203,7 +198,7 @@ impl SortPreservingMergeStream { } }); - Box::pin(RecordBatchStreamAdapter::new(schema_clone, stream)) + Box::pin(RecordBatchStreamAdapter::new(schema_clone, stream).with_metrics(cloned_metrics)) } /// If the stream at the given index is not exhausted, and the last cursor for the @@ -489,7 +484,6 @@ impl SortPreservingMergeStream { } self.loser_tree[cmp_node] = winner; } - self.loser_tree_adjusted = true; } /// Resets the poll count by incrementing the reset epoch. @@ -601,7 +595,6 @@ impl SortPreservingMergeStream { } self.loser_tree[0] = winner; - self.loser_tree_adjusted = true; } } diff --git a/datafusion/physical-plan/src/stream.rs b/datafusion/physical-plan/src/stream.rs index 9d0b964886afd..ea410f56692cf 100644 --- a/datafusion/physical-plan/src/stream.rs +++ b/datafusion/physical-plan/src/stream.rs @@ -416,6 +416,8 @@ pin_project! { // adapter itself is dropped. #[pin] stream: Option, + + metrics: Option, } } @@ -447,8 +449,15 @@ impl RecordBatchStreamAdapter { Self { schema, stream: Some(stream), + metrics: None, } } + + pub fn with_metrics(mut self, baseline_metrics: BaselineMetrics) -> Self { + self.metrics = Some(baseline_metrics); + + self + } } impl std::fmt::Debug for RecordBatchStreamAdapter { @@ -468,7 +477,12 @@ where fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); let Some(inner) = this.stream.as_mut().as_pin_mut() else { - return Poll::Ready(None); + let poll = Poll::Ready(None); + return if let Some(metrics) = &this.metrics { + metrics.record_poll(poll) + } else { + poll + } }; let item = ready!(inner.poll_next(cx)); if item.is_none() { @@ -481,7 +495,14 @@ where *this.stream.as_mut().get_unchecked_mut() = None; } } - Poll::Ready(item) + + let poll = Poll::Ready(item); + + if let Some(metrics) = &this.metrics { + metrics.record_poll(poll) + } else { + poll + } } fn size_hint(&self) -> (usize, Option) { From 0adc2e57cdcea60240df242ccab203d2df56bea2 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:14:44 +0300 Subject: [PATCH 09/29] align with main --- .../src/approx_distinct.rs | 2 ++ .../src/equivalence/properties/mod.rs | 13 +++++++- .../src/equivalence/properties/union.rs | 32 ++++++++++++++++++- datafusion/physical-plan/src/stream.rs | 25 ++------------- datafusion/physical-plan/src/union.rs | 4 +-- .../sqllogictest/test_files/aggregate.slt | 16 ++++++++++ 6 files changed, 64 insertions(+), 28 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index a8dbd8611d857..0e35b47d643c2 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -799,6 +799,7 @@ impl AggregateUDFImpl for ApproxDistinct { | DataType::Utf8View | DataType::Binary | DataType::BinaryView + | DataType::FixedSizeBinary(_) | DataType::LargeBinary => Box::new(HLLAccumulator::new()), DataType::Null => { Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0)))) @@ -868,6 +869,7 @@ fn is_hll_groups_type(data_type: &DataType) -> bool { | DataType::Utf8View | DataType::Binary | DataType::BinaryView + | DataType::FixedSizeBinary(_) | DataType::LargeBinary ) } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index bb74cd1d9c7b3..17c3898fd9c89 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -1314,7 +1314,18 @@ impl EquivalenceProperties { if let (Some(data_type), Some(AcrossPartitions::Uniform(Some(value)))) = (data_type, &mut eq_class.constant) { - *value = value.cast_to(&data_type)?; + match value.cast_to(&data_type) { + Ok(cast_value) => *value = cast_value, + Err(_) => { + // This is optimizer metadata. If a stale constant + // value cannot be represented after schema rewrite, + // drop the constant instead of failing planning. + eq_class.constant = None; + } + } + } + if eq_class.is_trivial() { + continue; } eq_classes.push(eq_class); } diff --git a/datafusion/physical-expr/src/equivalence/properties/union.rs b/datafusion/physical-expr/src/equivalence/properties/union.rs index d77129472a8ba..ea4094e75159a 100644 --- a/datafusion/physical-expr/src/equivalence/properties/union.rs +++ b/datafusion/physical-expr/src/equivalence/properties/union.rs @@ -311,7 +311,7 @@ mod tests { use crate::equivalence::tests::{create_test_schema, parse_sort_expr}; use crate::expressions::col; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datafusion_common::ScalarValue; use itertools::Itertools; @@ -899,6 +899,36 @@ mod tests { Ok(()) } + #[test] + fn test_union_drops_unrepresentable_constant_value_after_schema_rewrite() -> Result<()> + { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "ticker", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )])); + let output_schema = Arc::new(Schema::new(vec![Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )])); + + let ticker = col("ticker", &input_schema)?; + let stale_value = ScalarValue::Utf8(Some("ESU6".to_owned())); + let const_expr = ConstExpr::new( + Arc::clone(&ticker), + AcrossPartitions::Uniform(Some(stale_value)), + ); + + let mut input = EquivalenceProperties::new(input_schema); + input.add_constants(vec![const_expr])?; + + let union_props = calculate_union(vec![input], output_schema)?; + assert!(union_props.constants().is_empty()); + + Ok(()) + } + /// Return a new schema with the same types, but new field names /// /// The new field names are the old field names with `text` appended. diff --git a/datafusion/physical-plan/src/stream.rs b/datafusion/physical-plan/src/stream.rs index ea410f56692cf..9d0b964886afd 100644 --- a/datafusion/physical-plan/src/stream.rs +++ b/datafusion/physical-plan/src/stream.rs @@ -416,8 +416,6 @@ pin_project! { // adapter itself is dropped. #[pin] stream: Option, - - metrics: Option, } } @@ -449,15 +447,8 @@ impl RecordBatchStreamAdapter { Self { schema, stream: Some(stream), - metrics: None, } } - - pub fn with_metrics(mut self, baseline_metrics: BaselineMetrics) -> Self { - self.metrics = Some(baseline_metrics); - - self - } } impl std::fmt::Debug for RecordBatchStreamAdapter { @@ -477,12 +468,7 @@ where fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); let Some(inner) = this.stream.as_mut().as_pin_mut() else { - let poll = Poll::Ready(None); - return if let Some(metrics) = &this.metrics { - metrics.record_poll(poll) - } else { - poll - } + return Poll::Ready(None); }; let item = ready!(inner.poll_next(cx)); if item.is_none() { @@ -495,14 +481,7 @@ where *this.stream.as_mut().get_unchecked_mut() = None; } } - - let poll = Poll::Ready(item); - - if let Some(metrics) = &this.metrics { - metrics.record_poll(poll) - } else { - poll - } + Poll::Ready(item) } fn size_hint(&self) -> (usize, Option) { diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index d6f664c0059bc..b330c305833ff 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -130,9 +130,7 @@ impl UnionExec { // The schema of the inputs and the union schema is consistent when: // - They have the same number of fields, and // - Their fields have same types at the same indices. - // Here, we know that schemas are consistent and the call below can - // not return an error. - let cache = Self::compute_properties(&inputs, schema).unwrap(); + let cache = Self::compute_properties(&inputs, schema)?; Ok(Arc::new(UnionExec { inputs, metrics: ExecutionPlanMetricsSet::new(), diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 37ee3d8a95843..c5970bde9c954 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -1953,6 +1953,22 @@ SELECT g, approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'BinaryView')) FRO 4 1 +# FixedSizeBinary non-grouped +query I +SELECT approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'FixedSizeBinary(1)')) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +# FixedSizeBinary grouped +query II +SELECT g, approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'FixedSizeBinary(1)')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + + # Integers (Int32): group 1 -> {10,20}=2, group 2 -> {30,40}=2, group 3 -> 0, group 4 -> {50}=1 query II SELECT g, approx_distinct(i) FROM approx_distinct_group_test GROUP BY g ORDER BY g; From f438d7b5661f642b011fda559d76b8be1be631f6 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:14:57 +0300 Subject: [PATCH 10/29] use observed stream --- datafusion/physical-plan/src/sorts/merge.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 699e5b49ada96..59c4625763de0 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -28,7 +28,7 @@ use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; -use crate::stream::RecordBatchStreamAdapter; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -198,7 +198,9 @@ impl SortPreservingMergeStream { } }); - Box::pin(RecordBatchStreamAdapter::new(schema_clone, stream).with_metrics(cloned_metrics)) + let stream = Box::pin(RecordBatchStreamAdapter::new(schema_clone, 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 From 6171c127489fa5f0914276764dc05d41acb95342 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:36:29 +0300 Subject: [PATCH 11/29] uncomment pr --- datafusion/physical-plan/src/sorts/merge.rs | 70 ++++++++++----------- 1 file changed, 34 insertions(+), 36 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 59c4625763de0..feac476eb2e26 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -653,40 +653,38 @@ mod tests { } } - // TODO - uncomment this test - // - // /// The merge loop's terminal drain (`while let Some(b) = emit()? { send }`) - // /// relies on `emit_in_progress_batch` returning buffered rows and then - // /// `None` once empty. This exercises that primitive directly. - // #[test] - // fn test_emit_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); - // let metrics = ExecutionPlanMetricsSet::new(); - // - // let mut stream = SortPreservingMergeStream::::create( - // Box::new(EmptyPartitionedStream), - // Arc::clone(&schema), - // BaselineMetrics::new(&metrics, 0), - // 16, - // Some(1), - // reservation, - // true, - // ); - // - // 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); - // - // let batch = stream - // .emit_in_progress_batch() - // .unwrap() - // .expect("buffered row should be emitted"); - // assert_eq!(batch.num_rows(), 1); - // assert!(stream.in_progress.is_empty()); - // assert!(stream.emit_in_progress_batch().unwrap().is_none()); - // } + /// The merge loop's terminal drain (`while let Some(b) = emit()? { send }`) + /// relies on `emit_in_progress_batch` returning buffered rows and then + /// `None` once empty. This exercises that primitive directly. + #[test] + fn test_emit_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); + let metrics = ExecutionPlanMetricsSet::new(); + + let mut stream = SortPreservingMergeStream::::new( + Box::new(EmptyPartitionedStream), + Arc::clone(&schema), + BaselineMetrics::new(&metrics, 0), + 16, + Some(1), + reservation, + true, + ); + + 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); + + let batch = stream + .emit_in_progress_batch() + .unwrap() + .expect("buffered row should be emitted"); + assert_eq!(batch.num_rows(), 1); + assert!(stream.in_progress.is_empty()); + assert!(stream.emit_in_progress_batch().unwrap().is_none()); + } } From e56c139951153ea6725bf3685edee2b5db026c83 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:41:08 +0300 Subject: [PATCH 12/29] add into stream --- datafusion/physical-plan/src/sorts/builder.rs | 5 +++ datafusion/physical-plan/src/sorts/merge.rs | 39 ++++++------------- .../src/sorts/streaming_merge.rs | 8 ++-- 3 files changed, 20 insertions(+), 32 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index ec767be30e079..75eb2ff980325 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -124,6 +124,11 @@ impl BatchBuilder { self.indices.is_empty() } + /// Returns the schema of this [`BatchBuilder`] + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + /// Try to interleave all columns using the given index slice. fn try_interleave_columns( &self, diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index feac476eb2e26..0b12532fac9b1 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -166,34 +166,17 @@ impl SortPreservingMergeStream { } } - pub(crate) fn create( - streams: CursorStream, - schema: SchemaRef, - metrics: BaselineMetrics, - batch_size: usize, - fetch: Option, - reservation: MemoryReservation, - enable_round_robin_tie_breaker: bool, - ) -> SendableRecordBatchStream + pub(crate) fn into_stream(self) -> SendableRecordBatchStream where C: 'static, { - let schema_clone = Arc::clone(&schema); + let schema_clone = Arc::clone(self.in_progress.schema()); - let cloned_metrics = metrics.clone(); + let cloned_metrics = self.metrics.clone(); + let mut this = self; let stream = Gen::new(|co| async move { - let mut s = SortPreservingMergeStream::new( - streams, - schema, - metrics, - batch_size, - fetch, - reservation, - enable_round_robin_tie_breaker, - ); - - if let Err(e) = s.run(&co).await { + if let Err(e) = this.run(&co).await { co.yield_(Err(e)).await; } }); @@ -271,12 +254,12 @@ impl SortPreservingMergeStream { break; } - if self.in_progress.len() >= self.batch_size { - if let Some(batch) = self.emit_in_progress_batch()? { - drop(timer); - co.yield_(Ok(batch)).await; - timer = elapsed_compute.timer(); - } + if self.in_progress.len() >= self.batch_size + && let Some(batch) = self.emit_in_progress_batch()? + { + drop(timer); + co.yield_(Ok(batch)).await; + timer = elapsed_compute.timer(); } // Adjust the loser tree if necessary, returning control if needed diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 2e83662994f53..6120c918e2582 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(SortPreservingMergeStream::create( + return Ok(SortPreservingMergeStream::new( Box::new(streams), $schema, $tracking_metrics, @@ -54,7 +54,7 @@ macro_rules! merge_helper { $fetch, $reservation, $enable_round_robin_tie_breaker, - )); + ).into_stream()); }}; } @@ -254,7 +254,7 @@ impl<'a> StreamingMergeBuilder<'a> { streams, reservation.new_empty(), )?; - Ok(SortPreservingMergeStream::create( + Ok(SortPreservingMergeStream::new( Box::new(streams), schema, metrics, @@ -262,6 +262,6 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, reservation, enable_round_robin_tie_breaker, - )) + ).into_stream()) } } From 72f3513db27e51fc36b9c946ad0337944776c0bb Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:42:34 +0300 Subject: [PATCH 13/29] format --- datafusion/physical-plan/src/sorts/streaming_merge.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 6120c918e2582..e96138ef1306c 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -54,7 +54,8 @@ macro_rules! merge_helper { $fetch, $reservation, $enable_round_robin_tie_breaker, - ).into_stream()); + ) + .into_stream()); }}; } @@ -262,6 +263,7 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, reservation, enable_round_robin_tie_breaker, - ).into_stream()) + ) + .into_stream()) } } From 03441d00cc9012110ad7b721f3a732447f9a1238 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:43:35 +0300 Subject: [PATCH 14/29] revert test to main --- datafusion/physical-plan/src/sorts/merge.rs | 25 ++++++++++++--------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 0b12532fac9b1..8b7071670a3c6 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -593,8 +593,8 @@ mod tests { use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; + use futures::task::noop_waker_ref; use std::cmp::Ordering; - use std::task::{Context, Poll}; #[derive(Debug)] struct EmptyPartitionedStream; @@ -636,11 +636,8 @@ mod tests { } } - /// The merge loop's terminal drain (`while let Some(b) = emit()? { send }`) - /// relies on `emit_in_progress_batch` returning buffered rows and then - /// `None` once empty. This exercises that primitive directly. #[test] - fn test_emit_drains_buffered_rows() { + 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); @@ -661,13 +658,19 @@ mod tests { .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 batch = stream - .emit_in_progress_batch() - .unwrap() - .expect("buffered row should be emitted"); - assert_eq!(batch.num_rows(), 1); + let waker = noop_waker_ref(); + let mut cx = Context::from_waker(waker); + + 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!(stream.emit_in_progress_batch().unwrap().is_none()); + assert!(matches!(stream.poll_next_inner(&mut cx), Poll::Ready(None))); } } From 3e5ab00d21af05a2720978e3b91a4bf1d676672c Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:50:47 +0300 Subject: [PATCH 15/29] fix test --- datafusion/physical-plan/src/sorts/merge.rs | 26 +++++++++------------ 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 8b7071670a3c6..5e5776b299dce 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -593,7 +593,7 @@ mod tests { use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; - use futures::task::noop_waker_ref; + use futures::TryStreamExt; use std::cmp::Ordering; #[derive(Debug)] @@ -636,8 +636,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); @@ -653,24 +653,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); } } From 3182a97c68688f24288bb6553c04b6f8f5dd48d2 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:34:19 +0300 Subject: [PATCH 16/29] update with main --- datafusion/physical-plan/src/sorts/merge.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 5e5776b299dce..bb08cb044adc1 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -290,15 +290,15 @@ impl SortPreservingMergeStream { 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, uninitiated_partitions: &mut Vec, cx: &mut Context, ) -> Poll> { - // 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. - assert_eq!( self.loser_tree.len(), 0, From 48aab960aa176f51fd9db89511b8508438294d2f Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:44:14 +0300 Subject: [PATCH 17/29] move comment --- datafusion/physical-plan/src/sorts/merge.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bb08cb044adc1..030cfafa42290 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -209,10 +209,6 @@ impl SortPreservingMergeStream { } } - async fn poll_stream_and_wait(&mut self, idx: usize) -> Result<()> { - poll_fn(|cx| self.maybe_poll_stream(cx, idx)).await - } - fn emit_in_progress_batch(&mut self) -> Result> { let rows_before = self.in_progress.len(); let result = self.in_progress.build_record_batch(); @@ -262,16 +258,17 @@ impl SortPreservingMergeStream { timer = elapsed_compute.timer(); } - // Adjust the loser tree if necessary, returning control if needed 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() { drop(timer); - self.poll_stream_and_wait(winner).await?; + poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; timer = elapsed_compute.timer(); } + + // Adjusting the loser tree if necessary self.update_loser_tree(); } From 039d07324a807256c01206ce4583622355769e4f Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 9 Jul 2026 21:04:17 +0300 Subject: [PATCH 18/29] format `Cargo.toml` --- Cargo.toml | 2 +- datafusion/physical-plan/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index c847b76df2a7b..2d93292c9de89 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,7 +113,6 @@ arrow-ipc = { version = "59.1.0", default-features = false, features = [ ] } arrow-ord = { version = "59.1.0", default-features = false } arrow-schema = { version = "59.1.0", default-features = false } -genawaiter = {version = "0.99.1", default-features = false} async-trait = "0.1.89" bigdecimal = "0.4.8" bytes = "1.11" @@ -160,6 +159,7 @@ datafusion-session = { path = "datafusion/session", version = "54.0.0" } datafusion-spark = { path = "datafusion/spark", version = "54.0.0" } datafusion-sql = { path = "datafusion/sql", version = "54.0.0" } datafusion-substrait = { path = "datafusion/substrait", version = "54.0.0" } +genawaiter = { version = "0.99.1", default-features = false } doc-comment = "0.3" env_logger = "0.11" diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 343c15feedcb1..36e4b27806b4e 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -64,7 +64,6 @@ arrow-data = { workspace = true } arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } -genawaiter = { workspace = true, features = ["futures03"] } async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true } @@ -79,6 +78,7 @@ datafusion-physical-expr-common = { workspace = true } datafusion-proto-common = { workspace = true, optional = true } datafusion-proto-models = { workspace = true, optional = true } futures = { workspace = true } +genawaiter = { workspace = true, features = ["futures03"] } half = { workspace = true } hashbrown = { workspace = true } indexmap = { workspace = true } From dfd12265fcba818ee97fc2477fa3da8834eb9e36 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 12 Jul 2026 16:55:30 +0300 Subject: [PATCH 19/29] add passthrough last stream optimization --- datafusion/physical-plan/src/coalesce/mod.rs | 7 ++ .../physical-plan/src/coalesce_batches.rs | 34 +++++-- datafusion/physical-plan/src/sorts/builder.rs | 9 ++ datafusion/physical-plan/src/sorts/merge.rs | 95 ++++++++++++++++++- datafusion/physical-plan/src/sorts/stream.rs | 18 ++++ 5 files changed, 149 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index ea1a87d091481..d610cf14dbb25 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -68,6 +68,13 @@ 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 59b3138b55430..3cbe36910b4df 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(), + Ok(Box::pin(CoalesceBatchesStream::new( + self.input.execute(partition, context)?, self.target_batch_size, self.fetch, - ), - baseline_metrics: BaselineMetrics::new(&self.metrics, partition), - completed: false, - })) + 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,26 @@ 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..7561428a0bda0 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -114,6 +114,15 @@ 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() diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 030cfafa42290..e518ae43d06bf 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,11 +18,12 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. +use std::cmp::Ordering; use std::fmt::Debug; use std::future::poll_fn; use std::sync::Arc; use std::task::{Context, Poll}; - +use arrow::compute::BatchCoalescer; use crate::SendableRecordBatchStream; use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; @@ -32,10 +33,13 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use futures::{StreamExt, TryStreamExt}; use datafusion_common::Result; use datafusion_execution::memory_pool::MemoryReservation; use genawaiter::sync::{Co, Gen}; +use crate::coalesce::LimitedBatchCoalescer; +use crate::coalesce_batches::CoalesceBatchesStream; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -134,6 +138,9 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, + + number_of_exhausted_streams: usize, + is_exhausted: Vec, } impl SortPreservingMergeStream { @@ -163,6 +170,8 @@ impl SortPreservingMergeStream { fetch, produced: 0, enable_round_robin_tie_breaker, + number_of_exhausted_streams: 0, + is_exhausted: vec![false; stream_count], } } @@ -200,7 +209,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)); @@ -238,7 +253,8 @@ impl SortPreservingMergeStream { let elapsed_compute = self.metrics.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); - 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; @@ -265,6 +281,8 @@ impl SortPreservingMergeStream { if self.cursors[winner].is_none() { drop(timer); poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; + + // TODO - check if only 1 stream left timer = elapsed_compute.timer(); } @@ -274,6 +292,21 @@ impl SortPreservingMergeStream { drop(timer); + let last_stream_idx = self.loser_tree[0]; + + self.passthrough_last_stream(co, last_stream_idx).await?; + + Ok(()) + } + + async fn passthrough_last_stream(&mut self, co: &Co>, 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; + // 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`. @@ -281,9 +314,54 @@ 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()? { - co.yield_(Ok(batch)).await; + while !self.in_progress.is_empty() { + // 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); + co.yield_(Ok(last_batch)).await; + timer = elapsed_compute.timer(); + } + + 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); + co.yield_(Ok(last_batch)).await; + timer = elapsed_compute.timer(); + } + + 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 + }; + + // TODO - Add memory accounting + let mut coalescer = CoalesceBatchesStream::new(last_stream, self.batch_size, self.fetch.map(|x| x - self.produced), self.metrics.clone()) + // Use batch size as the coalescer + .with_biggest_coalesce_batch_size(None); + + drop(timer); + while let Some(batch) = coalescer.next().await { + co.yield_(batch).await; + } + Ok(()) } @@ -592,6 +670,8 @@ mod tests { }; use futures::TryStreamExt; use std::cmp::Ordering; + use futures::stream::Fuse; + use crate::EmptyRecordBatchStream; #[derive(Debug)] struct EmptyPartitionedStream; @@ -603,6 +683,11 @@ 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..a4d8111e00aab 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -47,6 +47,8 @@ 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 +69,14 @@ 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 +219,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 +293,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<'_>, From b9a5d058fcab30aabe5655deae5f34e97a438268 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 12 Jul 2026 17:39:21 +0300 Subject: [PATCH 20/29] add batch wise skip --- datafusion/physical-plan/src/sorts/builder.rs | 21 ++++- datafusion/physical-plan/src/sorts/cursor.rs | 60 +++++++++++++- datafusion/physical-plan/src/sorts/merge.rs | 82 +++++++++++++++++-- .../src/sorts/streaming_merge.rs | 5 ++ 4 files changed, 155 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 7561428a0bda0..7ad4f5b59c99f 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -214,14 +214,27 @@ impl BatchBuilder { /// /// Returns `None` if no pending rows pub fn build_record_batch(&mut self) -> Result> { - if self.is_empty() { + self.build_record_batch_with_up_to_n_rows(self.indices.len()) + } + + /// Drains the in_progress row indexes up to n rows, 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(crate) fn build_record_batch_with_up_to_n_rows(&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| { - self.try_interleave_columns(&self.indices[..rows_to_emit]) - })?; + retry_interleave(batch_size, batch_size, |rows_to_emit| { + self.try_interleave_columns(&self.indices[..rows_to_emit]) + })?; Ok(Some(self.finish_record_batch(rows_to_emit, columns)?)) } diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index ff2753c19bc18..a3d23fa252e04 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). @@ -100,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). @@ -117,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() @@ -126,6 +153,10 @@ impl Cursor { false } } + + pub(crate) fn len(&self) -> usize { + self.values.len() - self.offset + } } impl PartialEq for Cursor { @@ -291,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 @@ -520,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 e518ae43d06bf..4da3c460c1a8a 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -226,7 +226,10 @@ 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_with_up_to_n_rows(rows_to_emit); self.produced += rows_before - self.in_progress.len(); result } @@ -282,12 +285,35 @@ impl SortPreservingMergeStream { drop(timer); poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; - // TODO - check if only 1 stream left 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(); } - // Adjusting the loser tree if necessary - self.update_loser_tree(); } drop(timer); @@ -418,6 +444,48 @@ 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) { @@ -624,15 +692,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; } 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(); From b744378d737431dbdb01a1160d0e5b9cc53b1d90 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 12 Jul 2026 17:46:50 +0300 Subject: [PATCH 21/29] fix condition --- datafusion/physical-plan/src/sorts/merge.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4da3c460c1a8a..43e38c61a3fa8 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -257,7 +257,7 @@ impl SortPreservingMergeStream { 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() { + while self.number_of_exhausted_streams + 1 < self.streams.partitions() { let stream_idx = self.loser_tree[0]; if !self.advance_cursors(stream_idx) { break; From 10c1133a87dc9b73f6c304fd7cf693f775ef9fb0 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Sun, 12 Jul 2026 19:41:22 +0300 Subject: [PATCH 22/29] fix test --- datafusion/physical-plan/src/sorts/merge.rs | 29 ++++++++++++++++----- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 43e38c61a3fa8..835bd85c6524c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -25,7 +25,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use arrow::compute::BatchCoalescer; use crate::SendableRecordBatchStream; -use crate::metrics::BaselineMetrics; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; @@ -316,10 +316,23 @@ impl SortPreservingMergeStream { } - drop(timer); - 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(co, last_stream_idx).await?; Ok(()) @@ -378,9 +391,13 @@ impl SortPreservingMergeStream { last_stream }; - // TODO - Add memory accounting - let mut coalescer = CoalesceBatchesStream::new(last_stream, self.batch_size, self.fetch.map(|x| x - self.produced), self.metrics.clone()) - // Use batch size as the coalescer + 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); From 6063261d3f8726072183a608b023e334c5aefcf2 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 12:09:23 +0300 Subject: [PATCH 23/29] fix infinite loop --- datafusion/physical-plan/src/coalesce/mod.rs | 9 +- .../physical-plan/src/coalesce_batches.rs | 20 +++- datafusion/physical-plan/src/sorts/builder.rs | 14 ++- datafusion/physical-plan/src/sorts/merge.rs | 106 +++++++++++------- datafusion/physical-plan/src/sorts/stream.rs | 14 ++- 5 files changed, 105 insertions(+), 58 deletions(-) diff --git a/datafusion/physical-plan/src/coalesce/mod.rs b/datafusion/physical-plan/src/coalesce/mod.rs index d610cf14dbb25..2d8883e9c29c9 100644 --- a/datafusion/physical-plan/src/coalesce/mod.rs +++ b/datafusion/physical-plan/src/coalesce/mod.rs @@ -68,9 +68,14 @@ impl LimitedBatchCoalescer { } } - pub fn with_biggest_coalesce_batch_size(self, biggest_coalesce_batch_size: Option) -> Self { + 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), + inner: self + .inner + .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size), ..self } } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 70c5aa5b0b16c..e35572c8b9f5f 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -202,8 +202,8 @@ impl ExecutionPlan for CoalesceBatchesExec { ) -> Result { Ok(Box::pin(CoalesceBatchesStream::new( self.input.execute(partition, context)?, - self.target_batch_size, - self.fetch, + self.target_batch_size, + self.fetch, BaselineMetrics::new(&self.metrics, partition), ))) } @@ -312,7 +312,12 @@ impl Stream for CoalesceBatchesStream { } impl CoalesceBatchesStream { - pub(crate) fn new(input: SendableRecordBatchStream, target_batch_size: usize, fetch: Option, baseline_metrics: BaselineMetrics) -> Self { + pub(crate) fn new( + input: SendableRecordBatchStream, + target_batch_size: usize, + fetch: Option, + baseline_metrics: BaselineMetrics, + ) -> Self { CoalesceBatchesStream { coalescer: LimitedBatchCoalescer::new( input.schema(), @@ -325,9 +330,14 @@ impl CoalesceBatchesStream { } } - pub(crate) fn with_biggest_coalesce_batch_size(self, biggest_coalesce_batch_size: Option) -> Self { + 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), + coalescer: self + .coalescer + .with_biggest_coalesce_batch_size(biggest_coalesce_batch_size), ..self } } diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 7ad4f5b59c99f..744cf032d2ff7 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -120,7 +120,8 @@ impl BatchBuilder { let row_idx = cursor.row_idx; cursor.row_idx += n; - self.indices.extend((0..n).map(|i| (cursor.batch_idx, row_idx + i))); + self.indices + .extend((0..n).map(|i| (cursor.batch_idx, row_idx + i))); } /// Returns the number of in-progress rows in this [`BatchBuilder`] @@ -224,7 +225,10 @@ impl BatchBuilder { /// retries with progressively fewer rows until it succeeds. /// /// Returns `None` if no pending rows - pub(crate) fn build_record_batch_with_up_to_n_rows(&mut self, n: usize) -> Result> { + pub(crate) fn build_record_batch_with_up_to_n_rows( + &mut self, + n: usize, + ) -> Result> { if self.is_empty() || n == 0 { return Ok(None); } @@ -232,9 +236,9 @@ impl BatchBuilder { let batch_size = self.indices.len().min(n); let (rows_to_emit, columns) = - retry_interleave(batch_size, batch_size, |rows_to_emit| { - self.try_interleave_columns(&self.indices[..rows_to_emit]) - })?; + retry_interleave(batch_size, batch_size, |rows_to_emit| { + self.try_interleave_columns(&self.indices[..rows_to_emit]) + })?; Ok(Some(self.finish_record_batch(rows_to_emit, columns)?)) } diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 835bd85c6524c..a2f6c4d8d6187 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,28 +18,26 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::cmp::Ordering; -use std::fmt::Debug; -use std::future::poll_fn; -use std::sync::Arc; -use std::task::{Context, Poll}; -use arrow::compute::BatchCoalescer; use crate::SendableRecordBatchStream; -use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; +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 futures::{StreamExt, TryStreamExt}; use datafusion_common::Result; use datafusion_execution::memory_pool::MemoryReservation; +use futures::StreamExt; -use genawaiter::sync::{Co, Gen}; -use crate::coalesce::LimitedBatchCoalescer; use crate::coalesce_batches::CoalesceBatchesStream; +use genawaiter::sync::{Co, Gen}; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -215,7 +213,7 @@ impl SortPreservingMergeStream { 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)); @@ -228,8 +226,11 @@ impl SortPreservingMergeStream { let rows_before = self.in_progress.len(); // 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_with_up_to_n_rows(rows_to_emit); + let rows_to_emit = + (self.fetch.unwrap_or(usize::MAX) - self.produced).min(self.batch_size); + let result = self + .in_progress + .build_record_batch_with_up_to_n_rows(rows_to_emit); self.produced += rows_before - self.in_progress.len(); result } @@ -297,14 +298,22 @@ impl SortPreservingMergeStream { // 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"); + 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); + 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); @@ -313,7 +322,6 @@ impl SortPreservingMergeStream { // Adjusting the loser tree if necessary self.update_loser_tree(); } - } let last_stream_idx = self.loser_tree[0]; @@ -323,7 +331,7 @@ impl SortPreservingMergeStream { let mut remaining = cursor.len(); if let Some(fetch) = self.fetch { remaining = remaining - .min(fetch.saturating_sub(self.produced + self.in_progress.len())); + .min(fetch.saturating_sub(self.produced + self.in_progress.len())); } if remaining > 0 { self.in_progress.push_n_rows(last_stream_idx, remaining); @@ -333,12 +341,18 @@ impl SortPreservingMergeStream { drop(timer); - self.passthrough_last_stream(co, last_stream_idx).await?; + if self.fetch.is_none_or(|fetch| fetch > self.produced) { + self.passthrough_last_stream(co, last_stream_idx).await?; + } Ok(()) } - async fn passthrough_last_stream(&mut self, co: &Co>, last_stream_index: usize) -> Result<()> { + async fn passthrough_last_stream( + &mut self, + co: &Co>, + last_stream_index: usize, + ) -> Result<()> { let elapsed_compute = self.metrics.elapsed_compute().clone(); let mut timer = elapsed_compute.timer(); @@ -346,14 +360,10 @@ impl SortPreservingMergeStream { let mut last_batch = None; - // 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. - while !self.in_progress.is_empty() { + // 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 @@ -363,14 +373,23 @@ impl SortPreservingMergeStream { 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 last_stream.is_done() || self.fetch.is_some_and(|fetch| fetch <= self.produced) + { if let Some(last_batch) = last_batch.take() { drop(timer); co.yield_(Ok(last_batch)).await; - timer = elapsed_compute.timer(); + + // Not creating a timer since we are returning right away } return Ok(()); @@ -381,12 +400,14 @@ impl SortPreservingMergeStream { 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 + Box::pin(RecordBatchStreamAdapter::new( + schema, + stream.chain(last_stream).boxed(), + )) as SendableRecordBatchStream } else { last_stream }; @@ -395,10 +416,10 @@ impl SortPreservingMergeStream { last_stream, self.batch_size, self.fetch.map(|x| x - self.produced), - self.metrics.intermediate() + self.metrics.intermediate(), ) - // Don't allow for passthrough of batches with sizes other than the provided batch size - .with_biggest_coalesce_batch_size(None); + // 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 { @@ -461,8 +482,6 @@ 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. @@ -746,6 +765,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; @@ -754,9 +774,8 @@ mod tests { MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; use futures::TryStreamExt; - use std::cmp::Ordering; use futures::stream::Fuse; - use crate::EmptyRecordBatchStream; + use std::cmp::Ordering; #[derive(Debug)] struct EmptyPartitionedStream; @@ -768,9 +787,14 @@ mod tests { 1 } - fn take_partition(&mut self, partition_idx: usize) -> Fuse { + 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() + (Box::pin(EmptyRecordBatchStream::new(Arc::new(Schema::empty()))) + as SendableRecordBatchStream) + .fuse() } fn poll_next( diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index a4d8111e00aab..ac4ca1aee1d5c 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -47,7 +47,8 @@ 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 take_partition(&mut self, partition_idx: usize) + -> Fuse; fn poll_next( &mut self, @@ -69,11 +70,14 @@ impl std::fmt::Debug for FusedStreams { } impl FusedStreams { - fn take_stream(&mut self, stream_idx: usize) -> Fuse{ + 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(); + let empty_stream: Fuse = + (Box::pin(EmptyRecordBatchStream::new(stream_schema)) + as SendableRecordBatchStream) + .fuse(); mem::replace(&mut self.0[stream_idx], empty_stream) } @@ -219,7 +223,7 @@ impl PartitionedStream for RowCursorStream { self.streams.0.len() } - fn take_partition(&mut self, stream_idx: usize) -> Fuse{ + fn take_partition(&mut self, stream_idx: usize) -> Fuse { self.streams.take_stream(stream_idx) } @@ -293,7 +297,7 @@ impl PartitionedStream for FieldCursorStream { self.streams.0.len() } - fn take_partition(&mut self, stream_idx: usize) -> Fuse{ + fn take_partition(&mut self, stream_idx: usize) -> Fuse { self.streams.take_stream(stream_idx) } From c435c46a249e61fc822d40b8a4b53e42bd56db20 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:06:14 +0300 Subject: [PATCH 24/29] fix too much splitting batch size --- .../src/sorts/multi_level_merge.rs | 96 ++++++++++++------- 1 file changed, 63 insertions(+), 33 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 4d108ac046eb0..b7fd6066e6a1d 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -138,7 +138,7 @@ use futures::{Stream, StreamExt}; pub(crate) struct MultiLevelMergeBuilder { spill_manager: SpillManager, schema: SchemaRef, - sorted_spill_files: Vec, + sorted_spill_files: Vec<(SortedSpillFile, usize)>, sorted_streams: Vec, expr: LexOrdering, metrics: BaselineMetrics, @@ -171,7 +171,7 @@ impl MultiLevelMergeBuilder { Self { spill_manager, schema, - sorted_spill_files, + sorted_spill_files: sorted_spill_files.into_iter().map(|s| (s, batch_size)).collect(), sorted_streams, expr, metrics, @@ -191,8 +191,11 @@ impl MultiLevelMergeBuilder { async fn create_stream(mut self) -> Result { loop { - let mut stream = match self.merge_sorted_runs_within_mem_limit()? { - MergeStep::Stream(stream) => stream, + let (mut stream, batch_size) = match self.merge_sorted_runs_within_mem_limit()? { + MergeStep::Stream { + stream, + batch_size + } => (stream, batch_size), MergeStep::SplitThenRetry(index) => { // Couldn't reserve memory for the minimum of 2 streams. Re-spill the // larger of the two we're trying to merge with half its batch size so @@ -230,10 +233,10 @@ impl MultiLevelMergeBuilder { }; // Add the spill file - self.sorted_spill_files.push(SortedSpillFile { + self.sorted_spill_files.push((SortedSpillFile { file: spill_file, max_record_batch_memory, - }); + }, batch_size)); } } @@ -245,37 +248,50 @@ impl MultiLevelMergeBuilder { (0, 0) => { let empty_stream = Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); - Ok(MergeStep::Stream(self.observe_output(empty_stream))) + Ok(MergeStep::Stream { + stream: self.observe_output(empty_stream), + batch_size: self.batch_size, + }) } // Only in-memory stream, return that (0, 1) => { let output_stream = self.sorted_streams.remove(0); - Ok(MergeStep::Stream(self.observe_output(output_stream))) + Ok(MergeStep::Stream { + stream: self.observe_output(output_stream), + batch_size: self.batch_size, + }) } // Only single sorted spill file so return it (1, 0) => { - let spill_file = self.sorted_spill_files.remove(0); + let (spill_file, batch_size) = self.sorted_spill_files.remove(0); // Not reserving any memory for this disk as we are not holding it in memory let output_stream = self .spill_manager .read_spill_as_stream(spill_file.file, None)?; - Ok(MergeStep::Stream(self.observe_output(output_stream))) + Ok(MergeStep::Stream { + stream: self.observe_output(output_stream), + batch_size, + }) } // Only in memory streams, so merge them all in a single pass (0, _) => { let sorted_stream = mem::take(&mut self.sorted_streams); // No need to wrap with observed stream since merge sort will update the observed metrics - Ok(MergeStep::Stream(self.create_new_merge_sort( - sorted_stream, - // If we have no sorted spill files left, this is the last run - true, - true, - )?)) + Ok(MergeStep::Stream { + stream: self.create_new_merge_sort( + sorted_stream, + // If we have no sorted spill files left, this is the last run + true, + true, + self.batch_size, + )?, + batch_size: self.batch_size, + }) } // Need to merge multiple streams @@ -326,7 +342,9 @@ impl MultiLevelMergeBuilder { mem::swap(&mut self.reservation, &mut memory_reservation); } - for spill in sorted_spill_files { + let mut min_batch_size = self.batch_size; + + for (spill, batch_size) in sorted_spill_files { let stream = self .spill_manager .clone() @@ -336,12 +354,14 @@ impl MultiLevelMergeBuilder { Some(spill.max_record_batch_memory), )?; sorted_streams.push(stream); + min_batch_size = min_batch_size.min(batch_size); } let merge_sort_stream = self.create_new_merge_sort( sorted_streams, // If we have no sorted spill files left, this is the last run self.sorted_spill_files.is_empty(), is_only_merging_memory_streams, + min_batch_size, )?; // If we're only merging memory streams, we don't need to attach the memory reservation @@ -353,14 +373,20 @@ impl MultiLevelMergeBuilder { "when only merging memory streams, we should not have any memory reservation and let the merge sort handle the memory" ); - Ok(MergeStep::Stream(merge_sort_stream)) + Ok(MergeStep::Stream { + stream: merge_sort_stream, + batch_size: min_batch_size, + }) } else { // Attach the memory reservation to the stream to make sure we have enough memory // throughout the merge process as we bypassed the memory pool for the merge sort stream - Ok(MergeStep::Stream(Box::pin(StreamAttachedReservation::new( - merge_sort_stream, - memory_reservation, - )))) + Ok(MergeStep::Stream { + stream: Box::pin(StreamAttachedReservation::new( + merge_sort_stream, + memory_reservation, + )), + batch_size: min_batch_size, + }) } } } @@ -371,11 +397,12 @@ impl MultiLevelMergeBuilder { streams: Vec, is_output: bool, all_in_memory: bool, + batch_size: usize, ) -> Result { let mut builder = StreamingMergeBuilder::new() .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr) - .with_batch_size(self.batch_size) + .with_batch_size(batch_size) .with_fetch(self.fetch) .with_metrics(if is_output { // Only add the metrics to the last run @@ -433,9 +460,9 @@ impl MultiLevelMergeBuilder { } let per_spill = get_reserved_bytes_for_record_batch_size( - spill.max_record_batch_memory, + spill.0.max_record_batch_memory, // Size will be the same as the sliced size, bc it is a spilled batch. - spill.max_record_batch_memory, + spill.0.max_record_batch_memory, ) * buffer_len; total_needed += per_spill; @@ -478,8 +505,8 @@ impl MultiLevelMergeBuilder { // of them with a smaller batch size and retry, the smaller max // batch lowers the per-stream reservation enough to seat both. let split_index = usize::from( - self.sorted_spill_files[1].max_record_batch_memory - > self.sorted_spill_files[0].max_record_batch_memory, + self.sorted_spill_files[1].0.max_record_batch_memory + > self.sorted_spill_files[0].0.max_record_batch_memory, ); return Ok(SpillFilesToMerge::SplitThenRetry(split_index)); } @@ -520,7 +547,7 @@ impl MultiLevelMergeBuilder { // exactly as it started, just with the target file shrunk. let last = self.sorted_spill_files.len() - 1; self.sorted_spill_files.swap(index, last); - let target = self + let (target, spill_file_batch_size) = self .sorted_spill_files .pop() .expect("index is in bounds, so the vec is non-empty"); @@ -578,14 +605,14 @@ impl MultiLevelMergeBuilder { // narrower batches. Otherwise the merged stream would rebuild a full-size // (potentially giant) batch and, when spilled back as an intermediate run, // reintroduce the exact skew we just resolved. - self.batch_size = (self.batch_size / 2).max(1); + // self.batch_size = (self.batch_size / 2).max(1); // Push the re-spilled (smaller) file and swap it back into `index`, undoing // the swap-to-back above so the order is preserved. - self.sorted_spill_files.push(SortedSpillFile { + self.sorted_spill_files.push((SortedSpillFile { file, max_record_batch_memory: new_max, - }); + }, (spill_file_batch_size / 2).max(1))); let last = self.sorted_spill_files.len() - 1; self.sorted_spill_files.swap(index, last); @@ -603,7 +630,7 @@ impl MultiLevelMergeBuilder { /// Outcome of trying to reserve memory for one multi-level merge pass. enum SpillFilesToMerge { /// Enough memory: the spill files to read this pass and the read-ahead buffer size. - Ready(Vec, usize), + Ready(Vec<(SortedSpillFile, usize)>, usize), /// Could not seat the minimum of 2 streams. Re-spill the spill file at this index /// with a smaller (halved) batch size, then retry the pass. SplitThenRetry(usize), @@ -612,7 +639,10 @@ enum SpillFilesToMerge { /// What one iteration of the multi-level merge loop should do next. enum MergeStep { /// A merged stream is ready to be consumed (and possibly spilled back). - Stream(SendableRecordBatchStream), + Stream { + stream: SendableRecordBatchStream, + batch_size: usize, + }, /// Re-spill the spill file at this index smaller, then retry the merge step. SplitThenRetry(usize), } From 33c679a59bc3432c6f3661839a6198d17073bf8f Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:16:48 +0300 Subject: [PATCH 25/29] update --- Cargo.lock | 37 ++++++--------------- Cargo.toml | 1 - datafusion/physical-plan/Cargo.toml | 1 - datafusion/physical-plan/src/sorts/merge.rs | 2 +- 4 files changed, 11 insertions(+), 30 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bfef890400595..45d7e5b15e297 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -112,7 +112,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -123,7 +123,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2469,7 +2469,6 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "futures", - "genawaiter", "half", "hashbrown 0.17.1", "indexmap 2.14.0", @@ -2763,7 +2762,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2902,7 +2901,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3174,22 +3173,6 @@ dependencies = [ "prost-build", ] -[[package]] -name = "genawaiter" -version = "0.99.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c86bd0361bcbde39b13475e6e36cb24c329964aa2611be285289d1e4b751c1a0" -dependencies = [ - "futures-core", - "genawaiter-macro", -] - -[[package]] -name = "genawaiter-macro" -version = "0.99.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b32dfe1fdfc0bbde1f22a5da25355514b5e450c33a6af6770884c8750aedfbc" - [[package]] name = "generic-array" version = "0.14.7" @@ -4190,7 +4173,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -5348,7 +5331,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5810,7 +5793,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -5910,7 +5893,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -6079,7 +6062,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -6998,7 +6981,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index afd7a6bc8cef5..a59a9d69c4147 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,7 +159,6 @@ datafusion-session = { path = "datafusion/session", version = "54.0.0" } datafusion-spark = { path = "datafusion/spark", version = "54.0.0" } datafusion-sql = { path = "datafusion/sql", version = "54.0.0" } datafusion-substrait = { path = "datafusion/substrait", version = "54.0.0" } -genawaiter = { version = "0.99.1", default-features = false } doc-comment = "0.3" env_logger = "0.11" diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 36e4b27806b4e..c43ae81003ccc 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -78,7 +78,6 @@ datafusion-physical-expr-common = { workspace = true } datafusion-proto-common = { workspace = true, optional = true } datafusion-proto-models = { workspace = true, optional = true } futures = { workspace = true } -genawaiter = { workspace = true, features = ["futures03"] } half = { workspace = true } hashbrown = { workspace = true } indexmap = { workspace = true } diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 030cfafa42290..8ed8639aff821 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -175,7 +175,7 @@ impl SortPreservingMergeStream { let cloned_metrics = self.metrics.clone(); let mut this = self; - let stream = Gen::new(|co| async move { + let stream = async_try_stream(|co| async move { if let Err(e) = this.run(&co).await { co.yield_(Err(e)).await; } From eb70fbafe35ed499d0ff1c84cd08b6f86d3c6841 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:25:20 +0300 Subject: [PATCH 26/29] update to use the new `async_stream` --- datafusion/physical-plan/src/sorts/merge.rs | 23 +++++++++------------ 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 8ed8639aff821..ef4267449a4cb 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -32,11 +32,11 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use futures::Stream; use datafusion_common::Result; +use datafusion_execution::async_try_stream; use datafusion_execution::memory_pool::MemoryReservation; -use genawaiter::sync::{Co, Gen}; - /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -174,14 +174,7 @@ impl SortPreservingMergeStream { let cloned_metrics = self.metrics.clone(); - let mut this = self; - let stream = async_try_stream(|co| async move { - if let Err(e) = this.run(&co).await { - co.yield_(Err(e)).await; - } - }); - - let stream = Box::pin(RecordBatchStreamAdapter::new(schema_clone, stream)); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema_clone, self.create_stream())); Box::pin(ObservedStream::new(stream, cloned_metrics, None)) } @@ -216,7 +209,9 @@ impl SortPreservingMergeStream { result } - async fn run(&mut self, co: &Co>) -> Result<()> { + 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::>(); @@ -254,7 +249,7 @@ impl SortPreservingMergeStream { && let Some(batch) = self.emit_in_progress_batch()? { drop(timer); - co.yield_(Ok(batch)).await; + emitter.emit(batch).await; timer = elapsed_compute.timer(); } @@ -282,9 +277,11 @@ impl SortPreservingMergeStream { // 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()? { - co.yield_(Ok(batch)).await; + emitter.emit(batch).await; } Ok(()) + + }) } /// Initialize all partitions, return `Poll::Pending` if any partition returns `Poll::Pending` From 18215bd08c42cc7adf72b1bb2cf3467f8c6c35fe Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:26:23 +0300 Subject: [PATCH 27/29] format and lint --- datafusion/physical-plan/src/sorts/merge.rs | 121 ++++++++++---------- 1 file changed, 61 insertions(+), 60 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index ef4267449a4cb..986da549f75c8 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -32,10 +32,10 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use futures::Stream; use datafusion_common::Result; use datafusion_execution::async_try_stream; use datafusion_execution::memory_pool::MemoryReservation; +use futures::Stream; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -174,7 +174,10 @@ impl SortPreservingMergeStream { let cloned_metrics = self.metrics.clone(); - let stream = Box::pin(RecordBatchStreamAdapter::new(schema_clone, self.create_stream())); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); Box::pin(ObservedStream::new(stream, cloned_metrics, None)) } @@ -209,78 +212,76 @@ impl SortPreservingMergeStream { result } - fn create_stream(mut self) -> impl Stream> { + 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::>(); - // 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?; - poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) - .await?; + assert_eq!(uninitiated_partitions.len(), 0); - assert_eq!(uninitiated_partitions.len(), 0); + // If there are no more uninitiated partitions, set up the loser tree and continue + // to the next phase. - // 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 + drop(uninitiated_partitions); + self.init_loser_tree(); - // Claim the memory for the uninitiated partitions - 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(); - // 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(); + loop { + let stream_idx = self.loser_tree[0]; + if !self.advance_cursors(stream_idx) { + break; + } + self.in_progress.push_row(stream_idx); - loop { - 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; + } - // 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(); + } - 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(); - } + 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() { + drop(timer); + poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; + timer = elapsed_compute.timer(); + } - 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() { - 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(); } - // Adjusting the loser tree if necessary - self.update_loser_tree(); - } - - drop(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. - while let Some(batch) = self.emit_in_progress_batch()? { - emitter.emit(batch).await; - } - Ok(()) + drop(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. + while let Some(batch) = self.emit_in_progress_batch()? { + emitter.emit(batch).await; + } + Ok(()) }) } From 2f5812e22ff3ee54648ec772023702e8eac3f34c Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 20:50:57 +0300 Subject: [PATCH 28/29] align with use-yield --- datafusion/physical-plan/src/sorts/merge.rs | 218 ++++++++++---------- 1 file changed, 109 insertions(+), 109 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 1e7af353a58b7..657607def2dc0 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -32,8 +32,8 @@ 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::{async_try_stream, Emitter, TryEmitter}; use datafusion_execution::memory_pool::MemoryReservation; use futures::StreamExt; @@ -234,132 +234,132 @@ impl SortPreservingMergeStream { 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::>(); + // 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?; + poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) + .await?; - assert_eq!(uninitiated_partitions.len(), 0); + assert_eq!(uninitiated_partitions.len(), 0); - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. + // 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 - drop(uninitiated_partitions); - self.init_loser_tree(); + // Claim the memory for the uninitiated partitions + 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(); + // 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(); - } + // 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); - 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() { - drop(timer); - poll_fn(|cx| self.maybe_poll_stream(cx, winner)).await?; + // stop sorting if fetch has been reached + if self.fetch_reached() { + break; + } - timer = elapsed_compute.timer(); + 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(); + } - // 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(); - } + 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() { + 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(); } + } - let last_stream_idx = self.loser_tree[0]; + 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); - } + // 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); + drop(timer); - if self.fetch.is_none_or(|fetch| fetch > self.produced) { - self.passthrough_last_stream(co, last_stream_idx).await?; - } + if self.fetch.is_none_or(|fetch| fetch > self.produced) { + self.passthrough_last_stream(emitter, last_stream_idx).await?; + } - Ok(()) - }) - } + Ok(()) + }) + } - async fn passthrough_last_stream( - &mut self, - emitter: &Emitter>, - last_stream_index: usize, - ) -> Result<()> { - let elapsed_compute = self.metrics.elapsed_compute().clone(); - let mut timer = elapsed_compute.timer(); + 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 last_stream = self.streams.take_partition(last_stream_index); - let mut last_batch = None; + let mut last_batch = None; - // Continue while we still have rows in the in progress builder and not reached fetch limit + // 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) { @@ -368,7 +368,7 @@ impl SortPreservingMergeStream { // 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); - co.yield_(Ok(last_batch)).await; + emitter.emit(last_batch).await; timer = elapsed_compute.timer(); } From 22cb672979a4c74c744c1c97df6b94083ec06a89 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Mon, 13 Jul 2026 20:53:52 +0300 Subject: [PATCH 29/29] lint and format --- datafusion/physical-plan/src/sorts/builder.rs | 18 +- datafusion/physical-plan/src/sorts/merge.rs | 228 +++++++++--------- 2 files changed, 116 insertions(+), 130 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 744cf032d2ff7..1f0f40e3d582b 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -207,28 +207,14 @@ 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> { - self.build_record_batch_with_up_to_n_rows(self.indices.len()) - } - - /// Drains the in_progress row indexes up to n rows, 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(crate) fn build_record_batch_with_up_to_n_rows( - &mut self, - n: usize, - ) -> Result> { + pub fn build_record_batch(&mut self, n: usize) -> Result> { if self.is_empty() || n == 0 { return Ok(None); } diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 657607def2dc0..6bbb538d0c71c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -33,8 +33,8 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result}; -use datafusion_execution::{async_try_stream, Emitter, TryEmitter}; use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{TryEmitter, async_try_stream}; use futures::StreamExt; use crate::coalesce_batches::CoalesceBatchesStream; @@ -225,141 +225,141 @@ impl SortPreservingMergeStream { // 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_with_up_to_n_rows(rows_to_emit); + let result = self.in_progress.build_record_batch(rows_to_emit); self.produced += rows_before - self.in_progress.len(); result } - 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::>(); + 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?; + poll_fn(|cx| self.initialize_all_partitions(&mut uninitiated_partitions, cx)) + .await?; - assert_eq!(uninitiated_partitions.len(), 0); + assert_eq!(uninitiated_partitions.len(), 0); - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. + // 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 - drop(uninitiated_partitions); - self.init_loser_tree(); + // Claim the memory for the uninitiated partitions + 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(); + // 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); + // 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; - } + // 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(); - } + 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(); + } - 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() { - 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(); - } - } + 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() { + 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(); + } + } - let last_stream_idx = self.loser_tree[0]; + 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); - } - } + // 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); + drop(timer); - if self.fetch.is_none_or(|fetch| fetch > self.produced) { - self.passthrough_last_stream(emitter, last_stream_idx).await?; - } + if self.fetch.is_none_or(|fetch| fetch > self.produced) { + self.passthrough_last_stream(emitter, last_stream_idx) + .await?; + } - Ok(()) - }) - } + 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(); + 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 last_stream = self.streams.take_partition(last_stream_index); - let mut last_batch = None; + let mut last_batch = None; - // Continue while we still have rows in the in progress builder and not reached fetch limit + // 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) {