Interface for physical plan invariant checking.

datafusion/core/src/physical_planner.rs

@@ -66,6 +66,7 @@ use arrow::datatypes::{Schema, SchemaRef};
 use arrow_array::builder::StringBuilder;
 use arrow_array::RecordBatch;
 use datafusion_common::display::ToStringifiedPlan;
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion_common::{
     exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema,
     ScalarValue,
@@ -1880,18 +1881,12 @@ impl DefaultPhysicalPlanner {
                 .map_err(|e| {
                     DataFusionError::Context(optimizer.name().to_string(), Box::new(e))
                 })?;
-            if optimizer.schema_check() && new_plan.schema() != before_schema {
-                let e = DataFusionError::Internal(format!(
-                    "PhysicalOptimizer rule '{}' failed, due to generate a different schema, original schema: {:?}, new schema: {:?}",
-                    optimizer.name(),
-                    before_schema,
-                    new_plan.schema()
-                ));
-                return Err(DataFusionError::Context(
-                    optimizer.name().to_string(),
-                    Box::new(e),
-                ));
-            }
+
+            // confirm optimizer change did not violate invariants
+            InvariantChecker
+                .check(&new_plan, optimizer, before_schema)
+                .map_err(|e| e.context(optimizer.name().to_string()))?;
+
             trace!(
                 "Optimized physical plan by {}:\n{}\n",
                 optimizer.name(),
@@ -2006,6 +2001,45 @@ fn tuple_err<T, R>(value: (Result<T>, Result<R>)) -> Result<(T, R)> {
     }
 }
 
+#[derive(Default)]
+struct InvariantChecker;
+
+impl InvariantChecker {
+    /// Checks that the plan change is permitted, returning an Error if not.
+    ///
+    /// In debug mode, this recursively walks the entire physical plan and
+    /// performs additional checks using [`ExecutionPlan::check_node_invariants`].
+    pub fn check(
+        &mut self,
+        plan: &Arc<dyn ExecutionPlan>,
+        rule: &Arc<dyn PhysicalOptimizerRule + Send + Sync>,
+        previous_schema: Arc<Schema>,
+    ) -> Result<()> {
+        // Invariant: in most cases, the schema cannot be changed
+        // since the plan's output cannot change after the optimizer pass.
+        if rule.schema_check() && plan.schema() != previous_schema {
+            internal_err!("PhysicalOptimizer rule '{}' failed, due to generate a different schema, original schema: {:?}, new schema: {:?}",
+                rule.name(),
+                previous_schema,
+                plan.schema()
+            )?
+        }
+
+        #[cfg(debug_assertions)]
+        plan.visit(self)?;
+        Ok(())
+    }
+}
+
+impl<'n> TreeNodeVisitor<'n> for InvariantChecker {
+    type Node = Arc<dyn ExecutionPlan>;
+
+    fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
+        node.check_node_invariants()?;
+        Ok(TreeNodeRecursion::Continue)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::any::Any;
@@ -2026,6 +2060,7 @@ mod tests {
     use crate::execution::session_state::SessionStateBuilder;
     use arrow::array::{ArrayRef, DictionaryArray, Int32Array};
     use arrow::datatypes::{DataType, Field, Int32Type};
+    use datafusion_common::config::ConfigOptions;
     use datafusion_common::{assert_contains, DFSchemaRef, TableReference};
     use datafusion_execution::runtime_env::RuntimeEnv;
     use datafusion_execution::TaskContext;
@@ -2780,4 +2815,154 @@ digraph {
 
         assert_contains!(generated_graph, expected_tooltip);
     }
+
+    /// Extension Node which passes invariant checks
+    #[derive(Debug)]
+    struct OkExtensionNode(Vec<Arc<dyn ExecutionPlan>>);
+    impl ExecutionPlan for OkExtensionNode {
+        fn name(&self) -> &str {
+            "always ok"
+        }
+        fn with_new_children(
+            self: Arc<Self>,
+            children: Vec<Arc<dyn ExecutionPlan>>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(Arc::new(Self(children)))
+        }
+        fn schema(&self) -> SchemaRef {
+            Arc::new(Schema::empty())
+        }
+        fn as_any(&self) -> &dyn Any {
+            unimplemented!()
+        }
+        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+            self.0.iter().collect::<Vec<_>>()
+        }
+        fn properties(&self) -> &PlanProperties {
+            unimplemented!()
+        }
+        fn execute(
+            &self,
+            _partition: usize,
+            _context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            unimplemented!()
+        }
+    }
+    impl DisplayAs for OkExtensionNode {
+        fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(f, "{}", self.name())
+        }
+    }
+
+    /// Extension Node which fails invariant checks
+    #[derive(Debug)]
+    struct InvariantFailsExtensionNode;
+    impl ExecutionPlan for InvariantFailsExtensionNode {
+        fn name(&self) -> &str {
+            "InvariantFailsExtensionNode"
+        }
+        fn check_node_invariants(&self) -> Result<()> {
+            plan_err!("extension node failed it's user-defined invariant check")
+        }
+        fn schema(&self) -> SchemaRef {
+            Arc::new(Schema::empty())
+        }
+        fn with_new_children(
+            self: Arc<Self>,
+            _children: Vec<Arc<dyn ExecutionPlan>>,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn as_any(&self) -> &dyn Any {
+            unimplemented!()
+        }
+        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+            unimplemented!()
+        }
+        fn properties(&self) -> &PlanProperties {
+            unimplemented!()
+        }
+        fn execute(
+            &self,
+            _partition: usize,
+            _context: Arc<TaskContext>,
+        ) -> Result<SendableRecordBatchStream> {
+            unimplemented!()
+        }
+    }
+    impl DisplayAs for InvariantFailsExtensionNode {
+        fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(f, "{}", self.name())
+        }
+    }
+
+    /// Extension Optimizer rule that requires the schema check
+    #[derive(Debug)]
+    struct OptimizerRuleWithSchemaCheck;
+    impl PhysicalOptimizerRule for OptimizerRuleWithSchemaCheck {
+        fn optimize(
+            &self,
+            plan: Arc<dyn ExecutionPlan>,
+            _config: &ConfigOptions,
+        ) -> Result<Arc<dyn ExecutionPlan>> {
+            Ok(plan)
+        }
+        fn name(&self) -> &str {
+            "OptimizerRuleWithSchemaCheck"
+        }
+        fn schema_check(&self) -> bool {
+            true
+        }
+    }
+
+    #[test]
+    fn test_invariant_checker() -> Result<()> {
+        let rule: Arc<dyn PhysicalOptimizerRule + Send + Sync> =
+            Arc::new(OptimizerRuleWithSchemaCheck);
+
+        // ok plan
+        let ok_node: Arc<dyn ExecutionPlan> = Arc::new(OkExtensionNode(vec![]));
+        let child = Arc::clone(&ok_node);
+        let ok_plan = Arc::clone(&ok_node).with_new_children(vec![
+            Arc::clone(&child).with_new_children(vec![Arc::clone(&child)])?,
+            Arc::clone(&child),
+        ])?;
+
+        // Test: check should pass with same schema
+        let equal_schema = ok_plan.schema();
+        InvariantChecker.check(&ok_plan, &rule, equal_schema)?;
+
+        // Test: should fail with schema changed
+        let different_schema =
+            Arc::new(Schema::new(vec![Field::new("a", DataType::Boolean, false)]));
+        let expected_err = InvariantChecker
+            .check(&ok_plan, &rule, different_schema)
+            .unwrap_err();
+        assert!(expected_err.to_string().contains("PhysicalOptimizer rule 'OptimizerRuleWithSchemaCheck' failed, due to generate a different schema"));
+
+        // Test: should fail when extension node fails it's own invariant check
+        let failing_node: Arc<dyn ExecutionPlan> = Arc::new(InvariantFailsExtensionNode);
+        let expected_err = InvariantChecker
+            .check(&failing_node, &rule, ok_plan.schema())
+            .unwrap_err();
+        assert!(expected_err
+            .to_string()
+            .contains("extension node failed it's user-defined invariant check"));
+
+        // Test: should fail when descendent extension node fails
+        let failing_node: Arc<dyn ExecutionPlan> = Arc::new(InvariantFailsExtensionNode);
+        let invalid_plan = ok_node.with_new_children(vec![
+            Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?,
+            Arc::clone(&child),
+        ])?;
+        let expected_err = InvariantChecker
+            .check(&invalid_plan, &rule, ok_plan.schema())
+            .unwrap_err();
+        assert!(expected_err
+            .to_string()
+            .contains("extension node failed it's user-defined invariant check"));
+
+        Ok(())
+    }
 }

datafusion/physical-plan/src/execution_plan.rs

@@ -110,6 +110,16 @@ pub trait ExecutionPlan: Debug + DisplayAs + Send + Sync {
     /// trait, which is implemented for all `ExecutionPlan`s.
     fn properties(&self) -> &PlanProperties;
 
+    /// Returns an error if this individual node does not conform to its invariants.
+    /// These invariants are typically only checked in debug mode.
+    ///
+    /// A default set of invariants is provided in the default implementation.
+    /// Extension nodes can provide their own invariants.
+    fn check_node_invariants(&self) -> Result<()> {
+        // TODO
+        Ok(())
+    }
+
     /// Specifies the data distribution requirements for all the
     /// children for this `ExecutionPlan`, By default it's [[Distribution::UnspecifiedDistribution]] for each child,
     fn required_input_distribution(&self) -> Vec<Distribution> {