diff --git a/datafusion/optimizer/src/eliminate_cross_join.rs b/datafusion/optimizer/src/eliminate_cross_join.rs index 95b70da443d8..02c96bde1f96 100644 --- a/datafusion/optimizer/src/eliminate_cross_join.rs +++ b/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 diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index db7ad8475273..4af9e1c6e55e 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -483,6 +483,42 @@ fn map_children_mut Result>( }) } +/// 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( + plan: LogicalPlan, + f: F, +) -> Result> +where + F: FnMut(LogicalPlan) -> Result>, +{ + let child_schemas = plan + .inputs() + .into_iter() + .map(|child| Arc::clone(child.schema())) + .collect::>(); + + 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` 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