jayzhan211 commented on code in PR #24593:
URL: https://github.com/apache/datafusion/pull/24593#discussion_r3862885017


##########
datafusion/physical-plan/src/joins/sort_merge_join/exec.rs:
##########
@@ -707,82 +706,53 @@ impl ExecutionPlan for SortMergeJoinExec {
         }))
     }
 
-    /// Tries to swap the projection with its input [`SortMergeJoinExec`]. If 
it can be done,
-    /// it returns the new swapped version having the [`SortMergeJoinExec`] as 
the top plan.
-    /// Otherwise, it returns None.
+    /// Tries to push `projection` down through this join. If possible, 
returns a
+    /// new [`SortMergeJoinExec`] whose children are the projected inputs. 
Otherwise
+    /// the join applies the projection itself (see [`EmbeddedProjection`]).
     fn try_swapping_with_projection(
         &self,
         projection: &ProjectionExec,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         if self.projection.is_some() {
             return Ok(None);
         }
-        // Convert projected PhysicalExpr's to columns. If not possible, we 
cannot proceed.
-        let Some(projection_as_columns) = 
physical_to_column_exprs(projection.expr())
-        else {
-            return Ok(None);
-        };
 
-        let (far_right_left_col_ind, far_left_right_col_ind) = 
join_table_borders(
-            self.left().schema().fields().len(),
-            &projection_as_columns,
+        let schema = self.schema();
+        let (_, column_indices) = build_join_schema(
+            &self.left().schema(),
+            &self.right().schema(),
+            &self.join_type,
         );
 
-        // Pushing into the children needs each side's columns to stay 
together, which
-        // an arbitrary projection does not. The join can apply that one 
itself.
-        if !join_allows_pushdown(
-            &projection_as_columns,
-            &self.schema(),
-            far_right_left_col_ind,
-            far_left_right_col_ind,
-        ) {
-            return try_embed_projection(projection, self);
-        }
-
-        let left_field_size = self.left().schema().fields().len();
-        let left_projection = 
&projection_as_columns[0..=far_right_left_col_ind as usize];
-        let right_projection = &projection_as_columns[far_left_right_col_ind 
as usize..];
-
-        let Some(new_on) = update_join_on(
-            left_projection,
-            right_projection,
+        // Remaps the join keys and the filter's column indices to the
+        // projected children, and declines the pushdown if the projection
+        // drops a column the filter needs.
+        if let Some(JoinData {
+            projected_left_child,
+            projected_right_child,
+            join_filter,
+            join_on,
+        }) = try_pushdown_through_join_with_column_indices(

Review Comment:
   Thanks for your review, test added!



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to