From 4382945fdffcf805aca0388ad6fc2b207817aa0d Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Thu, 29 Aug 2024 11:34:25 +0300 Subject: [PATCH 01/13] Ready Update merge_fuzz.rs Update merge.rs Update merge.rs Update merge.rs Update merge_fuzz.rs Update merge.rs Update merge_fuzz.rs Update merge.rs Counter is not enough Termination logic with counter --- .../core/tests/fuzz_cases/merge_fuzz.rs | 202 +++++++++++++++++- datafusion/physical-plan/src/sorts/merge.rs | 44 +++- 2 files changed, 237 insertions(+), 9 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 95cd75f50a004..8c0e634e750e3 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -16,13 +16,21 @@ // under the License. //! Fuzz Test for various corner cases merging streams of RecordBatches -use std::sync::Arc; + +use std::any::Any; +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Duration; use arrow::{ array::{ArrayRef, Int32Array}, compute::SortOptions, record_batch::RecordBatch, }; +use arrow_array::UInt64Array; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::physical_plan::{ collect, expressions::{col, PhysicalSortExpr}, @@ -30,8 +38,20 @@ use datafusion::physical_plan::{ sorts::sort_preserving_merge::SortPreservingMergeExec, }; use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::{DataFusionError, Result}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties, +}; use test_utils::{batches_to_vec, partitions_to_sorted_vec, stagger_batch_with_seed}; +use futures::Stream; +use tokio::time::timeout; + #[tokio::test] async fn test_merge_2() { run_merge_test(vec![ @@ -160,3 +180,183 @@ fn concat(mut v1: Vec, v2: Vec) -> Vec { v1.extend(v2); v1 } + +#[derive(Debug, Clone)] +struct CongestedExec { + schema: Schema, + cache: PlanProperties, + batch_size: u64, + limit: u64, + 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(expr.clone(), SortOptions::default())) + .collect::>()]); + let mode = ExecutionMode::Unbounded; + PlanProperties::new(eq_properties, Partitioning::Hash(columns, 2), 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()), + batch_size: self.batch_size, + offset: 0, + congestion_cleared: self.congestion_cleared.clone(), + limit: self.limit, + 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(()) + } +} + +#[derive(Debug)] +pub struct CongestedStream { + schema: SchemaRef, + batch_size: u64, + offset: u64, + limit: u64, + congestion_cleared: Arc>, + partition: usize, +} + +impl CongestedStream { + fn create_record_batch( + schema: SchemaRef, + value: u64, + batch_size: u64, + ) -> RecordBatch { + let values = (0..batch_size).map(|_| value).collect::>(); + let array = UInt64Array::from(values); + let array_ref: ArrayRef = Arc::new(array); + RecordBatch::try_new(schema, vec![array_ref]).unwrap() + } +} + +impl Stream for CongestedStream { + type Item = Result; + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + match self.partition { + 0 => { + let cleared = self.congestion_cleared.lock().unwrap(); + if *cleared { + std::mem::drop(cleared); + if self.offset == self.limit { + return Poll::Ready(None); + } + let batch = CongestedStream::create_record_batch( + Arc::clone(&self.schema), + 0, + self.batch_size, + ); + self.offset += self.batch_size; + Poll::Ready(Some(Ok(batch))) + } else { + Poll::Pending + } + } + 1 => { + if self.offset == self.limit { + return Poll::Ready(None); + } + let batch = CongestedStream::create_record_batch( + Arc::clone(&self.schema), + 0, + self.batch_size, + ); + self.offset += self.batch_size; + let mut cleared = self.congestion_cleared.lock().unwrap(); + *cleared = true; + Poll::Ready(Some(Ok(batch))) + } + _ => 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(), + batch_size: 1, + cache: CongestedExec::compute_properties(Arc::new(schema.clone())), + limit: 1_000, + 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(async move { collect(Arc::new(spm), task_ctx).await }); + + let result = timeout(Duration::from_secs(1), 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(), + )), + } +} diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 85418ff36119d..235526ad70301 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::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}; +use hashbrown::HashSet; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -97,6 +101,9 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, + + /// Cursors which are not initialized yet + uninitiated_cursors: HashSet, } impl SortPreservingMergeStream { @@ -121,6 +128,7 @@ impl SortPreservingMergeStream { batch_size, fetch, produced: 0, + uninitiated_cursors: (0..stream_count).collect(), } } @@ -138,9 +146,13 @@ impl SortPreservingMergeStream { } match futures::ready!(self.streams.poll_next(cx, idx)) { - None => Poll::Ready(Ok(())), + None => { + self.uninitiated_cursors.retain(|&x| x != idx); + Poll::Ready(Ok(())) + } Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok((cursor, batch))) => { + self.uninitiated_cursors.retain(|&x| x != idx); self.cursors[idx] = Some(Cursor::new(cursor)); Poll::Ready(self.in_progress.push_batch(idx, batch)) } @@ -158,11 +170,27 @@ impl SortPreservingMergeStream { 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))); + loop { + let mut return_pending = false; + let uninitiated_cursors = self.uninitiated_cursors.clone(); + for i in uninitiated_cursors { + if self.uninitiated_cursors.contains(&i) { + match self.maybe_poll_stream(cx, i) { + Poll::Ready(Err(e)) => { + self.aborted = true; + return Poll::Ready(Some(Err(e))); + } + Poll::Pending => { + return_pending = true; + } + _ => {} + } + } + } + if self.uninitiated_cursors.is_empty() { + break; } + return Poll::Pending; } self.init_loser_tree(); } From d14dde3a534cf5826f4d8314396f62c4c13dfd93 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Mon, 2 Sep 2024 18:32:23 +0300 Subject: [PATCH 02/13] Add comments --- datafusion/core/tests/fuzz_cases/merge_fuzz.rs | 14 ++++++++------ datafusion/physical-plan/src/sorts/merge.rs | 4 +++- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 8c0e634e750e3..81a378c1fb3a7 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -181,6 +181,7 @@ fn concat(mut v1: Vec, v2: Vec) -> Vec { v1 } +/// It returns pending for the 1st partition until the 2nd partition is polled. #[derive(Debug, Clone)] struct CongestedExec { schema: Schema, @@ -235,7 +236,7 @@ impl ExecutionPlan for CongestedExec { Ok(Box::pin(CongestedStream { schema: Arc::new(self.schema.clone()), batch_size: self.batch_size, - offset: 0, + offset: (0, 0), congestion_cleared: self.congestion_cleared.clone(), limit: self.limit, partition, @@ -254,11 +255,12 @@ impl DisplayAs for CongestedExec { } } +/// It returns pending for the 1st partition until the 2nd partition is polled. #[derive(Debug)] pub struct CongestedStream { schema: SchemaRef, batch_size: u64, - offset: u64, + offset: (u64, u64), limit: u64, congestion_cleared: Arc>, partition: usize, @@ -288,7 +290,7 @@ impl Stream for CongestedStream { let cleared = self.congestion_cleared.lock().unwrap(); if *cleared { std::mem::drop(cleared); - if self.offset == self.limit { + if self.offset.0 == self.limit { return Poll::Ready(None); } let batch = CongestedStream::create_record_batch( @@ -296,14 +298,14 @@ impl Stream for CongestedStream { 0, self.batch_size, ); - self.offset += self.batch_size; + self.offset.0 += self.batch_size; Poll::Ready(Some(Ok(batch))) } else { Poll::Pending } } 1 => { - if self.offset == self.limit { + if self.offset.1 == self.limit { return Poll::Ready(None); } let batch = CongestedStream::create_record_batch( @@ -311,7 +313,7 @@ impl Stream for CongestedStream { 0, self.batch_size, ); - self.offset += self.batch_size; + self.offset.1 += self.batch_size; let mut cleared = self.congestion_cleared.lock().unwrap(); *cleared = true; Poll::Ready(Some(Ok(batch))) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 235526ad70301..9f98abb5cddc5 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -190,7 +190,9 @@ impl SortPreservingMergeStream { if self.uninitiated_cursors.is_empty() { break; } - return Poll::Pending; + if return_pending { + return Poll::Pending; + } } self.init_loser_tree(); } From 1567c0c3490b1b19d201a06aa5cccd8791c3ef41 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Mon, 2 Sep 2024 18:40:40 +0300 Subject: [PATCH 03/13] Increase threshold --- datafusion/core/tests/fuzz_cases/merge_fuzz.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 81a378c1fb3a7..241b32752f47b 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -350,7 +350,7 @@ async fn test_spm_congestion() -> Result<()> { let spm_task = SpawnedTask::spawn(async move { collect(Arc::new(spm), task_ctx).await }); - let result = timeout(Duration::from_secs(1), spm_task.join()).await; + let result = timeout(Duration::from_secs(10), spm_task.join()).await; match result { Ok(Ok(Ok(_batches))) => Ok(()), Ok(Ok(Err(e))) => Err(e), From 02cdbfec594b8eaf88da1dab3e59957aea910d56 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Tue, 3 Sep 2024 11:13:44 +0300 Subject: [PATCH 04/13] Debug --- .../core/tests/fuzz_cases/merge_fuzz.rs | 6 +-- datafusion/physical-plan/src/sorts/merge.rs | 50 +++++-------------- 2 files changed, 15 insertions(+), 41 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 241b32752f47b..4fb05f92d8833 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -285,6 +285,7 @@ impl Stream for CongestedStream { mut self: Pin<&mut Self>, _cx: &mut Context<'_>, ) -> Poll> { + println!("asd"); match self.partition { 0 => { let cleared = self.congestion_cleared.lock().unwrap(); @@ -347,10 +348,9 @@ async fn test_spm_congestion() -> Result<()> { )], Arc::new(source), ); - let spm_task = - SpawnedTask::spawn(async move { collect(Arc::new(spm), task_ctx).await }); + let spm_task = SpawnedTask::spawn(collect(Arc::new(spm), task_ctx)); - let result = timeout(Duration::from_secs(10), spm_task.join()).await; + let result = timeout(Duration::from_secs(5), spm_task.join()).await; match result { Ok(Ok(Ok(_batches))) => Ok(()), Ok(Ok(Err(e))) => Err(e), diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 9f98abb5cddc5..52c28522989de 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,23 +18,19 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -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 hashbrown::HashSet; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{ready, Context, Poll}; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] type CursorStream = Box>>; @@ -101,9 +97,6 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, - - /// Cursors which are not initialized yet - uninitiated_cursors: HashSet, } impl SortPreservingMergeStream { @@ -128,7 +121,6 @@ impl SortPreservingMergeStream { batch_size, fetch, produced: 0, - uninitiated_cursors: (0..stream_count).collect(), } } @@ -146,13 +138,9 @@ impl SortPreservingMergeStream { } match futures::ready!(self.streams.poll_next(cx, idx)) { - None => { - self.uninitiated_cursors.retain(|&x| x != idx); - Poll::Ready(Ok(())) - } + None => Poll::Ready(Ok(())), Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok((cursor, batch))) => { - self.uninitiated_cursors.retain(|&x| x != idx); self.cursors[idx] = Some(Cursor::new(cursor)); Poll::Ready(self.in_progress.push_batch(idx, batch)) } @@ -163,37 +151,23 @@ impl SortPreservingMergeStream { &mut self, cx: &mut Context<'_>, ) -> Poll>> { + println!("111"); if self.aborted { return Poll::Ready(None); } // try to initialize the loser tree if self.loser_tree.is_empty() { + println!("222"); // Ensure all non-exhausted streams have a cursor from which // rows can be pulled - loop { - let mut return_pending = false; - let uninitiated_cursors = self.uninitiated_cursors.clone(); - for i in uninitiated_cursors { - if self.uninitiated_cursors.contains(&i) { - match self.maybe_poll_stream(cx, i) { - Poll::Ready(Err(e)) => { - self.aborted = true; - return Poll::Ready(Some(Err(e))); - } - Poll::Pending => { - return_pending = true; - } - _ => {} - } - } - } - if self.uninitiated_cursors.is_empty() { - break; - } - if return_pending { - return Poll::Pending; + for i in 0..self.streams.partitions() { + println!("i: {}", i); + if let Err(e) = ready!(self.maybe_poll_stream(cx, i)) { + self.aborted = true; + return Poll::Ready(Some(Err(e))); } } + println!("333"); self.init_loser_tree(); } From eb068a07020c44f5212e087269390652bd47e397 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Tue, 3 Sep 2024 11:24:50 +0300 Subject: [PATCH 05/13] Remove merge.rs changes --- datafusion/core/tests/fuzz_cases/merge_fuzz.rs | 1 - datafusion/physical-plan/src/sorts/merge.rs | 4 ---- 2 files changed, 5 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 4fb05f92d8833..6b5afd865bf1a 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -285,7 +285,6 @@ impl Stream for CongestedStream { mut self: Pin<&mut Self>, _cx: &mut Context<'_>, ) -> Poll> { - println!("asd"); match self.partition { 0 => { let cleared = self.congestion_cleared.lock().unwrap(); diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 52c28522989de..85418ff36119d 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -151,23 +151,19 @@ impl SortPreservingMergeStream { &mut self, cx: &mut Context<'_>, ) -> Poll>> { - println!("111"); if self.aborted { return Poll::Ready(None); } // try to initialize the loser tree if self.loser_tree.is_empty() { - println!("222"); // Ensure all non-exhausted streams have a cursor from which // rows can be pulled for i in 0..self.streams.partitions() { - println!("i: {}", i); if let Err(e) = ready!(self.maybe_poll_stream(cx, i)) { self.aborted = true; return Poll::Ready(Some(Err(e))); } } - println!("333"); self.init_loser_tree(); } From 93a9c7cf050b313c8be5788c47add957608f91b2 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Tue, 3 Sep 2024 14:03:09 +0300 Subject: [PATCH 06/13] Use waker --- .../core/tests/fuzz_cases/merge_fuzz.rs | 2 +- datafusion/physical-plan/src/sorts/merge.rs | 25 +++++++++++++++---- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 6b5afd865bf1a..53c2acca4dd8c 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -349,7 +349,7 @@ async fn test_spm_congestion() -> Result<()> { ); let spm_task = SpawnedTask::spawn(collect(Arc::new(spm), task_ctx)); - let result = timeout(Duration::from_secs(5), spm_task.join()).await; + let result = timeout(Duration::from_secs(1), spm_task.join()).await; match result { Ok(Ok(Ok(_batches))) => Ok(()), Ok(Ok(Err(e))) => Err(e), diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 85418ff36119d..789b7f839d3a0 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -95,8 +95,12 @@ pub(crate) struct SortPreservingMergeStream { /// Optional number of rows to fetch fetch: Option, - /// number of rows produced + /// Number of rows produced produced: usize, + + /// A vector having partition indices in a priortiy order + /// to visit the partitions in a round-robin fashion + partitions_in_priority_order: Vec, } impl SortPreservingMergeStream { @@ -121,6 +125,7 @@ impl SortPreservingMergeStream { batch_size, fetch, produced: 0, + partitions_in_priority_order: (0..stream_count).collect::>(), } } @@ -137,6 +142,8 @@ impl SortPreservingMergeStream { return Poll::Ready(Ok(())); } + cx.waker().wake_by_ref(); + match futures::ready!(self.streams.poll_next(cx, idx)) { None => Poll::Ready(Ok(())), Some(Err(e)) => Poll::Ready(Err(e)), @@ -158,10 +165,18 @@ impl SortPreservingMergeStream { 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 partitions = self.partitions_in_priority_order.clone(); + for i in partitions { + match self.maybe_poll_stream(cx, i) { + Poll::Ready(Err(e)) => { + self.aborted = true; + return Poll::Ready(Some(Err(e))); + } + Poll::Pending => { + self.partitions_in_priority_order.rotate_left(1); + return Poll::Pending; + } + _ => {} } } self.init_loser_tree(); From 5640da893dc1c36d998255a4f87b4e156a651d61 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Tue, 3 Sep 2024 14:29:57 +0300 Subject: [PATCH 07/13] Update merge.rs --- datafusion/physical-plan/src/sorts/merge.rs | 31 +++++++++++---------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 789b7f839d3a0..3c6b082084166 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,19 +18,22 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. +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>>; @@ -95,12 +98,12 @@ pub(crate) struct SortPreservingMergeStream { /// Optional number of rows to fetch fetch: Option, - /// Number of rows produced + /// number of rows produced produced: usize, - /// A vector having partition indices in a priortiy order - /// to visit the partitions in a round-robin fashion - partitions_in_priority_order: Vec, + /// Unitiated partitions. They are stored in a vector to keep them in + /// a priortiy order to visit the partitions in a round-robin fashion + uninitiated_partitions: Vec, } impl SortPreservingMergeStream { @@ -125,7 +128,7 @@ impl SortPreservingMergeStream { batch_size, fetch, produced: 0, - partitions_in_priority_order: (0..stream_count).collect::>(), + uninitiated_partitions: (0..stream_count).collect(), } } @@ -142,8 +145,6 @@ impl SortPreservingMergeStream { return Poll::Ready(Ok(())); } - cx.waker().wake_by_ref(); - match futures::ready!(self.streams.poll_next(cx, idx)) { None => Poll::Ready(Ok(())), Some(Err(e)) => Poll::Ready(Err(e)), @@ -165,18 +166,20 @@ impl SortPreservingMergeStream { if self.loser_tree.is_empty() { // Ensure all non-exhausted streams have a cursor from which // rows can be pulled - let partitions = self.partitions_in_priority_order.clone(); - for i in partitions { + 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 => { - self.partitions_in_priority_order.rotate_left(1); + self.uninitiated_partitions.rotate_left(1); return Poll::Pending; } - _ => {} + _ => { + self.uninitiated_partitions.retain(|idx| *idx != i); + } } } self.init_loser_tree(); From 07bf172c3d5f25ad09cbcf0756bed4efbf3815d6 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Tue, 3 Sep 2024 14:41:48 +0300 Subject: [PATCH 08/13] Simplify the test --- .../core/tests/fuzz_cases/merge_fuzz.rs | 49 ++----------------- datafusion/physical-plan/src/sorts/merge.rs | 4 +- 2 files changed, 5 insertions(+), 48 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 53c2acca4dd8c..7b87ccde83176 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -29,7 +29,6 @@ use arrow::{ compute::SortOptions, record_batch::RecordBatch, }; -use arrow_array::UInt64Array; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::physical_plan::{ collect, @@ -186,8 +185,6 @@ fn concat(mut v1: Vec, v2: Vec) -> Vec { struct CongestedExec { schema: Schema, cache: PlanProperties, - batch_size: u64, - limit: u64, congestion_cleared: Arc>, } @@ -235,10 +232,7 @@ impl ExecutionPlan for CongestedExec { ) -> Result { Ok(Box::pin(CongestedStream { schema: Arc::new(self.schema.clone()), - batch_size: self.batch_size, - offset: (0, 0), congestion_cleared: self.congestion_cleared.clone(), - limit: self.limit, partition, })) } @@ -259,26 +253,10 @@ impl DisplayAs for CongestedExec { #[derive(Debug)] pub struct CongestedStream { schema: SchemaRef, - batch_size: u64, - offset: (u64, u64), - limit: u64, congestion_cleared: Arc>, partition: usize, } -impl CongestedStream { - fn create_record_batch( - schema: SchemaRef, - value: u64, - batch_size: u64, - ) -> RecordBatch { - let values = (0..batch_size).map(|_| value).collect::>(); - let array = UInt64Array::from(values); - let array_ref: ArrayRef = Arc::new(array); - RecordBatch::try_new(schema, vec![array_ref]).unwrap() - } -} - impl Stream for CongestedStream { type Item = Result; fn poll_next( @@ -289,34 +267,15 @@ impl Stream for CongestedStream { 0 => { let cleared = self.congestion_cleared.lock().unwrap(); if *cleared { - std::mem::drop(cleared); - if self.offset.0 == self.limit { - return Poll::Ready(None); - } - let batch = CongestedStream::create_record_batch( - Arc::clone(&self.schema), - 0, - self.batch_size, - ); - self.offset.0 += self.batch_size; - Poll::Ready(Some(Ok(batch))) + return Poll::Ready(None); } else { Poll::Pending } } 1 => { - if self.offset.1 == self.limit { - return Poll::Ready(None); - } - let batch = CongestedStream::create_record_batch( - Arc::clone(&self.schema), - 0, - self.batch_size, - ); - self.offset.1 += self.batch_size; let mut cleared = self.congestion_cleared.lock().unwrap(); *cleared = true; - Poll::Ready(Some(Ok(batch))) + Poll::Ready(None) } _ => unreachable!(), } @@ -335,9 +294,7 @@ async fn test_spm_congestion() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); let source = CongestedExec { schema: schema.clone(), - batch_size: 1, cache: CongestedExec::compute_properties(Arc::new(schema.clone())), - limit: 1_000, congestion_cleared: Arc::new(Mutex::new(false)), }; let spm = SortPreservingMergeExec::new( @@ -349,7 +306,7 @@ async fn test_spm_congestion() -> Result<()> { ); let spm_task = SpawnedTask::spawn(collect(Arc::new(spm), task_ctx)); - let result = timeout(Duration::from_secs(1), spm_task.join()).await; + let result = timeout(Duration::from_secs(3), spm_task.join()).await; match result { Ok(Ok(Ok(_batches))) => Ok(()), Ok(Ok(Err(e))) => Err(e), diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 3c6b082084166..90c27e916e60c 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -164,8 +164,7 @@ impl SortPreservingMergeStream { } // try to initialize the loser tree if self.loser_tree.is_empty() { - // Ensure all non-exhausted streams have a cursor from which - // rows can be pulled + // Ensure all non-exhausted streams have a cursor from which rows can be pulled let remaining_partitions = self.uninitiated_partitions.clone(); for i in remaining_partitions { match self.maybe_poll_stream(cx, i) { @@ -175,6 +174,7 @@ impl SortPreservingMergeStream { } Poll::Pending => { self.uninitiated_partitions.rotate_left(1); + cx.waker().wake_by_ref(); return Poll::Pending; } _ => { From 2e959231fa861ef67dff61593a259772ed6c98c8 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Tue, 3 Sep 2024 14:54:08 +0300 Subject: [PATCH 09/13] Update merge_fuzz.rs --- datafusion/core/tests/fuzz_cases/merge_fuzz.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 7b87ccde83176..30f4d6b8f3d56 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -260,14 +260,14 @@ pub struct CongestedStream { impl Stream for CongestedStream { type Item = Result; fn poll_next( - mut self: Pin<&mut Self>, + self: Pin<&mut Self>, _cx: &mut Context<'_>, ) -> Poll> { match self.partition { 0 => { let cleared = self.congestion_cleared.lock().unwrap(); if *cleared { - return Poll::Ready(None); + Poll::Ready(None) } else { Poll::Pending } From 9f32d2b7005bedc364b53739c66b815e9d99fe77 Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Wed, 4 Sep 2024 10:13:11 +0300 Subject: [PATCH 10/13] Update merge.rs --- datafusion/physical-plan/src/sorts/merge.rs | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 90c27e916e60c..811c93ce5892d 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -89,7 +89,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 @@ -101,8 +101,10 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, - /// Unitiated partitions. They are stored in a vector to keep them in - /// a priortiy order to visit the partitions in a round-robin fashion + /// This vector 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: Vec, } From c8b32a5b7b3fa972cbd52e7071217c89063c67ad Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Fri, 6 Sep 2024 10:29:30 +0300 Subject: [PATCH 11/13] Addresses the latest review --- .../core/tests/fuzz_cases/merge_fuzz.rs | 161 +---------------- datafusion/physical-plan/src/sorts/merge.rs | 13 +- .../src/sorts/sort_preserving_merge.rs | 169 +++++++++++++++++- 3 files changed, 180 insertions(+), 163 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 30f4d6b8f3d56..0ee8faeca27cf 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -17,19 +17,13 @@ //! Fuzz Test for various corner cases merging streams of RecordBatches -use std::any::Any; -use std::fmt::Formatter; -use std::pin::Pin; -use std::sync::{Arc, Mutex}; -use std::task::{Context, Poll}; -use std::time::Duration; +use std::sync::Arc; use arrow::{ array::{ArrayRef, Int32Array}, compute::SortOptions, record_batch::RecordBatch, }; -use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::physical_plan::{ collect, expressions::{col, PhysicalSortExpr}, @@ -37,20 +31,10 @@ use datafusion::physical_plan::{ sorts::sort_preserving_merge::SortPreservingMergeExec, }; use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_common::{DataFusionError, Result}; -use datafusion_common_runtime::SpawnedTask; -use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_expr_common::physical_expr::PhysicalExpr; -use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, PlanProperties, -}; +use datafusion_execution::RecordBatchStream; +use datafusion_physical_plan::ExecutionPlan; use test_utils::{batches_to_vec, partitions_to_sorted_vec, stagger_batch_with_seed}; -use futures::Stream; -use tokio::time::timeout; - #[tokio::test] async fn test_merge_2() { run_merge_test(vec![ @@ -179,142 +163,3 @@ fn concat(mut v1: Vec, v2: Vec) -> Vec { v1.extend(v2); v1 } - -/// It returns pending for the 1st partition until the 2nd partition is polled. -#[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(expr.clone(), SortOptions::default())) - .collect::>()]); - let mode = ExecutionMode::Unbounded; - PlanProperties::new(eq_properties, Partitioning::Hash(columns, 2), 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()), - congestion_cleared: self.congestion_cleared.clone(), - 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 1st partition until the 2nd partition is polled. -#[derive(Debug)] -pub struct CongestedStream { - schema: SchemaRef, - congestion_cleared: Arc>, - partition: usize, -} - -impl Stream for CongestedStream { - type Item = Result; - fn poll_next( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - match self.partition { - 0 => { - let cleared = self.congestion_cleared.lock().unwrap(); - if *cleared { - Poll::Ready(None) - } else { - Poll::Pending - } - } - 1 => { - 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(), - )), - } -} diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 811c93ce5892d..bd731488c3328 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -164,9 +164,10 @@ 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 let remaining_partitions = self.uninitiated_partitions.clone(); for i in remaining_partitions { match self.maybe_poll_stream(cx, i) { @@ -175,11 +176,19 @@ impl SortPreservingMergeStream { 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. self.uninitiated_partitions.rotate_left(1); + // 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); } } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 131fa71217cce..61140f0880bfa 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,155 @@ 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(expr.clone(), 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: self.congestion_cleared.clone(), + 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(), + )), + } + } } From 4bdbafa6bebaab40cdb50e96c05f98fd5feea80a Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Fri, 6 Sep 2024 10:32:35 +0300 Subject: [PATCH 12/13] fix clippy --- datafusion/core/tests/fuzz_cases/merge_fuzz.rs | 2 -- datafusion/physical-plan/src/sorts/sort_preserving_merge.rs | 6 ++++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 0ee8faeca27cf..4eb1070e6c857 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -31,8 +31,6 @@ use datafusion::physical_plan::{ sorts::sort_preserving_merge::SortPreservingMergeExec, }; use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_execution::RecordBatchStream; -use datafusion_physical_plan::ExecutionPlan; use test_utils::{batches_to_vec, partitions_to_sorted_vec, stagger_batch_with_seed}; #[tokio::test] diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 61140f0880bfa..4d333175bf758 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -1174,7 +1174,9 @@ mod tests { let mut eq_properties = EquivalenceProperties::new(schema); eq_properties.add_new_orderings(vec![columns .iter() - .map(|expr| PhysicalSortExpr::new(expr.clone(), SortOptions::default())) + .map(|expr| { + PhysicalSortExpr::new(Arc::clone(expr), SortOptions::default()) + }) .collect::>()]); let mode = ExecutionMode::Unbounded; PlanProperties::new(eq_properties, Partitioning::Hash(columns, 3), mode) @@ -1208,7 +1210,7 @@ mod tests { Ok(Box::pin(CongestedStream { schema: Arc::new(self.schema.clone()), none_polled_once: false, - congestion_cleared: self.congestion_cleared.clone(), + congestion_cleared: Arc::clone(&self.congestion_cleared), partition, })) } From e0bf209f095618961309b288b711c209b51b3f6c Mon Sep 17 00:00:00 2001 From: berkaysynnada Date: Fri, 6 Sep 2024 15:58:58 +0300 Subject: [PATCH 13/13] Use VecDeque for rotation --- datafusion/physical-plan/src/sorts/merge.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index bd731488c3328..875922ac34b54 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,6 +18,7 @@ //! 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}; @@ -101,11 +102,11 @@ pub(crate) struct SortPreservingMergeStream { /// number of rows produced produced: usize, - /// This vector contains partition indices in order. When a partition is polled and returns `Poll::Ready`, + /// 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: Vec, + uninitiated_partitions: VecDeque, } impl SortPreservingMergeStream { @@ -179,7 +180,10 @@ impl SortPreservingMergeStream { // 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. - self.uninitiated_partitions.rotate_left(1); + 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.