feat: use spawned tasks to reduce call stack depth and avoid busy waiting

Author: pepijnve

datafusion/common/src/error.rs

@@ -350,6 +350,12 @@ impl From<GenericError> for DataFusionError {
     }
 }
 
+impl From<JoinError> for DataFusionError {
+    fn from(e: JoinError) -> Self {
+        DataFusionError::ExecutionJoin(e)
+    }
+}
+
 impl Display for DataFusionError {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         let error_prefix = self.error_prefix();

datafusion/physical-plan/src/joins/cross_join.rs

@@ -18,6 +18,7 @@
 //! Defines the cross join plan for loading the left side of the cross join
 //! and producing batches in parallel for the right partitions
 
+use futures::FutureExt;
 use std::{any::Any, sync::Arc, task::Poll};
 
 use super::utils::{
@@ -47,6 +48,7 @@ use datafusion_execution::TaskContext;
 use datafusion_physical_expr::equivalence::join_equivalence_properties;
 
 use async_trait::async_trait;
+use datafusion_common_runtime::SpawnedTask;
 use futures::{ready, Stream, StreamExt, TryStreamExt};
 
 /// Data of the left side that is buffered into memory
@@ -303,12 +305,13 @@ impl ExecutionPlan for CrossJoinExec {
 
         let left_fut = self.left_fut.try_once(|| {
             let left_stream = self.left.execute(0, context)?;
-
-            Ok(load_left_input(
-                left_stream,
-                join_metrics.clone(),
-                reservation,
-            ))
+            let task = load_left_input(left_stream, join_metrics.clone(), reservation);
+            Ok(async move {
+                // Spawn a task the first time the stream is polled for the build phase.
+                // This ensures the consumer of the join does not poll unnecessarily
+                // while the build is ongoing
+                SpawnedTask::spawn(task).map(|r| r?).await
+            })
         })?;
 
         if enforce_batch_size_in_joins {

datafusion/physical-plan/src/joins/hash_join.rs

@@ -82,8 +82,9 @@ use datafusion_physical_expr::PhysicalExprRef;
 use datafusion_physical_expr_common::datum::compare_op_for_nested;
 
 use ahash::RandomState;
+use datafusion_common_runtime::SpawnedTask;
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
-use futures::{ready, Stream, StreamExt, TryStreamExt};
+use futures::{ready, FutureExt, Stream, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 
 /// Hard-coded seed to ensure hash values from the hash join differ from `RepartitionExec`, avoiding collisions.
@@ -810,15 +811,22 @@ impl ExecutionPlan for HashJoinExec {
                 let reservation =
                     MemoryConsumer::new("HashJoinInput").register(context.memory_pool());
 
-                Ok(collect_left_input(
+                let task = collect_left_input(
                     self.random_state.clone(),
                     left_stream,
                     on_left.clone(),
                     join_metrics.clone(),
                     reservation,
                     need_produce_result_in_final(self.join_type),
                     self.right().output_partitioning().partition_count(),
-                ))
+                );
+
+                Ok(async move {
+                    // Spawn a task the first time the stream is polled for the build phase.
+                    // This ensures the consumer of the join does not poll unnecessarily
+                    // while the build is ongoing
+                    SpawnedTask::spawn(task).map(|r| r?).await
+                })
             })?,
             PartitionMode::Partitioned => {
                 let left_stream = self.left.execute(partition, Arc::clone(&context))?;
@@ -827,15 +835,22 @@ impl ExecutionPlan for HashJoinExec {
                     MemoryConsumer::new(format!("HashJoinInput[{partition}]"))
                         .register(context.memory_pool());
 
-                OnceFut::new(collect_left_input(
+                let task = collect_left_input(
                     self.random_state.clone(),
                     left_stream,
                     on_left.clone(),
                     join_metrics.clone(),
                     reservation,
                     need_produce_result_in_final(self.join_type),
                     1,
-                ))
+                );
+
+                OnceFut::new(async move {
+                    // Spawn a task the first time the stream is polled for the build phase.
+                    // This ensures the consumer of the join does not poll unnecessarily
+                    // while the build is ongoing
+                    SpawnedTask::spawn(task).map(|r| r?).await
+                })
             }
             PartitionMode::Auto => {
                 return plan_err!(

datafusion/physical-plan/src/joins/nested_loop_join.rs

@@ -17,6 +17,7 @@
 
 //! [`NestedLoopJoinExec`]: joins without equijoin (equality predicates).
 
+use futures::FutureExt;
 use std::any::Any;
 use std::fmt::Formatter;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -61,6 +62,7 @@ use datafusion_physical_expr::equivalence::{
     join_equivalence_properties, ProjectionMapping,
 };
 
+use datafusion_common_runtime::SpawnedTask;
 use futures::{ready, Stream, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 
@@ -499,13 +501,19 @@ impl ExecutionPlan for NestedLoopJoinExec {
         let inner_table = self.inner_table.try_once(|| {
             let stream = self.left.execute(0, Arc::clone(&context))?;
 
-            Ok(collect_left_input(
+            let task = collect_left_input(
                 stream,
                 join_metrics.clone(),
                 load_reservation,
                 need_produce_result_in_final(self.join_type),
                 self.right().output_partitioning().partition_count(),
-            ))
+            );
+            Ok(async move {
+                // Spawn a task the first time the stream is polled for the build phase.
+                // This ensures the consumer of the join does not poll unnecessarily
+                // while the build is ongoing
+                SpawnedTask::spawn(task).map(|r| r?).await
+            })
         })?;
 
         let batch_size = context.session_config().batch_size();

datafusion/physical-plan/src/sorts/sort.rs

@@ -54,6 +54,7 @@ use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::LexOrdering;
 
+use datafusion_common_runtime::SpawnedTask;
 use futures::{StreamExt, TryStreamExt};
 use log::{debug, trace};
 
@@ -1126,15 +1127,22 @@ impl ExecutionPlan for SortExec {
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
                     futures::stream::once(async move {
-                        while let Some(batch) = input.next().await {
-                            let batch = batch?;
-                            topk.insert_batch(batch)?;
-                            if topk.finished {
-                                break;
+                        // Spawn a task the first time the stream is polled for the sort phase.
+                        // This ensures the consumer of the sort does not poll unnecessarily
+                        // while the sort is ongoing
+                        SpawnedTask::spawn(async move {
+                            while let Some(batch) = input.next().await {
+                                let batch = batch?;
+                                topk.insert_batch(batch)?;
+                                if topk.finished {
+                                    break;
+                                }
                             }
-                        }
-                        topk.emit()
+                            topk.emit()
+                        })
+                        .await
                     })
+                    .map(|s| s?)
                     .try_flatten(),
                 )))
             }
@@ -1152,12 +1160,19 @@ impl ExecutionPlan for SortExec {
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
                     futures::stream::once(async move {
-                        while let Some(batch) = input.next().await {
-                            let batch = batch?;
-                            sorter.insert_batch(batch).await?;
-                        }
-                        sorter.sort().await
+                        // Spawn a task the first time the stream is polled for the sort phase.
+                        // This ensures the consumer of the sort does not poll unnecessarily
+                        // while the sort is ongoing
+                        SpawnedTask::spawn(async move {
+                            while let Some(batch) = input.next().await {
+                                let batch = batch?;
+                                sorter.insert_batch(batch).await?;
+                            }
+                            sorter.sort().await
+                        })
+                        .await
                     })
+                    .map(|s| s?)
                     .try_flatten(),
                 )))
             }