perf: collapse chained projections in a single pass

Long chains of `ProjectionExec`s (and their `LogicalPlan::Projection`
counterparts) — produced by e.g. a pipeline of `SELECT *, <expr> AS dN
FROM ...` stacked N times — were being collapsed one level at a time,
forcing O(N) intermediate `ProjectionExec` constructions (each
recomputing equivalence properties through the projection mapping) and
O(N) re-runs of the logical `OptimizeProjections` rule.

Physical: replace the pairwise recursion in
`ProjectionExec::try_swapping_with_projection` with a single iterative
chain collapse that walks all consecutive `ProjectionExec`s and builds
only one final `ProjectionExec`. Leaf pushdown into a non-`Projection`
input (e.g. `DataSourceExec`) is preserved by handing the unified
projection back to `remove_unnecessary_projections` once at the end.

Logical: wrap `merge_consecutive_projections` in an internal loop so an
N-deep chain folds in a single rule application instead of waiting for
the outer fixpoint optimizer to re-run the rule N times.

Bench (`physical_plan_chained_case_projection_hotspot`, N=80 chained
projections, depth-23 nested CASE per step, 30-wide OR filter on top of
a LEFT JOIN):
- before: 623 ms
- after : 364 ms  (-41.5%, criterion p<0.05, CI [-41.83%, -41.33%])

At the parser-depth limit (depth=45) the bench measures ~1.5 s/iter
post-fix — the same shape was multi-second to several-second territory
on the prior code.

Verification: 24/24 physical-plan projection unit tests pass; 54/54
optimizer projection tests; 23/23 `physical_optimizer::projection_pushdown`
integration tests.

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

Author: Dandandan

datafusion/core/benches/sql_planner_extended.rs

@@ -324,6 +324,61 @@ fn build_non_case_left_join_df_with_push_down_filter(
     rt.block_on(async { ctx.sql(&query).await.unwrap() })
 }
 
+/// Build a query of the shape:
+///   join + wide-OR filter + N chained CTEs, each adding one column
+///   defined by a depth-K nested CASE ladder over the same input column.
+///
+/// The chained projections force the physical planner to walk a stack of
+/// `ProjectionExec`s whose expression trees all reference the same upstream
+/// column, which exercises the `ProjectionPushdown` physical rule heavily.
+fn build_chained_case_projection_query(
+    chained_steps: usize,
+    case_depth: usize,
+    or_width: usize,
+) -> String {
+    let mut q = String::new();
+    q.push_str("WITH s0 AS (\n  SELECT l.c0, l.c1 FROM t l LEFT JOIN t r ON l.c0 = r.c0");
+    if or_width > 0 {
+        q.push_str("\n  WHERE (");
+        for i in 0..or_width {
+            if i > 0 {
+                q.push_str(" OR ");
+            }
+            let _ = write!(&mut q, "l.c1 = '{i}'");
+        }
+        q.push(')');
+    }
+    q.push_str("\n)");
+
+    for n in 1..=chained_steps {
+        q.push_str(",\n");
+        let _ = write!(&mut q, "s{n} AS (SELECT *, ");
+        for d in 0..case_depth {
+            let _ = write!(&mut q, "CASE WHEN c0 = '{d}' THEN 'label' ELSE ");
+        }
+        q.push_str("c0");
+        for _ in 0..case_depth {
+            q.push_str(" END");
+        }
+        let _ = write!(&mut q, " AS d{n} FROM s{prev})", prev = n - 1);
+    }
+
+    let _ = write!(&mut q, "\nSELECT * FROM s{chained_steps}");
+    q
+}
+
+fn build_chained_case_projection_df(
+    rt: &Runtime,
+    chained_steps: usize,
+    case_depth: usize,
+    or_width: usize,
+) -> DataFrame {
+    let ctx = SessionContext::new();
+    register_string_table(&ctx, 100, 1000);
+    let query = build_chained_case_projection_query(chained_steps, case_depth, or_width);
+    rt.block_on(async { ctx.sql(&query).await.unwrap() })
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     let baseline_ctx = SessionContext::new();
     let case_heavy_ctx = SessionContext::new();
@@ -460,6 +515,44 @@ fn criterion_benchmark(c: &mut Criterion) {
         }
     }
     control_group.finish();
+
+    // Hot-spot bench: N chained CTE projections, each adding one column
+    // defined by a depth-K nested CASE ladder, on top of a join + wide-OR
+    // filter. Exercises the physical `ProjectionPushdown` rule.
+    let chained_df = build_chained_case_projection_df(&rt, 80, 23, 30);
+    c.bench_function(
+        "physical_plan_chained_case_projection_hotspot",
+        |b| {
+            b.iter(|| {
+                let df_clone = chained_df.clone();
+                black_box(rt.block_on(async {
+                    df_clone.create_physical_plan().await.unwrap()
+                }));
+            })
+        },
+    );
+
+    let mut chained_group =
+        c.benchmark_group("physical_plan_chained_case_projection_sweep");
+    for &steps in &[10usize, 20, 40] {
+        for &case_depth in &[5usize, 10, 15] {
+            let df = build_chained_case_projection_df(&rt, steps, case_depth, 30);
+            let label = format!("steps={steps},depth={case_depth}");
+            chained_group.bench_with_input(
+                BenchmarkId::new("physical_plan", &label),
+                &df,
+                |b, df| {
+                    b.iter(|| {
+                        let df_clone = df.clone();
+                        black_box(rt.block_on(async {
+                            df_clone.create_physical_plan().await.unwrap()
+                        }));
+                    })
+                },
+            );
+        }
+    }
+    chained_group.finish();
 }
 
 criterion_group!(benches, criterion_benchmark);

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -536,6 +536,31 @@ fn optimize_subqueries(
 /// - `Ok(None)`: Signals that merge is not beneficial (and has not taken place).
 /// - `Err(error)`: An error occurred during the function call.
 fn merge_consecutive_projections(proj: Projection) -> Result<Transformed<Projection>> {
+    // Iteratively merge as long as the current Projection's input is another
+    // Projection that can be folded. Without this, a chain of N consecutive
+    // projections (e.g. SQL of the form `SELECT *, <expr> AS dN FROM ...`
+    // stacked N times) requires N applications of this rule (re-run by the
+    // outer optimizer loop, each walking the full plan tree) before the
+    // chain is fully collapsed.
+    let mut current = proj;
+    let mut transformed_any = false;
+    loop {
+        let result = merge_consecutive_projections_one_level(current)?;
+        current = result.data;
+        if !result.transformed {
+            return Ok(if transformed_any {
+                Transformed::yes(current)
+            } else {
+                Transformed::no(current)
+            });
+        }
+        transformed_any = true;
+    }
+}
+
+fn merge_consecutive_projections_one_level(
+    proj: Projection,
+) -> Result<Transformed<Projection>> {
     let Projection {
         expr,
         input,

datafusion/physical-plan/src/projection.rs

@@ -382,12 +382,15 @@ impl ExecutionPlan for ProjectionExec {
         &self,
         projection: &ProjectionExec,
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-        let maybe_unified = try_unifying_projections(projection, self)?;
-        if let Some(new_plan) = maybe_unified {
-            // To unify 3 or more sequential projections:
-            remove_unnecessary_projections(new_plan).data().map(Some)
-        } else {
-            Ok(Some(Arc::new(projection.clone())))
+        // Collapse the entire run of consecutive `ProjectionExec`s in a single
+        // pass instead of unifying one level at a time. The pairwise approach
+        // had to allocate (and recompute equivalence properties for) one
+        // intermediate `ProjectionExec` per level, which is the dominant cost
+        // for plans with many chained projections (e.g. a long pipeline of
+        // `SELECT *, <expr> AS dN FROM ...`).
+        match try_collapse_projection_chain(projection)? {
+            Some(plan) => Ok(Some(plan)),
+            None => Ok(Some(Arc::new(projection.clone()))),
         }
     }
 
@@ -1015,54 +1018,93 @@ pub fn update_join_filter(
 }
 
 /// Unifies `projection` with its input (which is also a [`ProjectionExec`]).
-fn try_unifying_projections(
-    projection: &ProjectionExec,
-    child: &ProjectionExec,
+/// Iteratively collapse a chain of consecutive [`ProjectionExec`]s starting
+/// at `outer` into a single `ProjectionExec`. Returns `None` if the very
+/// first level cannot be unified (matching the previous behavior of
+/// [`try_unifying_projections`] + recursive `remove_unnecessary_projections`).
+///
+/// Functionally equivalent to repeatedly applying [`try_unifying_projections`]
+/// and re-entering [`remove_unnecessary_projections`] on the result, but
+/// avoids constructing an intermediate `ProjectionExec` (and recomputing its
+/// [`PlanProperties`] / equivalence properties) for every level it walks
+/// through. For long projection chains the saved work scales linearly with
+/// the chain length.
+fn try_collapse_projection_chain(
+    outer: &ProjectionExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    let mut projected_exprs = vec![];
-    let mut column_ref_map: HashMap<Column, usize> = HashMap::new();
+    let mut current_exprs: Vec<ProjectionExpr> = outer.expr().to_vec();
+    let mut current_input: Arc<dyn ExecutionPlan> = Arc::clone(outer.input());
+    let mut collapsed_any = false;
 
-    // Collect the column references usage in the outer projection.
-    projection.expr().iter().for_each(|proj_expr| {
-        proj_expr
-            .expr
-            .apply(|expr| {
-                Ok({
-                    if let Some(column) = expr.downcast_ref::<Column>() {
-                        *column_ref_map.entry(column.clone()).or_default() += 1;
-                    }
-                    TreeNodeRecursion::Continue
-                })
-            })
-            .unwrap();
-    });
-    // Merging these projections is not beneficial, e.g
-    // If an expression is not trivial (KeepInPlace) and it is referred more than 1, unifies projections will be
-    // beneficial as caching mechanism for non-trivial computations.
-    // See discussion in: https://github.com/apache/datafusion/issues/8296
-    if column_ref_map.iter().any(|(column, count)| {
-        *count > 1
-            && !child.expr()[column.index()]
-                .expr
-                .placement()
-                .should_push_to_leaves()
-    }) {
-        return Ok(None);
-    }
-    for proj_expr in projection.expr() {
-        // If there is no match in the input projection, we cannot unify these
-        // projections. This case will arise if the projection expression contains
-        // a `PhysicalExpr` variant `update_expr` doesn't support.
-        let Some(expr) = update_expr(&proj_expr.expr, child.expr(), true)? else {
-            return Ok(None);
+    loop {
+        // Only continue if the current input is itself a projection.
+        let Some(inner_proj) = current_input.downcast_ref::<ProjectionExec>()
+        else {
+            break;
         };
-        projected_exprs.push(ProjectionExpr {
-            expr,
-            alias: proj_expr.alias.clone(),
+
+        // Replicate the "merging is beneficial" guard from
+        // `try_unifying_projections`: if any column referenced more than once
+        // in `current_exprs` resolves to a non-trivial expression in the
+        // inner projection, fusing them would duplicate that computation.
+        let mut column_ref_map: HashMap<Column, usize> = HashMap::new();
+        for proj_expr in &current_exprs {
+            proj_expr.expr.apply(|expr| {
+                if let Some(column) = expr.downcast_ref::<Column>() {
+                    *column_ref_map.entry(column.clone()).or_default() += 1;
+                }
+                Ok(TreeNodeRecursion::Continue)
+            })?;
+        }
+        let inner_exprs = inner_proj.expr();
+        let blocked = column_ref_map.iter().any(|(column, count)| {
+            *count > 1
+                && !inner_exprs[column.index()]
+                    .expr
+                    .placement()
+                    .should_push_to_leaves()
         });
+        if blocked {
+            break;
+        }
+
+        // Substitute each outer expression through the inner projection.
+        let mut new_exprs: Vec<ProjectionExpr> =
+            Vec::with_capacity(current_exprs.len());
+        let mut can_unify = true;
+        for proj_expr in &current_exprs {
+            match update_expr(&proj_expr.expr, inner_exprs, true)? {
+                Some(expr) => new_exprs.push(ProjectionExpr {
+                    expr,
+                    alias: proj_expr.alias.clone(),
+                }),
+                None => {
+                    can_unify = false;
+                    break;
+                }
+            }
+        }
+        if !can_unify {
+            break;
+        }
+
+        current_exprs = new_exprs;
+        current_input = Arc::clone(inner_proj.input());
+        collapsed_any = true;
     }
-    ProjectionExec::try_new(projected_exprs, Arc::clone(child.input()))
-        .map(|e| Some(Arc::new(e) as _))
+
+    if !collapsed_any {
+        return Ok(None);
+    }
+
+    // After collapsing the projection chain, hand the unified projection back
+    // to `remove_unnecessary_projections` once. This gives the non-Projection
+    // input (e.g. `DataSourceExec`) the chance to absorb the projection via
+    // its own `try_swapping_with_projection` impl — the same behaviour the
+    // pairwise recursion used to provide for the final step.
+    let unified: Arc<dyn ExecutionPlan> =
+        Arc::new(ProjectionExec::try_new(current_exprs, current_input)?);
+    remove_unnecessary_projections(unified).data().map(Some)
 }
 
 /// Collect all column indices from the given projection expressions.