You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/21 04:40:23 UTC
(datafusion) 01/01: perf: collapse chained projections in a single optimizer pass; reduce memory usage / recursion (#22389)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22389-4055e4417e80a6ef726b1ff11b65949f313c8912
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0da89616e43299546ed8405cfa900d0f97cb2515
Author: Daniël Heres <da...@gmail.com>
AuthorDate: Thu May 21 06:39:59 2026 +0200
perf: collapse chained projections in a single optimizer pass; reduce memory usage / recursion (#22389)
## Which issue does this PR close?
- Closes #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 #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) <no...@anthropic.com>
---
datafusion/core/benches/sql_planner_extended.rs | 61 +++++++++++
.../optimizer/src/optimize_projections/mod.rs | 24 +++++
.../optimizer/tests/optimizer_integration.rs | 7 +-
datafusion/physical-expr/src/projection.rs | 13 ++-
datafusion/physical-plan/src/projection.rs | 111 +++++++++++----------
5 files changed, 160 insertions(+), 56 deletions(-)
diff --git a/datafusion/core/benches/sql_planner_extended.rs b/datafusion/core/benches/sql_planner_extended.rs
index d4955313c7..b016d758f3 100644
--- a/datafusion/core/benches/sql_planner_extended.rs
+++ b/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);
diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs
index bc923706a4..59109a822b 100644
--- a/datafusion/optimizer/src/optimize_projections/mod.rs
+++ b/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,
diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs
index e61e646793..4e33bf6b3a 100644
--- a/datafusion/optimizer/tests/optimizer_integration.rs
+++ b/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"))]
"#,
);
diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs
index ca999479fa..8320983c10 100644
--- a/datafusion/physical-expr/src/projection.rs
+++ b/datafusion/physical-expr/src/projection.rs
@@ -963,8 +963,6 @@ 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 {}",
@@ -972,6 +970,17 @@ pub fn update_expr(
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
diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs
index e5b91fbb1c..951ed618e5 100644
--- a/datafusion/physical-plan/src/projection.rs
+++ b/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 ¤t_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 ¤t_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.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org