perf: reuse projection schema in OptimizeProjections instead of recomputing it - #24281
Draft
zhuqi-lucas wants to merge 4 commits into
Draft
perf: reuse projection schema in OptimizeProjections instead of recomputing it#24281zhuqi-lucas wants to merge 4 commits into
OptimizeProjections instead of recomputing it#24281zhuqi-lucas wants to merge 4 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Improves OptimizeProjections performance by avoiding repeated recomputation of projection output schemas when pruning projection expressions, instead reusing/slicing the already-computed Projection.schema.
Changes:
- Update
rewrite_projection_given_requirementsto build pruned projections viaProjection::try_new_with_schemausing a sliced schema from the existing projection schema. - Add
project_schema_by_indiceshelper to project fields + functional dependencies while reusing schema metadata. - Add a unit test validating that sliced schemas match
projection_schemarecomputation across representative index subsets.
Suppressed comments (1)
datafusion/optimizer/src/optimize_projections/mod.rs:1285
project_schema_by_indicesalso projects functional dependencies and preserves schema-level metadata, but the test currently only compares fields and qualifiers. Adding assertions for functional dependencies and schema metadata will better protect the behavior this PR relies on.
// Output fields (name, data type, nullability, field metadata) must
// match the from-scratch computation exactly.
assert_eq!(
reused.fields(),
recomputed.fields(),
"fields differ for indices {indices:?}"
);
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+1249
to
+1252
| binary_expr(col("b"), Operator::Plus, col("c")), | ||
| col("c").alias("c_alias"), | ||
| lit(1_i64).alias("one"), | ||
| Expr::Column(Column::new(Some(TableReference::bare("test")), "b")), |
…puting it `rewrite_projection_given_requirements` rebuilt the pruned projection with `Projection::try_new`, which recomputes the output schema from scratch via `projection_schema`: it calls `Expr::to_field` for every retained expression, and column resolution (`DFSchema::field_from_column`) is a linear scan, so recomputing a projection's schema is O(exprs * schema_width) and runs on every projection on every optimizer pass. This is especially costly for wide `SELECT *`-style projections over wide schemas. The retained expressions are a subset of the projection's original expressions, so their output fields are unchanged by pruning unreferenced sibling columns. Select those fields from the existing projection schema and construct the pruned projection with `try_new_with_schema`, mirroring the schema reuse already done in `merge_consecutive_projections`. When nothing is pruned the schema Arc is reused as-is. This turns the per-projection schema cost from O(exprs * width) into O(k). Behavior-preserving: the sliced schema is identical to the recomputed one. Adds `project_schema_by_indices_matches_recompute` asserting that equivalence across expression subsets; the full datafusion-optimizer suite still passes.
Addresses review feedback on the schema-reuse test: - The comment claimed a nullable literal, but lit(1_i64) is non-nullable, so nullability propagation was never actually exercised. Swapped it for a NULL Int64 literal and added assertions pinning the premise that the literal is nullable while the input columns are not. - project_schema_by_indices also carries schema-level metadata and projects functional dependencies through the kept indices, but the test only compared fields and qualifiers. Both are now asserted against the from-scratch computation for every subset.
zhuqi-lucas
force-pushed
the
optimize-projections-reuse-schema
branch
from
August 12, 2026 08:56
3c29a28 to
dd18a14
Compare
zhuqi-lucas
marked this pull request as draft
August 12, 2026 09:49
This was referenced Aug 12, 2026
`LogicalPlan::map_expressions` replaces a projection's expressions while keeping its existing schema, so `SimplifyExpressions` could leave the two out of step: constant folding turns a function call, whose field the planner derived as nullable, into a non-null literal, whose field is not, and the schema keeps the pre-folding answer. That was invisible because `OptimizeProjections` rebuilds the projections it touches with `Projection::try_new`, deriving the schema again and normalising it back. Which meant whether a stale schema reached the final plan depended on which rules happened to fire, and it blocked deriving a pruned projection's schema by reuse rather than recomputation. Derive the schema here instead, only when the expressions actually changed. The final plans are unchanged, since the normalisation that `OptimizeProjections` was doing simply happens earlier now: no snapshot or expected plan in the tree needed updating.
zhuqi-lucas
marked this pull request as ready for review
August 13, 2026 07:12
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24281 +/- ##
=======================================
Coverage 81.13% 81.13%
=======================================
Files 1112 1112
Lines 386716 386802 +86
Branches 386716 386802 +86
=======================================
+ Hits 313765 313839 +74
- Misses 54479 54481 +2
- Partials 18472 18482 +10 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Author
|
run benchmark sql_planner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing optimize-projections-reuse-schema (75deafe) to 186f96f (merge-base) diff Run configurationrun benchmark sql_plannerResults will be posted here when complete File an issue against this benchmark runner |
zhuqi-lucas
marked this pull request as draft
August 13, 2026 08:51
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing optimize-projections-reuse-schema (75deafe) to 186f96f (merge-base) diff Run configurationrun benchmark sql_plannerCPU Details (lscpu)Details
Resource Usagesql_planner — base (merge-base)
sql_planner — branch
File an issue against this benchmark runner |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this close?
Closes #24264. Answers #24284 in the process.
Rationale for this change
rewrite_projection_given_requirements, the core ofOptimizeProjections, prunes a projection's expressions to the subset actually required and then rebuilds it withProjection::try_new.try_newrecomputes the output schema viaprojection_schema, callingExpr::to_fieldfor every retained expression; column resolution is a linear scan over the input schema, so this isO(exprs * schema_width)per projection, per pass. The retained expressions are a subset of the ones the projection already has, so the answer is already sitting inproj.schema.Reusing it turned out not to be safe as-is, which is what the first revision of this PR got wrong and what #24284 is about.
LogicalPlan::map_expressionsreplaces a projection's expressions while keeping its existing schema, soSimplifyExpressionscan leave the two out of step. Constant folding turnsarrow_cast([...], 'LargeList(...)'), a function call whose field the planner derived as nullable, into a non-null literal whose field is not, while the schema keeps the pre-folding answer.OptimizeProjectionscallingtry_newwas quietly normalising that back, so whether a stale schema survived into the final plan depended on which rules happened to fire.So this PR fixes that first, then does the optimisation.
What changes are included in this PR?
1.
SimplifyExpressionsderives the projection schema after rewriting (simplify_exprs.rs), and only when the expressions actually changed.Final plans are unchanged: the normalisation
OptimizeProjectionswas performing simply happens earlier now. Nothing in the tree needed updating, no snapshot and no expected plan, which is the clearest evidence this is an equivalence rather than a behaviour change.2.
rewrite_projection_given_requirementsderives the pruned schema by selecting the already-computed fields from the existing projection schema (project_schema_by_indices) and builds withProjection::try_new_with_schema. When nothing is pruned, the existingArcis reused as-is. Functional dependencies are projected through the kept indices.Cost goes from
O(exprs * width)toO(k), and toO(1)when nothing is pruned. Unlike making the recompute cheaper, this removes the work rather than speeding it up: noto_fieldcall, no field allocation, no cache and no heuristics.Correctness
With (1) in place,
proj.schemais in step withproj.expr, so slicing it at the retained indices produces exactly whatprojection_schemawould recompute: fieldicorresponds to expressioni, andRequiredIndicesyields a sorted, deduplicated subset.project_schema_by_indices_matches_recomputeasserts that, for a mixed expression list (plain column, computed binary expr, alias, NULL literal, qualified column) and every representative index subset, the sliced schema matchesprojection_schemaon fields, qualifiers, schema metadata and functional dependencies, and that the identity subset reuses the sameArc.The two failures the first revision of this PR introduced are fixed by (1), not worked around:
roundtrip_literal_list,roundtrip_literal_struct,roundtrip_literal_named_struct,roundtrip_literal_renamed_structindatafusion-substrait, which compare plan schemas across a roundtripschema_evolution_nested.slt, where the projection feedsCOPY (SELECT ...) TO ... STORED AS PARQUET, so a stale nullability reached the written file andDESCRIBEreportedYESinstead ofNOFull local runs:
datafusion-substrait49 + 200 + 3,datafusion-optimizer760 + 26 + 5,datafusion-expr248 + 55,datafusion-common547,datafusion-sql88 + 572 + 12, andschema_evolution_nested.slt1/1. All green, with no test or snapshot modified.Are there any user-facing changes?
No. Optimized plans are unchanged.