From 5cd969b51018d006dab81891b3d1c294105bbaef Mon Sep 17 00:00:00 2001
From: Qi Zhu <zhuqilucas@gmail.com>
Date: Thu, 13 Aug 2026 15:09:52 +0800
Subject: [PATCH] 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<LogicalPlan> {
+    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 {