Fix RightSemi joins

Author: Vedin

datafusion/physical-optimizer/src/join_selection.rs

@@ -27,12 +27,12 @@ use crate::PhysicalOptimizerRule;
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::error::Result;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
-use datafusion_common::{internal_err, JoinSide, JoinType};
+use datafusion_common::{internal_err, DataFusionError, JoinSide, JoinType};
 use datafusion_expr_common::sort_properties::SortProperties;
 use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::LexOrdering;
 use datafusion_physical_plan::execution_plan::EmissionType;
-use datafusion_physical_plan::joins::utils::ColumnIndex;
+use datafusion_physical_plan::joins::utils::{check_join_is_valid, ColumnIndex};
 use datafusion_physical_plan::joins::{
     CrossJoinExec, GraceHashJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode,
     StreamJoinPartitionMode, SymmetricHashJoinExec,
@@ -173,6 +173,13 @@ pub(crate) fn try_collect_left(
     let left = hash_join.left();
     let right = hash_join.right();
 
+    // Skip collect-left rewrite if the join currently has inconsistent schemas (e.g. required
+    // columns were projected away temporarily). This mirrors the legacy hash join behavior where
+    // collect-left is only attempted once the join inputs are fully valid.
+    if check_join_is_valid(&left.schema(), &right.schema(), hash_join.on()).is_err() {
+        return Ok(None);
+    }
+
     let left_can_collect = ignore_threshold
         || supports_collect_by_thresholds(
             &**left,
@@ -191,33 +198,23 @@ pub(crate) fn try_collect_left(
             if hash_join.join_type().supports_swap()
                 && should_swap_join_order(&**left, &**right)?
             {
-                Ok(Some(hash_join.swap_inputs(PartitionMode::CollectLeft)?))
+                match hash_join.swap_inputs(PartitionMode::CollectLeft) {
+                    Ok(plan) => Ok(Some(plan)),
+                    Err(err) if is_missing_join_columns(&err) => Ok(None),
+                    Err(err) => Err(err),
+                }
             } else {
-                Ok(Some(Arc::new(HashJoinExec::try_new(
-                    Arc::clone(left),
-                    Arc::clone(right),
-                    hash_join.on().to_vec(),
-                    hash_join.filter().cloned(),
-                    hash_join.join_type(),
-                    hash_join.projection.clone(),
-                    PartitionMode::CollectLeft,
-                    hash_join.null_equality(),
-                )?)))
+                build_collect_left_exec(hash_join, left, right)
             }
         }
-        (true, false) => Ok(Some(Arc::new(HashJoinExec::try_new(
-            Arc::clone(left),
-            Arc::clone(right),
-            hash_join.on().to_vec(),
-            hash_join.filter().cloned(),
-            hash_join.join_type(),
-            hash_join.projection.clone(),
-            PartitionMode::CollectLeft,
-            hash_join.null_equality(),
-        )?))),
+        (true, false) => build_collect_left_exec(hash_join, left, right),
         (false, true) => {
             if hash_join.join_type().supports_swap() {
-                hash_join.swap_inputs(PartitionMode::CollectLeft).map(Some)
+                match hash_join.swap_inputs(PartitionMode::CollectLeft) {
+                    Ok(plan) => Ok(Some(plan)),
+                    Err(err) if is_missing_join_columns(&err) => Ok(None),
+                    Err(err) => Err(err),
+                }
             } else {
                 Ok(None)
             }
@@ -226,6 +223,35 @@ pub(crate) fn try_collect_left(
     }
 }
 
+fn is_missing_join_columns(err: &DataFusionError) -> bool {
+    matches!(
+        err,
+        DataFusionError::Plan(msg)
+            if msg.contains("The left or right side of the join does not have all columns")
+    )
+}
+
+fn build_collect_left_exec(
+    hash_join: &HashJoinExec,
+    left: &Arc<dyn ExecutionPlan>,
+    right: &Arc<dyn ExecutionPlan>,
+) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+    match HashJoinExec::try_new(
+        Arc::clone(left),
+        Arc::clone(right),
+        hash_join.on().to_vec(),
+        hash_join.filter().cloned(),
+        hash_join.join_type(),
+        hash_join.projection.clone(),
+        PartitionMode::CollectLeft,
+        hash_join.null_equality(),
+    ) {
+        Ok(exec) => Ok(Some(Arc::new(exec))),
+        Err(err) if is_missing_join_columns(&err) => Ok(None),
+        Err(err) => Err(err),
+    }
+}
+
 /// Creates a partitioned hash join execution plan, swapping inputs if beneficial.
 ///
 /// Checks if the join order should be swapped based on the join type and input statistics.

datafusion/physical-plan/src/joins/grace_hash_join/exec.rs

@@ -53,16 +53,13 @@ use datafusion_common::{
     internal_err, plan_err, project_schema, JoinSide, JoinType, NullEquality, Result,
 };
 use datafusion_execution::TaskContext;
-use datafusion_functions_aggregate_common::min_max::{MaxAccumulator, MinAccumulator};
 use datafusion_physical_expr::equivalence::{
     join_equivalence_properties, ProjectionMapping,
 };
 use datafusion_physical_expr::expressions::{lit, DynamicFilterPhysicalExpr};
 use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef};
 
-use crate::joins::grace_hash_join::stream::{
-    GraceAccumulator, GraceHashJoinStream, SpillFut,
-};
+use crate::joins::grace_hash_join::stream::{GraceHashJoinStream, SpillFut};
 use crate::joins::hash_join::shared_bounds::SharedBoundsAccumulator;
 use crate::metrics::SpillMetrics;
 use crate::spill::spill_manager::SpillLocation;
@@ -105,7 +102,6 @@ pub struct GraceHashJoinExec {
     /// Set when dynamic filter pushdown is detected in handle_child_pushdown_result.
     /// HashJoinExec also needs to keep a shared bounds accumulator for coordinating updates.
     dynamic_filter: Option<HashJoinExecDynamicFilter>,
-    accumulator: Arc<GraceAccumulator>,
 }
 
 #[derive(Clone)]
@@ -183,9 +179,6 @@ impl GraceHashJoinExec {
             &on,
             projection.as_ref(),
         )?;
-        let partitions = left.output_partitioning().partition_count();
-        let accumulator = GraceAccumulator::new(partitions);
-
         let metrics = ExecutionPlanMetricsSet::new();
         // Initialize both dynamic filter and bounds accumulator to None
         // They will be set later if dynamic filtering is enabled
@@ -203,7 +196,6 @@ impl GraceHashJoinExec {
             null_equality,
             cache,
             dynamic_filter: None,
-            accumulator,
         })
     }
 
@@ -547,7 +539,6 @@ impl ExecutionPlan for GraceHashJoinExec {
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        let new_partition_count = children[0].output_partitioning().partition_count();
         Ok(Arc::new(GraceHashJoinExec {
             left: Arc::clone(&children[0]),
             right: Arc::clone(&children[1]),
@@ -570,12 +561,10 @@ impl ExecutionPlan for GraceHashJoinExec {
             )?,
             // Keep the dynamic filter, bounds accumulator will be reset
             dynamic_filter: self.dynamic_filter.clone(),
-            accumulator: GraceAccumulator::new(new_partition_count),
         }))
     }
 
     fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> {
-        let partition_count = self.left.output_partitioning().partition_count();
         Ok(Arc::new(GraceHashJoinExec {
             left: Arc::clone(&self.left),
             right: Arc::clone(&self.right),
@@ -591,7 +580,6 @@ impl ExecutionPlan for GraceHashJoinExec {
             cache: self.cache.clone(),
             // Reset dynamic filter and bounds accumulator to initial state
             dynamic_filter: None,
-            accumulator: GraceAccumulator::new(partition_count),
         }))
     }
 
@@ -654,7 +642,6 @@ impl ExecutionPlan for GraceHashJoinExec {
         let on = self.on.clone();
         let spill_left_clone = Arc::clone(&spill_left);
         let spill_right_clone = Arc::clone(&spill_right);
-        let accumulator_clone = Arc::clone(&self.accumulator);
         let join_metrics_clone = Arc::clone(&join_metrics);
         let spill_fut = OnceFut::new(async move {
             let (left_idx, right_idx) = partition_and_spill(
@@ -670,14 +657,16 @@ impl ExecutionPlan for GraceHashJoinExec {
                 partition,
             )
             .await?;
-            accumulator_clone
-                .report_partition(partition, left_idx.clone(), right_idx.clone())
-                .await;
             Ok(SpillFut::new(partition, left_idx, right_idx))
         });
 
+        let left_input_schema = self.left.schema();
+        let right_input_schema = self.right.schema();
+
         Ok(Box::pin(GraceHashJoinStream::new(
             self.schema(),
+            left_input_schema,
+            right_input_schema,
             spill_fut,
             spill_left,
             spill_right,
@@ -689,7 +678,6 @@ impl ExecutionPlan for GraceHashJoinExec {
             column_indices_after_projection,
             join_metrics,
             context,
-            Arc::clone(&self.accumulator),
         )))
     }
 
@@ -845,7 +833,6 @@ impl ExecutionPlan for GraceHashJoinExec {
                         filter: dynamic_filter,
                         bounds_accumulator: OnceLock::new(),
                     }),
-                    accumulator: Arc::clone(&self.accumulator),
                 });
                 result = result.with_updated_node(new_node as Arc<dyn ExecutionPlan>);
             }