diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 95cd75f50a004..4eb1070e6c857 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -16,6 +16,7 @@ // under the License. //! Fuzz Test for various corner cases merging streams of RecordBatches + use std::sync::Arc; use arrow::{ diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 85418ff36119d..875922ac34b54 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,19 +18,23 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{ready, Context, Poll}; + use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; use crate::RecordBatchStream; + use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_execution::memory_pool::MemoryReservation; + use futures::Stream; -use std::pin::Pin; -use std::sync::Arc; -use std::task::{ready, Context, Poll}; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -86,7 +90,7 @@ pub(crate) struct SortPreservingMergeStream { /// been updated loser_tree_adjusted: bool, - /// target batch size + /// Target batch size batch_size: usize, /// Cursors for each input partition. `None` means the input is exhausted @@ -97,6 +101,12 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, + + /// This queue contains partition indices in order. When a partition is polled and returns `Poll::Ready`, + /// it is removed from the vector. If a partition returns `Poll::Pending`, it is moved to the end of the + /// vector to ensure the next iteration starts with a different partition, preventing the same partition + /// from being continuously polled. + uninitiated_partitions: VecDeque, } impl SortPreservingMergeStream { @@ -121,6 +131,7 @@ impl SortPreservingMergeStream { batch_size, fetch, produced: 0, + uninitiated_partitions: (0..stream_count).collect(), } } @@ -154,14 +165,36 @@ impl SortPreservingMergeStream { if self.aborted { return Poll::Ready(None); } - // try to initialize the loser tree + // 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() { - // Ensure all non-exhausted streams have a cursor from which - // rows can be pulled - for i in 0..self.streams.partitions() { - if let Err(e) = ready!(self.maybe_poll_stream(cx, i)) { - self.aborted = true; - return Poll::Ready(Some(Err(e))); + let remaining_partitions = self.uninitiated_partitions.clone(); + for i in remaining_partitions { + match self.maybe_poll_stream(cx, i) { + Poll::Ready(Err(e)) => { + self.aborted = true; + return Poll::Ready(Some(Err(e))); + } + Poll::Pending => { + // If a partition returns Poll::Pending, to avoid continuously polling it + // and potentially increasing upstream buffer sizes, we move it to the + // back of the polling queue. + if let Some(front) = self.uninitiated_partitions.pop_front() { + // This pop_front can never return `None`. + self.uninitiated_partitions.push_back(front); + } + // This function could remain in a pending state, so we manually wake it here. + // However, this approach can be investigated further to find a more natural way + // to avoid disrupting the runtime scheduler. + cx.waker().wake_by_ref(); + return Poll::Pending; + } + _ => { + // If the polling result is Poll::Ready(Some(batch)) or Poll::Ready(None), + // we remove this partition from the queue so it is not polled again. + self.uninitiated_partitions.retain(|idx| *idx != i); + } } } self.init_loser_tree(); diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 131fa71217cce..4d333175bf758 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -300,6 +300,11 @@ impl ExecutionPlan for SortPreservingMergeExec { #[cfg(test)] mod tests { + use std::fmt::Formatter; + use std::pin::Pin; + use std::sync::Mutex; + use std::task::{Context, Poll}; + use std::time::Duration; use super::*; use crate::coalesce_partitions::CoalescePartitionsExec; @@ -310,16 +315,23 @@ mod tests { use crate::stream::RecordBatchReceiverStream; use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec}; use crate::test::{self, assert_is_pending, make_partition}; - use crate::{collect, common}; + use crate::{collect, common, ExecutionMode}; use arrow::array::{ArrayRef, Int32Array, StringArray, TimestampNanosecondArray}; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; - use datafusion_common::{assert_batches_eq, assert_contains}; + use arrow_schema::SchemaRef; + use datafusion_common::{assert_batches_eq, assert_contains, DataFusionError}; + use datafusion_common_runtime::SpawnedTask; use datafusion_execution::config::SessionConfig; + use datafusion_execution::RecordBatchStream; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::EquivalenceProperties; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; - use futures::{FutureExt, StreamExt}; + use futures::{FutureExt, Stream, StreamExt}; + use tokio::time::timeout; #[tokio::test] async fn test_merge_interleave() { @@ -1141,4 +1153,157 @@ mod tests { collected.as_slice() ); } + + /// It returns pending for the 2nd partition until the 3rd partition is polled. The 1st + /// partition is exhausted from the start, and if it is polled more than one, it panics. + #[derive(Debug, Clone)] + struct CongestedExec { + schema: Schema, + cache: PlanProperties, + congestion_cleared: Arc>, + } + + impl CongestedExec { + fn compute_properties(schema: SchemaRef) -> PlanProperties { + let columns = schema + .fields + .iter() + .enumerate() + .map(|(i, f)| Arc::new(Column::new(f.name(), i)) as Arc) + .collect::>(); + let mut eq_properties = EquivalenceProperties::new(schema); + eq_properties.add_new_orderings(vec![columns + .iter() + .map(|expr| { + PhysicalSortExpr::new(Arc::clone(expr), SortOptions::default()) + }) + .collect::>()]); + let mode = ExecutionMode::Unbounded; + PlanProperties::new(eq_properties, Partitioning::Hash(columns, 3), mode) + } + } + + impl ExecutionPlan for CongestedExec { + fn name(&self) -> &'static str { + Self::static_name() + } + fn as_any(&self) -> &dyn Any { + self + } + fn properties(&self) -> &PlanProperties { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + Ok(self) + } + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(CongestedStream { + schema: Arc::new(self.schema.clone()), + none_polled_once: false, + congestion_cleared: Arc::clone(&self.congestion_cleared), + partition, + })) + } + } + + impl DisplayAs for CongestedExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CongestedExec",).unwrap() + } + } + Ok(()) + } + } + + /// It returns pending for the 2nd partition until the 3rd partition is polled. The 1st + /// partition is exhausted from the start, and if it is polled more than once, it panics. + #[derive(Debug)] + pub struct CongestedStream { + schema: SchemaRef, + none_polled_once: bool, + congestion_cleared: Arc>, + partition: usize, + } + + impl Stream for CongestedStream { + type Item = Result; + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + match self.partition { + 0 => { + if self.none_polled_once { + panic!("Exhausted stream is polled more than one") + } else { + self.none_polled_once = true; + Poll::Ready(None) + } + } + 1 => { + let cleared = self.congestion_cleared.lock().unwrap(); + if *cleared { + Poll::Ready(None) + } else { + Poll::Pending + } + } + 2 => { + let mut cleared = self.congestion_cleared.lock().unwrap(); + *cleared = true; + Poll::Ready(None) + } + _ => unreachable!(), + } + } + } + + impl RecordBatchStream for CongestedStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + } + + #[tokio::test] + async fn test_spm_congestion() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); + let source = CongestedExec { + schema: schema.clone(), + cache: CongestedExec::compute_properties(Arc::new(schema.clone())), + congestion_cleared: Arc::new(Mutex::new(false)), + }; + let spm = SortPreservingMergeExec::new( + vec![PhysicalSortExpr::new( + Arc::new(Column::new("c1", 0)), + SortOptions::default(), + )], + Arc::new(source), + ); + let spm_task = SpawnedTask::spawn(collect(Arc::new(spm), task_ctx)); + + let result = timeout(Duration::from_secs(3), spm_task.join()).await; + match result { + Ok(Ok(Ok(_batches))) => Ok(()), + Ok(Ok(Err(e))) => Err(e), + Ok(Err(_)) => Err(DataFusionError::Execution( + "SortPreservingMerge task panicked or was cancelled".to_string(), + )), + Err(_) => Err(DataFusionError::Execution( + "SortPreservingMerge caused a deadlock".to_string(), + )), + } + } }