From 573306a3b7fd6953058a88a9062e6e436e40a429 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Mon, 10 Aug 2026 23:25:06 +0800 Subject: [PATCH 1/3] perf: reuse projection schema in OptimizeProjections instead of recomputing 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. --- .../optimizer/src/optimize_projections/mod.rs | 140 +++++++++++++++++- 1 file changed, 135 insertions(+), 5 deletions(-) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 80aceb8cad44c..c894d0874cc72 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -24,7 +24,7 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::{ - Column, DFSchema, HashMap, JoinType, Result, assert_eq_or_internal_err, + Column, DFSchema, DFSchemaRef, HashMap, JoinType, Result, assert_eq_or_internal_err, get_required_group_by_exprs_indices, internal_datafusion_err, internal_err, }; use datafusion_expr::expr::Alias; @@ -842,10 +842,24 @@ fn rewrite_projection_given_requirements( config: &dyn OptimizerConfig, indices: &RequiredIndices, ) -> Result> { - let Projection { expr, input, .. } = proj; + let Projection { + expr, + input, + schema, + .. + } = proj; let exprs_used = indices.get_at_indices(&expr); + // The retained expressions are a subset of the original projection's + // expressions, so their output fields are unchanged by pruning unreferenced + // sibling columns. Derive the pruned output schema by selecting those fields + // from the existing projection schema instead of recomputing every field's + // type/nullability via `Expr::to_field` (see `projection_schema`), which is + // quadratic in the schema width. This mirrors the schema reuse already + // performed in `merge_consecutive_projections`. + let projected_schema = project_schema_by_indices(&schema, indices.indices())?; + let required_indices = RequiredIndices::new().with_exprs(input.schema(), exprs_used.iter()); @@ -856,13 +870,55 @@ fn rewrite_projection_given_requirements( if is_projection_unnecessary(&input, &exprs_used)? { Ok(Transformed::yes(input)) } else { - Projection::try_new(exprs_used, Arc::new(input)) - .map(LogicalPlan::Projection) - .map(Transformed::yes) + Projection::try_new_with_schema( + exprs_used, + Arc::new(input), + Arc::clone(&projected_schema), + ) + .map(LogicalPlan::Projection) + .map(Transformed::yes) } }) } +/// Builds the output schema of a projection that keeps only the fields at +/// `indices` (a sorted, deduplicated subset produced by [`RequiredIndices`]) of +/// `schema`, reusing the already-computed fields instead of recomputing each +/// field's type/nullability from the expressions. +/// +/// Pruning unreferenced sibling columns cannot change the retained fields, so +/// the sliced schema is identical to the one [`projection_schema`] would +/// recompute, at O(k) instead of O(exprs * schema_width). +/// +/// [`projection_schema`]: datafusion_expr::logical_plan::projection_schema +fn project_schema_by_indices( + schema: &DFSchemaRef, + indices: &[usize], +) -> Result { + // Nothing pruned: the output schema is unchanged, reuse it as-is. + if indices.len() == schema.fields().len() { + return Ok(Arc::clone(schema)); + } + + let qualified_fields = indices + .iter() + .map(|&i| { + let (qualifier, field) = schema.qualified_field(i); + (qualifier.cloned(), Arc::clone(field)) + }) + .collect::>(); + + let func_deps = schema + .functional_dependencies() + .project_functional_dependencies(indices, indices.len()); + + let projected = + DFSchema::new_with_metadata(qualified_fields, schema.metadata().clone())? + .with_functional_dependencies(func_deps)?; + + Ok(Arc::new(projected)) +} + /// Projection is unnecessary, when /// - input schema of the projection, output schema of the projection are same, and /// - all projection expressions are either Column or Literal @@ -1174,6 +1230,80 @@ mod tests { } } + /// `project_schema_by_indices` must produce exactly the schema that + /// recomputing it from scratch (`projection_schema`) would, for every subset + /// of a projection's expressions. This is the correctness invariant that lets + /// `rewrite_projection_given_requirements` reuse the parent schema instead of + /// re-deriving each field's type via `Expr::to_field`. + #[test] + fn project_schema_by_indices_matches_recompute() -> Result<()> { + use super::project_schema_by_indices; + use datafusion_expr::logical_plan::projection_schema; + + let input = Arc::new(test_table_scan()?); // columns: a, b, c (UInt32, NOT NULL) + + // A deliberately mixed expression list: plain column, computed binary + // expr, alias, nullable literal, and a qualified column. + let exprs = vec![ + col("a"), + 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")), + ]; + + let full_schema = + Arc::clone(&Projection::try_new(exprs.clone(), Arc::clone(&input))?.schema); + + // Every sorted, deduplicated subset RequiredIndices could produce, + // including the "nothing pruned" identity case. + let subsets: Vec> = vec![ + vec![0, 1, 2, 3, 4], // identity (SELECT * fast path) + vec![0], + vec![1], // computed expr only + vec![2], // alias only + vec![3], // literal only + vec![0, 2], + vec![1, 3], + vec![0, 1, 4], + vec![2, 3, 4], + ]; + + for indices in subsets { + let exprs_used: Vec = + indices.iter().map(|&i| exprs[i].clone()).collect(); + + let reused = project_schema_by_indices(&full_schema, &indices)?; + let recomputed = projection_schema(input.as_ref(), &exprs_used)?; + + // 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:?}" + ); + + // Column qualifiers (table references) must be preserved too. + let reused_quals: Vec<_> = reused.iter().map(|(q, _)| q.cloned()).collect(); + let recomputed_quals: Vec<_> = + recomputed.iter().map(|(q, _)| q.cloned()).collect(); + assert_eq!( + reused_quals, recomputed_quals, + "qualifiers differ for indices {indices:?}" + ); + } + + // The identity subset must reuse the very same Arc, not rebuild it. + let identity = project_schema_by_indices(&full_schema, &[0, 1, 2, 3, 4])?; + assert!( + Arc::ptr_eq(&identity, &full_schema), + "identity projection should reuse the existing schema Arc" + ); + + Ok(()) + } + #[test] fn merge_two_projection() -> Result<()> { let table_scan = test_table_scan()?; From dd18a14017b85c7fd7a3a9e9567493b048fae718 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Wed, 12 Aug 2026 16:24:06 +0800 Subject: [PATCH 2/3] test: cover nullability, metadata and functional dependencies 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. --- .../optimizer/src/optimize_projections/mod.rs | 33 +++++++++++++++++-- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index c894d0874cc72..cc9ffb3cfe3e6 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -960,7 +960,7 @@ mod tests { use crate::{OptimizerContext, OptimizerRule}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::{ - Column, DFSchema, DFSchemaRef, JoinType, Result, TableReference, + Column, DFSchema, DFSchemaRef, JoinType, Result, ScalarValue, TableReference, }; use datafusion_expr::ExprFunctionExt; use datafusion_expr::{ @@ -1243,18 +1243,31 @@ mod tests { let input = Arc::new(test_table_scan()?); // columns: a, b, c (UInt32, NOT NULL) // A deliberately mixed expression list: plain column, computed binary - // expr, alias, nullable literal, and a qualified column. + // expr, alias, nullable literal, and a qualified column. Every input + // column is NOT NULL, so the NULL literal is what makes nullability + // actually vary across the schema. let exprs = vec![ col("a"), binary_expr(col("b"), Operator::Plus, col("c")), col("c").alias("c_alias"), - lit(1_i64).alias("one"), + lit(ScalarValue::Int64(None)).alias("null_one"), Expr::Column(Column::new(Some(TableReference::bare("test")), "b")), ]; let full_schema = Arc::clone(&Projection::try_new(exprs.clone(), Arc::clone(&input))?.schema); + // Guard the premise above: if this ever stops holding, the subsets below + // would no longer exercise nullability propagation at all. + assert!( + full_schema.field(3).is_nullable(), + "the literal must be nullable for this test to cover nullability" + ); + assert!( + !full_schema.field(0).is_nullable(), + "input columns are expected to be NOT NULL" + ); + // Every sorted, deduplicated subset RequiredIndices could produce, // including the "nothing pruned" identity case. let subsets: Vec> = vec![ @@ -1292,6 +1305,20 @@ mod tests { reused_quals, recomputed_quals, "qualifiers differ for indices {indices:?}" ); + + // `project_schema_by_indices` also carries over schema-level + // metadata and projects functional dependencies through the kept + // indices, so both must match the from-scratch computation as well. + assert_eq!( + reused.metadata(), + recomputed.metadata(), + "schema metadata differs for indices {indices:?}" + ); + assert_eq!( + reused.functional_dependencies(), + recomputed.functional_dependencies(), + "functional dependencies differ for indices {indices:?}" + ); } // The identity subset must reuse the very same Arc, not rebuild it. From 8d4ece93d60c100e1fb24ef93e3d7db19a432353 Mon Sep 17 00:00:00 2001 From: Qi Zhu Date: Thu, 13 Aug 2026 15:09:52 +0800 Subject: [PATCH 3/3] fix: keep a projection's schema in step with its simplified expressions `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. --- .../simplify_expressions/simplify_exprs.rs | 46 +++++++++++++++---- 1 file changed, 37 insertions(+), 9 deletions(-) diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 0e72a17abc9f7..6e9190d8f59c7 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -146,18 +146,46 @@ impl SimplifyExpressions { )) }; - plan.map_expressions(|expr| { - // Preserve the aliasing of grouping sets. - if let Expr::GroupingSet(_) = &expr { - expr.map_children(&mut rewrite_expr) - } else { - rewrite_expr(expr) - } - })? - .transform_data(rewrite_aggregate_non_aggregate_aggr_expr) + let rewritten = plan + .map_expressions(|expr| { + // Preserve the aliasing of grouping sets. + if let Expr::GroupingSet(_) = &expr { + expr.map_children(&mut rewrite_expr) + } else { + rewrite_expr(expr) + } + })? + .transform_data(rewrite_aggregate_non_aggregate_aggr_expr)?; + + if !rewritten.transformed { + return Ok(rewritten); + } + + rewritten.map_data(refresh_projection_schema) } } +/// Recomputes a `Projection`'s output schema after its expressions were +/// simplified. +/// +/// `LogicalPlan::map_expressions` replaces a projection's expressions while +/// keeping its existing schema, so simplification can 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. The +/// schema keeps the pre-folding answer. +/// +/// Downstream rules that rebuild the projection with `Projection::try_new` +/// used to paper over that by deriving the schema again, which made the +/// discrepancy invisible but also made it depend on which rules happen to +/// fire. Deriving it here instead keeps a projection's schema in step with +/// its own expressions. +fn refresh_projection_schema(plan: LogicalPlan) -> Result { + let LogicalPlan::Projection(Projection { expr, input, .. }) = plan else { + return Ok(plan); + }; + Projection::try_new(expr, input).map(LogicalPlan::Projection) +} + impl SimplifyExpressions { #[expect(missing_docs)] pub fn new() -> Self {