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


##########
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:
   Nice improvement here. One optional thought: could we add an execution-level 
unit test for the successful remapping path, where the filter columns are 
retained but reordered? The SQL regression covers the fallback with an embedded 
projection, while the current unit test verifies the rewritten indices without 
executing the returned `SortMergeJoinExec`. Executing it would give us a bit 
more confidence that the remapped `JoinFilter` schema and columns are consumed 
correctly. The broader SortMergeJoin spill SQL tests already exercise filtered 
projected joins, so I see this as a small coverage improvement rather than 
something that should block this PR.



-- 
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