Add a checker to confirm physical optimizer rules will keep the physical plan schema immutable (#4316)

* Add a checker to confirm physical optimizer rules will keep the physical plan schema immutable

* fix debug print

* fix comments

Author: mingmwang

datafusion/core/src/physical_optimizer/aggregate_statistics.rs

@@ -18,7 +18,6 @@
 //! Utilizing exact statistics from sources to avoid scanning data
 use std::sync::Arc;
 
-use arrow::datatypes::Schema;
 use datafusion_expr::utils::COUNT_STAR_EXPANSION;
 
 use crate::execution::context::SessionConfig;
@@ -85,7 +84,7 @@ impl PhysicalOptimizerRule for AggregateStatistics {
                 // input can be entirely removed
                 Ok(Arc::new(ProjectionExec::try_new(
                     projections,
-                    Arc::new(EmptyExec::new(true, Arc::new(Schema::empty()))),
+                    Arc::new(EmptyExec::new(true, plan.schema())),
                 )?))
             } else {
                 optimize_children(self, plan, config)
@@ -98,6 +97,11 @@ impl PhysicalOptimizerRule for AggregateStatistics {
     fn name(&self) -> &str {
         "aggregate_statistics"
     }
+
+    /// This rule will change the nullable properties of the schema, disable the schema check.
+    fn schema_check(&self) -> bool {
+        false
+    }
 }
 
 /// assert if the node passed as argument is a final `AggregateExec` node that can be optimized:

datafusion/core/src/physical_optimizer/coalesce_batches.rs

@@ -72,4 +72,8 @@ impl PhysicalOptimizerRule for CoalesceBatches {
     fn name(&self) -> &str {
         "coalesce_batches"
     }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
 }

datafusion/core/src/physical_optimizer/enforcement.rs

@@ -101,6 +101,10 @@ impl PhysicalOptimizerRule for BasicEnforcement {
     fn name(&self) -> &str {
         "BasicEnforcement"
     }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
 }
 
 /// When the physical planner creates the Joins, the ordering of join keys is from the original query.

datafusion/core/src/physical_optimizer/hash_build_probe_order.rs

@@ -211,6 +211,10 @@ impl PhysicalOptimizerRule for HashBuildProbeOrder {
     fn name(&self) -> &str {
         "hash_build_probe_order"
     }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
 }
 
 #[cfg(test)]

datafusion/core/src/physical_optimizer/optimizer.rs

@@ -36,4 +36,10 @@ pub trait PhysicalOptimizerRule {
 
     /// A human readable name for this optimizer rule
     fn name(&self) -> &str;
+
+    /// A flag to indicate whether the physical planner should valid the rule will not
+    /// change the schema of the plan after the rewriting.
+    /// Some of the optimization rules might change the nullable properties of the schema
+    /// and should disable the schema check.
+    fn schema_check(&self) -> bool;
 }

datafusion/core/src/physical_optimizer/repartition.rs

@@ -233,6 +233,10 @@ impl PhysicalOptimizerRule for Repartition {
     fn name(&self) -> &str {
         "repartition"
     }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
 }
 #[cfg(test)]
 mod tests {

datafusion/core/src/physical_plan/planner.rs

@@ -1679,7 +1679,16 @@ impl DefaultPhysicalPlanner {
 
         let mut new_plan = plan;
         for optimizer in optimizers {
+            let before_schema = new_plan.schema();
             new_plan = optimizer.optimize(new_plan, &session_state.config)?;
+            if optimizer.schema_check() && new_plan.schema() != before_schema {
+                return Err(DataFusionError::Internal(format!(
+                        "PhysicalOptimizer rule '{}' failed, due to generate a different schema, original schema: {:?}, new schema: {:?}",
+                        optimizer.name(),
+                        before_schema,
+                        new_plan.schema()
+                    )));
+            }
             observer(new_plan.as_ref(), optimizer.as_ref())
         }
         debug!(