perf: collapse chained projections in a single optimizer pass; reduce memory usage / recursion (apache#22389)

## Which issue does this PR close?

- Closes apache#22391.

## Rationale for this change

Long chains of consecutive projections (e.g. `SELECT *, <expr> AS dN
FROM ...` stacked N times, with depth-K nested CASE per step) caused
planning to use multi-GB of RAM and tens of seconds of wall time,
scaling superlinearly in chain length and CASE depth. See apache#22391 for the
OOM/timing data.

## What changes are included in this PR?

Three independent fixes targeting the same workload:

1. **Physical chain collapse**
(`datafusion/physical-plan/src/projection.rs`): replace the pairwise
recursive unification in `ProjectionExec::try_swapping_with_projection`
with `try_collapse_projection_chain`, which walks the entire run of
consecutive `ProjectionExec`s and builds **one** final `ProjectionExec`
(saves N-1 intermediate constructions and their `compute_properties`
calls). Leaf pushdown into a non-`Projection` input is preserved by
calling `remove_unnecessary_projections` once at the end.

2. **Logical iterative merge**
(`datafusion/optimizer/src/optimize_projections/mod.rs`): wrap
`merge_consecutive_projections` in an internal loop so an N-deep
`LogicalPlan::Projection` chain collapses in a single rule application
instead of N outer fixpoint passes.

3. **`update_expr` Column-equality short-circuit**
(`datafusion/physical-expr/src/projection.rs`): when substituting a
Column with one that equals it (the pass-through case during chain
collapse), return `Transformed::no` so `transform_up` does not rebuild
the enclosing `CaseExpr` / `CaseBody`. This is the OOM fix — it
eliminates the O(N² · K²) cascade of CASE allocations.

## Are these changes tested?

- yes, existing test + added bench

  | Stage | Time | Δ vs master |
  |---|---:|---:|
  | master | 623 ms | — |
  | + chain collapse + logical merge loop | 364 ms | −41.5% |
| + `update_expr` Column-equality short-circuit | **155 ms** |
**−75.4%** |

  (criterion p<0.05, CI [−75.67%, −75.16%])

## Are there any user-facing changes?

No public API change.

---------

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

Author: Dandandan

datafusion/core/benches/sql_planner_extended.rs

@@ -324,6 +324,57 @@ fn build_non_case_left_join_df_with_push_down_filter(
     rt.block_on(async { ctx.sql(&query).await.unwrap() })
 }
 
+/// Join + wide-OR filter + N chained CTEs, each adding one column defined
+/// by a depth-K nested CASE ladder over the same input column. Exercises
+/// the physical `ProjectionPushdown` rule on long projection chains.
+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 +511,16 @@ fn criterion_benchmark(c: &mut Criterion) {
         }
     }
     control_group.finish();
+
+    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() }),
+            );
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -536,6 +536,30 @@ 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>> {
+    // Collapse the whole chain in one pass; otherwise an N-deep chain needs
+    // N outer optimizer passes to fully fold.
+    let mut current = proj;
+    let mut transformed_any = false;
+    loop {
+        let Transformed {
+            data, transformed, ..
+        } = merge_consecutive_projections_one_level(current)?;
+        current = data;
+        if !transformed {
+            break;
+        }
+        transformed_any = true;
+    }
+    Ok(if transformed_any {
+        Transformed::yes(current)
+    } else {
+        Transformed::no(current)
+    })
+}
+
+fn merge_consecutive_projections_one_level(
+    proj: Projection,
+) -> Result<Transformed<Projection>> {
     let Projection {
         expr,
         input,

datafusion/optimizer/tests/optimizer_integration.rs

@@ -824,10 +824,9 @@ fn extension_node_does_not_block_projection_pruning() -> Result<()> {
     OpaqueRequirementsExtension
       Sort: t.a ASC NULLS FIRST, t.ts ASC NULLS FIRST
         Projection: t.a, CAST(t.ts AS Timestamp(ms, "UTC")) AS ts
-          Projection: t.a, t.ts
-            Filter: __common_expr_3 > TimestampMillisecond(1000, Some("UTC")) AND __common_expr_3 < TimestampMillisecond(2000, Some("UTC"))
-              Projection: CAST(t.ts AS Timestamp(ms, "UTC")) AS __common_expr_3, t.a, t.ts
-                TableScan: t projection=[a, ts], partial_filters=[t.ts > TimestampNanosecond(1000000000, None), t.ts < TimestampNanosecond(2000000000, None), CAST(t.ts AS Timestamp(ms, "UTC")) > TimestampMillisecond(1000, Some("UTC")), CAST(t.ts AS Timestamp(ms, "UTC")) < TimestampMillisecond(2000, Some("UTC"))]
+          Filter: __common_expr_3 > TimestampMillisecond(1000, Some("UTC")) AND __common_expr_3 < TimestampMillisecond(2000, Some("UTC"))
+            Projection: CAST(t.ts AS Timestamp(ms, "UTC")) AS __common_expr_3, t.a, t.ts
+              TableScan: t projection=[a, ts], partial_filters=[t.ts > TimestampNanosecond(1000000000, None), t.ts < TimestampNanosecond(2000000000, None), CAST(t.ts AS Timestamp(ms, "UTC")) > TimestampMillisecond(1000, Some("UTC")), CAST(t.ts AS Timestamp(ms, "UTC")) < TimestampMillisecond(2000, Some("UTC"))]
     "#,
     );
 

datafusion/physical-expr/src/projection.rs

@@ -963,15 +963,24 @@ pub fn update_expr(
                 return Ok(Transformed::no(expr));
             };
             if unproject {
-                state = RewriteState::RewrittenValid;
-                // Update the index of `column`:
                 let projected_expr = projected_exprs.get(column.index()).ok_or_else(|| {
                     internal_datafusion_err!(
                         "Column index {} out of bounds for projected expressions of length {}",
                         column.index(),
                         projected_exprs.len()
                     )
                 })?;
+                // Skip rebuilding the parent if substituting with an equal
+                // Column (e.g. pass-through `c0@0` -> `c0@0` during chained
+                // projection collapse). Without this, every CASE/BinaryExpr
+                // containing such a Column is reconstructed unnecessarily.
+                if let Some(projected_col) =
+                    projected_expr.expr.downcast_ref::<Column>()
+                    && projected_col == column
+                {
+                    return Ok(Transformed::no(expr));
+                }
+                state = RewriteState::RewrittenValid;
                 Ok(Transformed::yes(Arc::clone(&projected_expr.expr)))
             } else {
                 // default to invalid, in case we can't find the relevant column

datafusion/physical-plan/src/projection.rs

@@ -382,12 +382,9 @@ 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())))
+        match try_collapse_projection_chain(projection)? {
+            Some(plan) => Ok(Some(plan)),
+            None => Ok(Some(Arc::new(projection.clone()))),
         }
     }
 
@@ -1014,55 +1011,69 @@ pub fn update_join_filter(
     })
 }
 
-/// Unifies `projection` with its input (which is also a [`ProjectionExec`]).
-fn try_unifying_projections(
-    projection: &ProjectionExec,
-    child: &ProjectionExec,
+/// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns
+/// `None` if nothing could be merged.
+fn try_collapse_projection_chain(
+    outer: &ProjectionExec,
 ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
-    let mut projected_exprs = vec![];
+    let mut current_exprs: Vec<ProjectionExpr> = outer.expr().to_vec();
+    let mut current_input: Arc<dyn ExecutionPlan> = Arc::clone(outer.input());
     let mut column_ref_map: HashMap<Column, usize> = HashMap::new();
+    let mut collapsed_any = false;
+
+    'outer: while let Some(inner_proj) = current_input.downcast_ref::<ProjectionExec>() {
+        // Collect the column references usage in the outer projection.
+        column_ref_map.clear();
+        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();
+        // 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
+        let blocked = column_ref_map.iter().any(|(column, count)| {
+            *count > 1
+                && !inner_exprs[column.index()]
+                    .expr
+                    .placement()
+                    .should_push_to_leaves()
+        });
+        if blocked {
+            break;
+        }
 
-    // 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);
+        let mut new_phys: Vec<Arc<dyn PhysicalExpr>> =
+            Vec::with_capacity(current_exprs.len());
+        for proj_expr in &current_exprs {
+            // 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, inner_exprs, true)? else {
+                break 'outer;
+            };
+            new_phys.push(expr);
+        }
+        for (proj_expr, expr) in current_exprs.iter_mut().zip(new_phys) {
+            proj_expr.expr = expr;
+        }
+        current_input = Arc::clone(inner_proj.input());
+        collapsed_any = true;
     }
-    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);
-        };
-        projected_exprs.push(ProjectionExpr {
-            expr,
-            alias: proj_expr.alias.clone(),
-        });
+
+    if !collapsed_any {
+        return Ok(None);
     }
-    ProjectionExec::try_new(projected_exprs, Arc::clone(child.input()))
-        .map(|e| Some(Arc::new(e) as _))
+
+    // To unify 3 or more sequential projections:
+    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.