Introduce LogicalPlan invariants, begin automatically checking them

datafusion/expr/src/logical_plan/mod.rs

@@ -20,6 +20,8 @@ mod ddl;
 pub mod display;
 pub mod dml;
 mod extension;
+pub(crate) mod invariants;
+pub use invariants::assert_expected_schema;
 mod plan;
 mod statement;
 pub mod tree_node;

datafusion/expr/src/logical_plan/plan.rs

@@ -24,6 +24,7 @@ use std::hash::{Hash, Hasher};
 use std::sync::{Arc, OnceLock};
 
 use super::dml::CopyTo;
+use super::invariants::assert_unique_field_names;
 use super::DdlStatement;
 use crate::builder::{change_redundant_column, unnest_with_options};
 use crate::expr::{Placeholder, Sort as SortExpr, WindowFunction};
@@ -1125,6 +1126,14 @@ impl LogicalPlan {
         }
     }
 
+    /// These are invariants to hold true for each logical plan.
+    pub fn assert_invariants(&self) -> Result<()> {
+        // Refer to <https://datafusion.apache.org/contributor-guide/specification/invariants.html#relation-name-tuples-in-logical-fields-and-logical-columns-are-unique>
+        assert_unique_field_names(self)?;
+
+        Ok(())
+    }
+
     /// Helper for [Self::with_new_exprs] to use when no expressions are expected.
     #[inline]
     #[allow(clippy::needless_pass_by_value)] // expr is moved intentionally to ensure it's not used again

datafusion/optimizer/src/analyzer/mod.rs

@@ -137,6 +137,14 @@ impl Analyzer {
     where
         F: FnMut(&LogicalPlan, &dyn AnalyzerRule),
     {
+        // verify at the start, before the first LP analyzer pass.
+        assert_valid_semantic_plan(&plan).map_err(|e| {
+            DataFusionError::Context(
+                "check_plan_before_analyzers".to_string(),
+                Box::new(e),
+            )
+        })?;
+
         let start_time = Instant::now();
         let mut new_plan = plan;
 
@@ -164,18 +172,38 @@ impl Analyzer {
             log_plan(rule.name(), &new_plan);
             observer(&new_plan, rule.as_ref());
         }
-        // for easier display in explain output
-        check_plan(&new_plan).map_err(|e| {
+
+        // verify at the end, after the last LP analyzer pass.
+        assert_valid_semantic_plan(&new_plan).map_err(|e| {
             DataFusionError::Context("check_analyzed_plan".to_string(), Box::new(e))
         })?;
+
         log_plan("Final analyzed plan", &new_plan);
         debug!("Analyzer took {} ms", start_time.elapsed().as_millis());
         Ok(new_plan)
     }
 }
 
-/// Do necessary check and fail the invalid plan
-fn check_plan(plan: &LogicalPlan) -> Result<()> {
+/// These are invariants which should hold true before and after each analyzer.
+///
+/// This differs from [`LogicalPlan::assert_invariants`], which addresses if a singular
+/// LogicalPlan is valid. Instead this address if the analyzer (before and after)
+/// is valid based upon permitted changes.
+///
+/// Does not check elements which are mutated by the analyzers (e.g. the schema).
+fn assert_valid_semantic_plan(plan: &LogicalPlan) -> Result<()> {
+    plan.assert_invariants()?;
+
+    // TODO: should this be moved to LogicalPlan::assert_invariants?
+    assert_subqueries_are_valid(plan)?;
+
+    Ok(())
+}
+
+/// Asserts that the subqueries are structured properly with valid node placement.
+///
+/// Refer to [`check_subquery_expr`] for more details.
+fn assert_subqueries_are_valid(plan: &LogicalPlan) -> Result<()> {
     plan.apply_with_subqueries(|plan: &LogicalPlan| {
         plan.apply_expressions(|expr| {
             // recursively look for subqueries

datafusion/optimizer/src/analyzer/subquery.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::analyzer::check_plan;
+use crate::analyzer::assert_subqueries_are_valid;
 use crate::utils::collect_subquery_cols;
 use recursive::recursive;
 
@@ -37,7 +37,7 @@ pub fn check_subquery_expr(
     inner_plan: &LogicalPlan,
     expr: &Expr,
 ) -> Result<()> {
-    check_plan(inner_plan)?;
+    assert_subqueries_are_valid(inner_plan)?;
     if let Expr::ScalarSubquery(subquery) = expr {
         // Scalar subquery should only return one column
         if subquery.subquery.schema().fields().len() > 1 {

datafusion/optimizer/src/optimizer.rs

@@ -21,6 +21,7 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use chrono::{DateTime, Utc};
+use datafusion_expr::assert_expected_schema;
 use datafusion_expr::registry::FunctionRegistry;
 use log::{debug, warn};
 
@@ -355,6 +356,14 @@ impl Optimizer {
     where
         F: FnMut(&LogicalPlan, &dyn OptimizerRule),
     {
+        // verify LP is valid, before the first LP optimizer pass.
+        plan.assert_invariants().map_err(|e| {
+            DataFusionError::Context(
+                "check_plan_before_optimizers".to_string(),
+                Box::new(e),
+            )
+        })?;
+
         let start_time = Instant::now();
         let options = config.options();
         let mut new_plan = plan;
@@ -384,9 +393,16 @@ impl Optimizer {
                     // rule handles recursion itself
                     None => optimize_plan_node(new_plan, rule.as_ref(), config),
                 }
-                // verify the rule didn't change the schema
                 .and_then(|tnr| {
-                    assert_schema_is_the_same(rule.name(), &starting_schema, &tnr.data)?;
+                    // verify after each optimizer pass.
+                    assert_valid_optimization(rule.name(), &tnr.data, &starting_schema)
+                        .map_err(|e| {
+                        DataFusionError::Context(
+                            "check_optimized_plan".to_string(),
+                            Box::new(e),
+                        )
+                    })?;
+
                     Ok(tnr)
                 });
 
@@ -445,35 +461,36 @@ impl Optimizer {
             }
             i += 1;
         }
+
+        // verify LP is valid, after the last optimizer pass.
+        new_plan.assert_invariants().map_err(|e| {
+            DataFusionError::Context(
+                "check_plan_after_optimizers".to_string(),
+                Box::new(e),
+            )
+        })?;
+
         log_plan("Final optimized plan", &new_plan);
         debug!("Optimizer took {} ms", start_time.elapsed().as_millis());
         Ok(new_plan)
     }
 }
 
-/// Returns an error if `new_plan`'s schema is different than `prev_schema`
+/// These are invariants which should hold true before and after each optimization.
 ///
-/// It ignores metadata and nullability.
-pub(crate) fn assert_schema_is_the_same(
+/// This differs from [`LogicalPlan::assert_invariants`], which addresses if a singular
+/// LogicalPlan is valid. Instead this address if the optimization (before and after)
+/// is valid based upon permitted changes.
+fn assert_valid_optimization(
     rule_name: &str,
-    prev_schema: &DFSchema,
-    new_plan: &LogicalPlan,
+    plan: &LogicalPlan,
+    prev_schema: &Arc<DFSchema>,
 ) -> Result<()> {
-    let equivalent = new_plan.schema().equivalent_names_and_types(prev_schema);
-
-    if !equivalent {
-        let e = DataFusionError::Internal(format!(
-            "Failed due to a difference in schemas, original schema: {:?}, new schema: {:?}",
-            prev_schema,
-            new_plan.schema()
-        ));
-        Err(DataFusionError::Context(
-            String::from(rule_name),
-            Box::new(e),
-        ))
-    } else {
-        Ok(())
-    }
+    // verify invariant: optimizer rule didn't change the schema
+    // Refer to <https://datafusion.apache.org/contributor-guide/specification/invariants.html#logical-schema-is-invariant-under-logical-optimization>
+    assert_expected_schema(rule_name, prev_schema, plan)?;
+
+    Ok(())
 }
 
 #[cfg(test)]