diff --git a/Cargo.lock b/Cargo.lock index 5b43435ec0a7b..6179841f3c9e9 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]] @@ -2468,6 +2468,7 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "futures", + "genawaiter", "half", "hashbrown 0.17.1", "indexmap 2.14.0", @@ -2761,7 +2762,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -2900,7 +2901,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]] @@ -3172,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" @@ -4172,7 +4189,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 +5347,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5792,7 +5809,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 +5909,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -6061,7 +6078,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6980,7 +6997,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..2d93292c9de89 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -159,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 c43ae81003ccc..36e4b27806b4e 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -78,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 } 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 4117789777fe8..030cfafa42290 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,22 +18,24 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::pin::Pin; +use std::fmt::Debug; +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::SendableRecordBatchStream; use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_execution::memory_pool::MemoryReservation; -use futures::Stream; +use genawaiter::sync::{Co, Gen}; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -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 @@ -93,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, @@ -150,9 +134,6 @@ 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, } impl SortPreservingMergeStream { @@ -171,8 +152,6 @@ impl SortPreservingMergeStream { in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation), streams, metrics, - done: false, - drain_in_progress_on_done: false, cursors: (0..stream_count).map(|_| None).collect(), prev_cursors: (0..stream_count).map(|_| None).collect(), round_robin_tie_breaker_mode: false, @@ -180,15 +159,33 @@ impl SortPreservingMergeStream { current_reset_epoch: 0, poll_reset_epochs: vec![0; stream_count], loser_tree: vec![], - loser_tree_adjusted: false, batch_size, fetch, produced: 0, - uninitiated_partitions: (0..stream_count).collect(), enable_round_robin_tie_breaker, } } + pub(crate) fn into_stream(self) -> SendableRecordBatchStream + where + C: 'static, + { + let schema_clone = Arc::clone(self.in_progress.schema()); + + let cloned_metrics = self.metrics.clone(); + + let mut this = self; + let stream = Gen::new(|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)); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + /// If the stream at the given index is not exhausted, and the last cursor for the /// stream is finished, poll the stream for the next RecordBatch and create a new /// cursor for the stream from the returned result @@ -219,90 +216,86 @@ 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, 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::>(); - // Once all partitions have set their corresponding cursors for the loser tree, - // we skip the following block. Until then, this function may be called multiple - // times and can return Poll::Pending if any partition returns Poll::Pending. - if self.loser_tree.is_empty() { - ready!(self.initialize_all_partitions(cx))?; - assert_eq!( - self.uninitiated_partitions.len(), - 0, - "all partitions should be initialized" - ); - - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. - - // Claim the memory for the uninitiated partitions - self.uninitiated_partitions.shrink_to_fit(); - self.init_loser_tree(); - } + 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, so there are no // calls to `timer.done()` below. let elapsed_compute = self.metrics.elapsed_compute().clone(); - let _timer = elapsed_compute.timer(); + let mut timer = elapsed_compute.timer(); loop { - // Adjust the loser tree if necessary, returning control if needed - if !self.loser_tree_adjusted { - let winner = self.loser_tree[0]; - // Fast path: skip the `maybe_poll_stream` call (and its `Poll` - // plumbing) unless the winner's cursor is exhausted and needs a - // fresh batch — it is live for almost every row. - if self.cursors[winner].is_none() { - match ready!(self.maybe_poll_stream(cx, winner)) { - Ok(()) => {} - Err(e) => { - self.done = true; - return Poll::Ready(Some(Err(e))); - } - } - } - self.update_loser_tree(); + let stream_idx = self.loser_tree[0]; + if !self.advance_cursors(stream_idx) { + break; } + self.in_progress.push_row(stream_idx); - 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; - } + // 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); + co.yield_(Ok(batch)).await; + timer = elapsed_compute.timer(); } - return Poll::Ready(self.emit_in_progress_batch().transpose()); + 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(); } + + 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()? { + co.yield_(Ok(batch)).await; + } + Ok(()) } /// Initialize all partitions, return `Poll::Pending` if any partition returns `Poll::Pending` /// /// This DOES NOT return `Poll::Pending` as soon as the first uninitiated partition returns `Poll::Pending` /// so we can continue to initialize the remaining partitions - fn initialize_all_partitions(&mut self, cx: &mut Context) -> Poll> { + fn initialize_all_partitions( + &mut self, + uninitiated_partitions: &mut Vec, + cx: &mut Context, + ) -> Poll> { assert_eq!( self.loser_tree.len(), 0, @@ -311,11 +304,10 @@ impl SortPreservingMergeStream { // Manual indexing since we're iterating over the vector and shrinking it in the loop let mut idx = 0; - while idx < self.uninitiated_partitions.len() { - let partition_idx = self.uninitiated_partitions[idx]; + while idx < uninitiated_partitions.len() { + let partition_idx = uninitiated_partitions[idx]; match self.maybe_poll_stream(cx, partition_idx) { Poll::Ready(Err(e)) => { - self.done = true; return Poll::Ready(Err(e)); } Poll::Pending => { @@ -331,12 +323,12 @@ impl SortPreservingMergeStream { // place which we'll try in the next loop iteration // swap_remove will change the partition poll order, but that shouldn't // make a difference since we're waiting for all streams to be ready. - self.uninitiated_partitions.swap_remove(idx); + uninitiated_partitions.swap_remove(idx); } } } - if self.uninitiated_partitions.is_empty() { + if uninitiated_partitions.is_empty() { Poll::Ready(Ok(())) } else { // There are still uninitiated partitions so return pending. @@ -474,7 +466,6 @@ impl SortPreservingMergeStream { } self.loser_tree[cmp_node] = winner; } - self.loser_tree_adjusted = true; } /// Resets the poll count by incrementing the reset epoch. @@ -586,25 +577,6 @@ impl SortPreservingMergeStream { } self.loser_tree[0] = winner; - self.loser_tree_adjusted = true; - } -} - -impl Stream for SortPreservingMergeStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - let poll = self.poll_next_inner(cx); - self.metrics.record_poll(poll) - } -} - -impl RecordBatchStream for SortPreservingMergeStream { - fn schema(&self) -> SchemaRef { - Arc::clone(self.in_progress.schema()) } } @@ -618,7 +590,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)] @@ -661,8 +633,8 @@ mod tests { } } - #[test] - fn test_done_drains_buffered_rows() { + #[tokio::test] + async fn test_done_drains_buffered_rows() { let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); let pool: Arc = Arc::new(UnboundedMemoryPool::default()); let reservation = MemoryConsumer::new("test").register(&pool); @@ -678,24 +650,20 @@ mod tests { true, ); + // Simulate rows left buffered in `in_progress` (as happens when + // `build_record_batch` emits a partial batch on offset overflow). With + // an empty input stream the merge loop breaks immediately, so the only + // way these rows reach the consumer is the generator's final drain loop. let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]) .unwrap(); stream.in_progress.push_batch(0, batch).unwrap(); stream.in_progress.push_row(0); - stream.done = true; - stream.drain_in_progress_on_done = true; - let waker = noop_waker_ref(); - let mut cx = Context::from_waker(waker); + // Drive the actual stream and confirm the buffered row is drained. + let batches: Vec = stream.into_stream().try_collect().await.unwrap(); - match stream.poll_next_inner(&mut cx) { - Poll::Ready(Some(Ok(batch))) => assert_eq!(batch.num_rows(), 1), - other => { - panic!("expected buffered rows to be drained after done, got {other:?}") - } - } - assert!(stream.in_progress.is_empty()); - assert!(matches!(stream.poll_next_inner(&mut cx), Poll::Ready(None))); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); } } diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index ade24ff0534ff..e96138ef1306c 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -46,7 +46,7 @@ macro_rules! merge_helper { ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ let streams = FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); - return Ok(Box::pin(SortPreservingMergeStream::new( + return Ok(SortPreservingMergeStream::new( Box::new(streams), $schema, $tracking_metrics, @@ -54,7 +54,8 @@ macro_rules! merge_helper { $fetch, $reservation, $enable_round_robin_tie_breaker, - ))); + ) + .into_stream()); }}; } @@ -254,7 +255,7 @@ impl<'a> StreamingMergeBuilder<'a> { streams, reservation.new_empty(), )?; - Ok(Box::pin(SortPreservingMergeStream::new( + Ok(SortPreservingMergeStream::new( Box::new(streams), schema, metrics, @@ -262,6 +263,7 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, reservation, enable_round_robin_tie_breaker, - ))) + ) + .into_stream()) } }