Parallel merge sort (#6162) (#6308)

* Parallel merge sort (#6162)

* Fix test

Author: tustvold

datafusion/core/src/physical_plan/common.rs

@@ -21,6 +21,7 @@ use super::SendableRecordBatchStream;
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::TaskContext;
 use crate::execution::memory_pool::MemoryReservation;
+use crate::physical_plan::stream::RecordBatchReceiverStream;
 use crate::physical_plan::{displayable, ColumnStatistics, ExecutionPlan, Statistics};
 use arrow::datatypes::Schema;
 use arrow::ipc::writer::{FileWriter, IpcWriteOptions};
@@ -131,6 +132,31 @@ pub(crate) fn spawn_execution(
     })
 }
 
+/// If running in a tokio context spawns the execution of `stream` to a separate task
+/// allowing it to execute in parallel with an intermediate buffer of size `buffer`
+pub(crate) fn spawn_buffered(
+    mut input: SendableRecordBatchStream,
+    buffer: usize,
+) -> SendableRecordBatchStream {
+    // Use tokio only if running from a tokio context (#2201)
+    let handle = match tokio::runtime::Handle::try_current() {
+        Ok(handle) => handle,
+        Err(_) => return input,
+    };
+
+    let schema = input.schema();
+    let (sender, receiver) = mpsc::channel(buffer);
+    let join = handle.spawn(async move {
+        while let Some(item) = input.next().await {
+            if sender.send(item).await.is_err() {
+                return;
+            }
+        }
+    });
+
+    RecordBatchReceiverStream::create(&schema, receiver, join)
+}
+
 /// Computes the statistics for an in-memory RecordBatch
 ///
 /// Only computes statistics that are in arrows metadata (num rows, byte size and nulls)

datafusion/core/src/physical_plan/sorts/sort.rs

@@ -25,7 +25,7 @@ use crate::execution::memory_pool::{
     human_readable_size, MemoryConsumer, MemoryReservation,
 };
 use crate::execution::runtime_env::RuntimeEnv;
-use crate::physical_plan::common::{batch_byte_size, IPCWriter};
+use crate::physical_plan::common::{batch_byte_size, spawn_buffered, IPCWriter};
 use crate::physical_plan::expressions::PhysicalSortExpr;
 use crate::physical_plan::metrics::{
     BaselineMetrics, CompositeMetricsSet, MemTrackingMetrics, MetricsSet,
@@ -284,11 +284,13 @@ impl ExternalSorter {
                     self.partition_id,
                     &self.runtime.memory_pool,
                 );
-                sort_batch_stream(batch, self.expr.clone(), self.fetch, metrics)
+                Ok(spawn_buffered(
+                    sort_batch_stream(batch, self.expr.clone(), self.fetch, metrics)?,
+                    1,
+                ))
             })
             .collect::<Result<_>>()?;
 
-        // TODO: Run batch sorts concurrently (#6162)
         // TODO: Pushdown fetch to streaming merge (#6000)
 
         streaming_merge(

datafusion/core/src/physical_plan/sorts/sort_preserving_merge.rs

@@ -22,18 +22,17 @@ use std::sync::Arc;
 
 use arrow::datatypes::SchemaRef;
 use log::{debug, trace};
-use tokio::sync::mpsc;
 
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::TaskContext;
+use crate::physical_plan::common::spawn_buffered;
 use crate::physical_plan::metrics::{
     ExecutionPlanMetricsSet, MemTrackingMetrics, MetricsSet,
 };
 use crate::physical_plan::sorts::streaming_merge;
-use crate::physical_plan::stream::RecordBatchReceiverStream;
 use crate::physical_plan::{
-    common::spawn_execution, expressions::PhysicalSortExpr, DisplayFormatType,
-    Distribution, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
+    expressions::PhysicalSortExpr, DisplayFormatType, Distribution, ExecutionPlan,
+    Partitioning, SendableRecordBatchStream, Statistics,
 };
 use datafusion_physical_expr::{EquivalenceProperties, PhysicalSortRequirement};
 
@@ -181,29 +180,12 @@ impl ExecutionPlan for SortPreservingMergeExec {
                 result
             }
             _ => {
-                // Use tokio only if running from a tokio context (#2201)
-                let receivers = match tokio::runtime::Handle::try_current() {
-                    Ok(_) => (0..input_partitions)
-                        .map(|part_i| {
-                            let (sender, receiver) = mpsc::channel(1);
-                            let join_handle = spawn_execution(
-                                self.input.clone(),
-                                sender,
-                                part_i,
-                                context.clone(),
-                            );
-
-                            RecordBatchReceiverStream::create(
-                                &schema,
-                                receiver,
-                                join_handle,
-                            )
-                        })
-                        .collect(),
-                    Err(_) => (0..input_partitions)
-                        .map(|partition| self.input.execute(partition, context.clone()))
-                        .collect::<Result<_>>()?,
-                };
+                let receivers = (0..input_partitions)
+                    .map(|partition| {
+                        let stream = self.input.execute(partition, context.clone())?;
+                        Ok(spawn_buffered(stream, 1))
+                    })
+                    .collect::<Result<_>>()?;
 
                 debug!("Done setting up sender-receiver for SortPreservingMergeExec::execute");
 
@@ -262,6 +244,7 @@ mod tests {
     use crate::physical_plan::memory::MemoryExec;
     use crate::physical_plan::metrics::MetricValue;
     use crate::physical_plan::sorts::sort::SortExec;
+    use crate::physical_plan::stream::RecordBatchReceiverStream;
     use crate::physical_plan::{collect, common};
     use crate::prelude::{SessionConfig, SessionContext};
     use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec};
@@ -812,7 +795,7 @@ mod tests {
         let mut streams = Vec::with_capacity(partition_count);
 
         for partition in 0..partition_count {
-            let (sender, receiver) = mpsc::channel(1);
+            let (sender, receiver) = tokio::sync::mpsc::channel(1);
             let mut stream = batches.execute(partition, task_ctx.clone()).unwrap();
             let join_handle = tokio::spawn(async move {
                 while let Some(batch) = stream.next().await {