Add schema-aware optimizer child rewrites

datafusion/optimizer/src/eliminate_cross_join.rs

@@ -16,11 +16,12 @@
 // under the License.
 
 //! [`EliminateCrossJoin`] converts `CROSS JOIN` to `INNER JOIN` if join predicates are available.
+use crate::optimizer::map_children_recompute_schema_if_needed;
 use crate::{OptimizerConfig, OptimizerRule};
 use std::sync::Arc;
 
 use crate::join_key_set::JoinKeySet;
-use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
+use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};
 use datafusion_common::{NullEquality, Result};
 use datafusion_expr::expr::{BinaryExpr, Expr};
 use datafusion_expr::logical_plan::{
@@ -255,15 +256,12 @@ fn rewrite_children(
     let transformed_plan = plan
         .map_uncorrelated_subqueries(|input| optimizer.rewrite(input, config))?
         .transform_sibling(|plan| {
-            plan.map_children(|input| optimizer.rewrite(input, config))
+            map_children_recompute_schema_if_needed(plan, |input| {
+                optimizer.rewrite(input, config)
+            })
         })?;
 
-    // recompute schema if the plan was transformed
-    if transformed_plan.transformed {
-        transformed_plan.map_data(|plan| plan.recompute_schema())
-    } else {
-        Ok(transformed_plan)
-    }
+    Ok(transformed_plan)
 }
 
 /// Recursively accumulate possible_join_keys and inputs from inner joins

datafusion/optimizer/src/optimizer.rs

@@ -483,6 +483,42 @@ fn map_children_mut<F: FnMut(&mut LogicalPlan) -> Result<bool>>(
     })
 }
 
+/// Rewrites the direct children of `plan` and refreshes its schema when a
+/// rewritten child has a different schema.
+///
+/// Logical plan nodes cache schemas derived from their children. This helper
+/// keeps that cache in sync for owned-plan rewrites without recomputing the
+/// schema when a child was transformed but its schema stayed the same.
+pub(crate) fn map_children_recompute_schema_if_needed<F>(
+    plan: LogicalPlan,
+    f: F,
+) -> Result<Transformed<LogicalPlan>>
+where
+    F: FnMut(LogicalPlan) -> Result<Transformed<LogicalPlan>>,
+{
+    let child_schemas = plan
+        .inputs()
+        .into_iter()
+        .map(|child| Arc::clone(child.schema()))
+        .collect::<Vec<_>>();
+
+    let transformed_plan = plan.map_children(f)?;
+    if !transformed_plan.transformed {
+        return Ok(transformed_plan);
+    }
+
+    let schema_changed = child_schemas
+        .iter()
+        .zip(transformed_plan.data.inputs())
+        .any(|(old_schema, child)| old_schema.as_ref() != child.schema().as_ref());
+
+    if schema_changed {
+        transformed_plan.map_data(LogicalPlan::recompute_schema)
+    } else {
+        Ok(transformed_plan)
+    }
+}
+
 /// Rewrites a plan tree in place using `Arc::make_mut` for
 /// copy-on-write semantics on `Arc<LogicalPlan>` children.
 ///
@@ -905,6 +941,46 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn owned_rewrite_recomputes_parent_schema_when_child_schema_changes() -> Result<()> {
+        let plan = LogicalPlanBuilder::from(test_table_scan()?)
+            .filter(lit(true))?
+            .build()?;
+
+        let transformed =
+            super::map_children_recompute_schema_if_needed(plan, |child| {
+                let child = LogicalPlanBuilder::from(child)
+                    .project([col("a")])?
+                    .build()?;
+                Ok(Transformed::yes(child))
+            })?;
+
+        assert!(transformed.transformed);
+        assert_eq!(transformed.data.schema().fields().len(), 1);
+        Ok(())
+    }
+
+    #[test]
+    fn owned_rewrite_preserves_parent_schema_when_child_schema_is_unchanged() -> Result<()>
+    {
+        let plan = LogicalPlanBuilder::from(test_table_scan()?)
+            .filter(lit(true))?
+            .build()?;
+        let original_schema = Arc::clone(plan.schema());
+
+        let transformed =
+            super::map_children_recompute_schema_if_needed(plan, |child| {
+                let child = LogicalPlanBuilder::from(child)
+                    .project([col("a"), col("b"), col("c")])?
+                    .build()?;
+                Ok(Transformed::yes(child))
+            })?;
+
+        assert!(transformed.transformed);
+        assert_eq!(transformed.data.schema().as_ref(), original_schema.as_ref());
+        Ok(())
+    }
+
     #[test]
     fn optimizer_detects_plan_equal_to_the_initial() -> Result<()> {
         // Run a goofy optimizer, which rotates projection columns