fix: wrap join operators with cooperative() for cancellation support

Author: erratic-pattern

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

@@ -25,6 +25,7 @@ use super::utils::{
     BatchTransformer, BuildProbeJoinMetrics, NoopBatchTransformer, OnceAsync, OnceFut,
     StatefulStreamResult,
 };
+use crate::coop::cooperative;
 use crate::execution_plan::{boundedness_from_children, EmissionType};
 use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use crate::projection::{
@@ -324,7 +325,7 @@ impl ExecutionPlan for CrossJoinExec {
         })?;
 
         if enforce_batch_size_in_joins {
-            Ok(Box::pin(CrossJoinStream {
+            Ok(Box::pin(cooperative(CrossJoinStream {
                 schema: Arc::clone(&self.schema),
                 left_fut,
                 right: stream,
@@ -333,9 +334,9 @@ impl ExecutionPlan for CrossJoinExec {
                 state: CrossJoinStreamState::WaitBuildSide,
                 left_data: RecordBatch::new_empty(self.left().schema()),
                 batch_transformer: BatchSplitter::new(batch_size),
-            }))
+            })))
         } else {
-            Ok(Box::pin(CrossJoinStream {
+            Ok(Box::pin(cooperative(CrossJoinStream {
                 schema: Arc::clone(&self.schema),
                 left_fut,
                 right: stream,
@@ -344,7 +345,7 @@ impl ExecutionPlan for CrossJoinExec {
                 state: CrossJoinStreamState::WaitBuildSide,
                 left_data: RecordBatch::new_empty(self.left().schema()),
                 batch_transformer: NoopBatchTransformer::new(),
-            }))
+            })))
         }
     }
 

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

@@ -33,6 +33,7 @@ use super::{
     PartitionMode, SharedBitmapBuilder,
 };
 use super::{JoinOn, JoinOnRef};
+use crate::coop::cooperative;
 use crate::execution_plan::{boundedness_from_children, EmissionType};
 use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64};
 use crate::projection::{
@@ -880,7 +881,7 @@ impl ExecutionPlan for HashJoinExec {
             None => self.column_indices.clone(),
         };
 
-        Ok(Box::pin(HashJoinStream {
+        Ok(Box::pin(cooperative(HashJoinStream {
             schema: self.schema(),
             on_right,
             filter: self.filter.clone(),
@@ -895,7 +896,7 @@ impl ExecutionPlan for HashJoinExec {
             batch_size,
             hashes_buffer: vec![],
             right_side_ordered: self.right.output_ordering().is_some(),
-        }))
+        })))
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

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

@@ -30,6 +30,7 @@ use super::utils::{
     StatefulStreamResult,
 };
 use crate::common::can_project;
+use crate::coop::cooperative;
 use crate::execution_plan::{boundedness_from_children, EmissionType};
 use crate::joins::utils::{
     adjust_indices_by_join_type, apply_join_filter_to_indices, build_batch_from_indices,
@@ -530,7 +531,7 @@ impl ExecutionPlan for NestedLoopJoinExec {
             None => self.column_indices.clone(),
         };
 
-        Ok(Box::pin(NestedLoopJoinStream {
+        Ok(Box::pin(cooperative(NestedLoopJoinStream {
             schema: self.schema(),
             filter: self.filter.clone(),
             join_type: self.join_type,
@@ -544,7 +545,7 @@ impl ExecutionPlan for NestedLoopJoinExec {
             left_data: None,
             join_result_status: None,
             intermediate_batch_size: batch_size,
-        }))
+        })))
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

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

@@ -33,6 +33,7 @@ use std::sync::atomic::Ordering::Relaxed;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+use crate::coop::cooperative;
 use crate::execution_plan::{boundedness_from_children, EmissionType};
 use crate::expressions::PhysicalSortExpr;
 use crate::joins::utils::{
@@ -501,7 +502,7 @@ impl ExecutionPlan for SortMergeJoinExec {
             .register(context.memory_pool());
 
         // create join stream
-        Ok(Box::pin(SortMergeJoinStream::try_new(
+        Ok(Box::pin(cooperative(SortMergeJoinStream::try_new(
             context.session_config().spill_compression(),
             Arc::clone(&self.schema),
             self.sort_options.clone(),
@@ -516,7 +517,7 @@ impl ExecutionPlan for SortMergeJoinExec {
             SortMergeJoinMetrics::new(partition, &self.metrics),
             reservation,
             context.runtime_env(),
-        )?))
+        )?)))
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

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

@@ -33,6 +33,7 @@ use std::task::{Context, Poll};
 use std::vec;
 
 use crate::common::SharedMemoryReservation;
+use crate::coop::cooperative;
 use crate::execution_plan::{boundedness_from_children, emission_type_from_children};
 use crate::joins::hash_join::{equal_rows_arr, update_hash};
 use crate::joins::stream_join_utils::{
@@ -533,7 +534,7 @@ impl ExecutionPlan for SymmetricHashJoinExec {
         }
 
         if enforce_batch_size_in_joins {
-            Ok(Box::pin(SymmetricHashJoinStream {
+            Ok(Box::pin(cooperative(SymmetricHashJoinStream {
                 left_stream,
                 right_stream,
                 schema: self.schema(),
@@ -551,9 +552,9 @@ impl ExecutionPlan for SymmetricHashJoinExec {
                 state: SHJStreamState::PullRight,
                 reservation,
                 batch_transformer: BatchSplitter::new(batch_size),
-            }))
+            })))
         } else {
-            Ok(Box::pin(SymmetricHashJoinStream {
+            Ok(Box::pin(cooperative(SymmetricHashJoinStream {
                 left_stream,
                 right_stream,
                 schema: self.schema(),
@@ -571,7 +572,7 @@ impl ExecutionPlan for SymmetricHashJoinExec {
                 state: SHJStreamState::PullRight,
                 reservation,
                 batch_transformer: NoopBatchTransformer::new(),
-            }))
+            })))
         }
     }