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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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/18] 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 }