Use expression tree traversal order instead of sorting for complex expression column collection

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: Dandandan

datafusion/physical-plan/src/projection.rs

@@ -51,7 +51,6 @@ use datafusion_execution::TaskContext;
 use datafusion_expr::ExpressionPlacement;
 use datafusion_physical_expr::equivalence::ProjectionMapping;
 use datafusion_physical_expr::projection::Projector;
-use datafusion_physical_expr::utils::collect_columns;
 use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql};
 use datafusion_physical_expr_common::sort_expr::{
     LexOrdering, LexRequirement, PhysicalSortExpr,
@@ -1082,14 +1081,11 @@ fn try_unifying_projections(
 /// Collect all column indices from the given projection expressions.
 fn collect_column_indices(exprs: &[ProjectionExpr]) -> Vec<usize> {
     // Collect column indices in a deterministic order that preserves the
-    // projection's column ordering when possible. For simple Column
-    // expressions, we use the column index directly (preserving the
-    // projection's desired output order). For complex expressions with
-    // multiple column references, we sort indices for determinism since
-    // collect_columns returns a HashSet with non-deterministic iteration.
+    // projection's column ordering. For simple Column expressions, we use
+    // the column index directly. For complex expressions, we walk the
+    // expression tree to collect column references in traversal order.
     // This allows the embedded projection to match the desired output
-    // column order for simple column reorderings, avoiding a residual
-    // ProjectionExec.
+    // column order, avoiding a residual ProjectionExec.
     let mut seen = std::collections::HashSet::new();
     let mut indices = Vec::new();
     for proj_expr in exprs {
@@ -1099,18 +1095,20 @@ fn collect_column_indices(exprs: &[ProjectionExpr]) -> Vec<usize> {
                 indices.push(col.index());
             }
         } else {
-            // Complex expression: collect all referenced columns in sorted
-            // order for determinism.
-            let mut expr_indices: Vec<usize> = collect_columns(&proj_expr.expr)
-                .into_iter()
-                .map(|c| c.index())
-                .collect();
-            expr_indices.sort();
-            for idx in expr_indices {
-                if seen.insert(idx) {
-                    indices.push(idx);
-                }
-            }
+            // Complex expression: collect all referenced columns in
+            // expression tree traversal order (deterministic) to preserve
+            // the natural ordering of column references.
+            proj_expr
+                .expr
+                .apply(|expr| {
+                    if let Some(col) = expr.as_any().downcast_ref::<Column>() {
+                        if seen.insert(col.index()) {
+                            indices.push(col.index());
+                        }
+                    }
+                    Ok(TreeNodeRecursion::Continue)
+                })
+                .expect("closure always returns OK");
         }
     }
     indices
@@ -1226,7 +1224,8 @@ mod tests {
             expr,
             alias: "b-(1+a)".to_string(),
         }]);
-        assert_eq!(column_indices, vec![1, 7]);
+        // Tree traversal order: b@7 is visited before a@1
+        assert_eq!(column_indices, vec![7, 1]);
         Ok(())
     }
 

datafusion/sqllogictest/test_files/lateral_join.slt

@@ -527,8 +527,8 @@ logical_plan
 physical_plan
 01)SortPreservingMergeExec: [id@0 ASC NULLS LAST]
 02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true]
-03)----ProjectionExec: expr=[id@0 as id, CASE WHEN __always_true@2 IS NULL THEN 0 ELSE cnt@1 END as cnt]
-04)------HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, t1_id@1)], projection=[id@0, cnt@1, __always_true@3]
+03)----ProjectionExec: expr=[id@0 as id, CASE WHEN __always_true@1 IS NULL THEN 0 ELSE cnt@2 END as cnt]
+04)------HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, t1_id@1)], projection=[id@0, __always_true@3, cnt@1]
 05)--------DataSourceExec: partitions=1, partition_sizes=[1]
 06)--------ProjectionExec: expr=[count(Int64(1))@1 as cnt, t1_id@0 as t1_id, true as __always_true]
 07)----------AggregateExec: mode=FinalPartitioned, gby=[t1_id@0 as t1_id], aggr=[count(Int64(1))]

datafusion/sqllogictest/test_files/tpch/plans/q9.slt.part

@@ -81,8 +81,8 @@ physical_plan
 04)------AggregateExec: mode=FinalPartitioned, gby=[nation@0 as nation, o_year@1 as o_year], aggr=[sum(profit.amount)]
 05)--------RepartitionExec: partitioning=Hash([nation@0, o_year@1], 4), input_partitions=4
 06)----------AggregateExec: mode=Partial, gby=[nation@0 as nation, o_year@1 as o_year], aggr=[sum(profit.amount)]
-07)------------ProjectionExec: expr=[n_name@0 as nation, date_part(YEAR, o_orderdate@1) as o_year, l_extendedprice@3 * (Some(1),20,0 - l_discount@4) - ps_supplycost@5 * l_quantity@2 as amount]
-08)--------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@3, n_nationkey@0)], projection=[n_name@7, o_orderdate@5, l_quantity@0, l_extendedprice@1, l_discount@2, ps_supplycost@4]
+07)------------ProjectionExec: expr=[n_name@0 as nation, date_part(YEAR, o_orderdate@1) as o_year, l_extendedprice@2 * (Some(1),20,0 - l_discount@3) - ps_supplycost@4 * l_quantity@5 as amount]
+08)--------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@3, n_nationkey@0)], projection=[n_name@7, o_orderdate@5, l_extendedprice@1, l_discount@2, ps_supplycost@4, l_quantity@0]
 09)----------------RepartitionExec: partitioning=Hash([s_nationkey@3], 4), input_partitions=4
 10)------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_orderkey@0, o_orderkey@0)], projection=[l_quantity@1, l_extendedprice@2, l_discount@3, s_nationkey@4, ps_supplycost@5, o_orderdate@7]
 11)--------------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4