diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index 80aceb8cad44c..cc9ffb3cfe3e6 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 @@ -904,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::{ @@ -1174,6 +1230,107 @@ 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. 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(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![ + 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:?}" + ); + + // `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. + 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()?; 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 {