Parallelise collect_partitioned

datafusion/core/src/datasource/memory.rs

@@ -17,7 +17,7 @@
 
 //! [`MemTable`] for querying `Vec<RecordBatch>` by DataFusion.
 
-use futures::{StreamExt, TryStreamExt};
+use futures::StreamExt;
 use std::any::Any;
 use std::sync::Arc;
 
@@ -26,17 +26,16 @@ use arrow::record_batch::RecordBatch;
 use async_trait::async_trait;
 use datafusion_expr::LogicalPlan;
 use tokio::sync::RwLock;
-use tokio::task;
 
 use crate::datasource::{TableProvider, TableType};
 use crate::error::{DataFusionError, Result};
 use crate::execution::context::SessionState;
 use crate::logical_expr::Expr;
 use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
-use crate::physical_plan::common;
 use crate::physical_plan::common::AbortOnDropSingle;
 use crate::physical_plan::memory::MemoryExec;
 use crate::physical_plan::ExecutionPlan;
+use crate::physical_plan::{collect_partitioned, common};
 use crate::physical_plan::{repartition::RepartitionExec, Partitioning};
 
 /// In-memory data source for presenting a `Vec<RecordBatch>` as a
@@ -199,25 +198,7 @@ impl TableProvider for MemTable {
             )?)
         };
 
-        // Get the task context from the session state.
-        let task_ctx = state.task_ctx();
-
-        // Execute the plan and collect the results into batches.
-        let mut tasks = vec![];
-        for idx in 0..plan.output_partitioning().partition_count() {
-            let stream = plan.execute(idx, task_ctx.clone())?;
-            let handle = task::spawn(async move {
-                stream.try_collect().await.map_err(DataFusionError::from)
-            });
-            tasks.push(AbortOnDropSingle::new(handle));
-        }
-        let results = futures::future::join_all(tasks)
-            .await
-            .into_iter()
-            .map(|result| {
-                result.map_err(|e| DataFusionError::Execution(format!("{e}")))?
-            })
-            .collect::<Result<Vec<Vec<RecordBatch>>>>()?;
+        let results = collect_partitioned(plan, state.task_ctx()).await?;
 
         // Write the results into the table.
         let mut all_batches = self.batches.write().await;

datafusion/core/src/physical_plan/mod.rs

@@ -33,7 +33,7 @@ pub use datafusion_expr::Accumulator;
 pub use datafusion_expr::ColumnarValue;
 pub use datafusion_physical_expr::aggregate::row_accumulator::RowAccumulator;
 pub use display::DisplayFormatType;
-use futures::stream::Stream;
+use futures::stream::{Stream, TryStreamExt};
 use std::fmt;
 use std::fmt::Debug;
 
@@ -443,11 +443,21 @@ pub async fn collect_partitioned(
     context: Arc<TaskContext>,
 ) -> Result<Vec<Vec<RecordBatch>>> {
     let streams = execute_stream_partitioned(plan, context)?;
-    let mut batches = Vec::with_capacity(streams.len());
-    for stream in streams {
-        batches.push(common::collect(stream).await?);
-    }
-    Ok(batches)
+
+    // Execute the plan and collect the results into batches.
+    let handles = streams
+        .into_iter()
+        .enumerate()
+        .map(|(idx, stream)| async move {
+            let handle = tokio::task::spawn(stream.try_collect());
+            AbortOnDropSingle::new(handle).await.map_err(|e| {
+                DataFusionError::Execution(format!(
+                    "collect_partitioned partition {idx} panicked: {e}"
+                ))
+            })?
+        });
+
+    futures::future::try_join_all(handles).await
 }
 
 /// Execute the [ExecutionPlan] and return a vec with one stream per output partition
@@ -665,6 +675,7 @@ pub mod values;
 pub mod windows;
 
 use crate::execution::context::TaskContext;
+use crate::physical_plan::common::AbortOnDropSingle;
 use crate::physical_plan::repartition::RepartitionExec;
 use crate::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 pub use datafusion_physical_expr::{