You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by al...@apache.org on 2024/09/06 14:10:07 UTC
(datafusion) branch main updated: Fix Possible Congestion Scenario in `SortPreservingMergeExec` (#12302)
This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new f2a8b072a7 Fix Possible Congestion Scenario in `SortPreservingMergeExec` (#12302)
f2a8b072a7 is described below
commit f2a8b072a70fce9061874ef99fe670564d197740
Author: Berkay Şahin <12...@users.noreply.github.com>
AuthorDate: Fri Sep 6 17:10:02 2024 +0300
Fix Possible Congestion Scenario in `SortPreservingMergeExec` (#12302)
* 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
* Add comments
* Increase threshold
* Debug
* Remove merge.rs changes
* Use waker
* Update merge.rs
* Simplify the test
* Update merge_fuzz.rs
* Update merge.rs
* Addresses the latest review
* fix clippy
* Use VecDeque for rotation
---
datafusion/core/tests/fuzz_cases/merge_fuzz.rs | 1 +
datafusion/physical-plan/src/sorts/merge.rs | 55 +++++--
.../src/sorts/sort_preserving_merge.rs | 171 ++++++++++++++++++++-
3 files changed, 213 insertions(+), 14 deletions(-)
diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs
index 95cd75f50a..4eb1070e6c 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 85418ff361..875922ac34 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<C> = Box<dyn PartitionedStream<Output = Result<(C, RecordBatch)>>>;
@@ -86,7 +90,7 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues> {
/// 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<C: CursorValues> {
/// 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<usize>,
}
impl<C: CursorValues> SortPreservingMergeStream<C> {
@@ -121,6 +131,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
batch_size,
fetch,
produced: 0,
+ uninitiated_partitions: (0..stream_count).collect(),
}
}
@@ -154,14 +165,36 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
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 131fa71217..4d333175bf 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<Mutex<bool>>,
+ }
+
+ 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<dyn PhysicalExpr>)
+ .collect::<Vec<_>>();
+ 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::<Vec<_>>()]);
+ 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<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn execute(
+ &self,
+ partition: usize,
+ _context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ 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<Mutex<bool>>,
+ partition: usize,
+ }
+
+ impl Stream for CongestedStream {
+ type Item = Result<RecordBatch>;
+ fn poll_next(
+ mut self: Pin<&mut Self>,
+ _cx: &mut Context<'_>,
+ ) -> Poll<Option<Self::Item>> {
+ 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(),
+ )),
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org