feat: run physical requirement enforcement first, as a PhysicalAnalyzer phase

datafusion/core/src/execution/session_state.rs

@@ -75,10 +75,13 @@ use datafusion_optimizer::{
 };
 use datafusion_physical_expr::create_physical_expr;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use datafusion_physical_optimizer::analyzer::PhysicalAnalyzer;
 use datafusion_physical_optimizer::optimizer::PhysicalOptimizer;
 use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::operator_statistics::StatisticsRegistry;
-use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session};
+use datafusion_session::{
+    PhysicalAnalyzerRule, PhysicalOptimizerContext, PhysicalOptimizerRule, Session,
+};
 #[cfg(feature = "sql")]
 use datafusion_sql::{
     parser::{DFParserBuilder, Statement},
@@ -168,6 +171,9 @@ struct SessionStateInner {
     type_planner: Option<Arc<dyn TypePlanner>>,
     /// Responsible for optimizing a logical plan
     optimizer: Optimizer,
+    /// Responsible for enforcing invariants on a physical execution plan
+    /// (distribution, ordering) before optimization
+    physical_analyzers: PhysicalAnalyzer,
     /// Responsible for optimizing a physical execution plan
     physical_optimizers: PhysicalOptimizer,
     /// Responsible for planning `LogicalPlan`s, and `ExecutionPlan`
@@ -272,6 +278,7 @@ impl Debug for SessionState {
         ret.field("query_planners", &self.inner.query_planner)
             .field("analyzer", &self.inner.analyzer)
             .field("optimizer", &self.inner.optimizer)
+            .field("physical_analyzers", &self.inner.physical_analyzers)
             .field("physical_optimizers", &self.inner.physical_optimizers)
             .field("table_functions", &self.inner.table_functions)
             .field("scalar_functions", &self.inner.scalar_functions)
@@ -312,6 +319,10 @@ impl Session for SessionState {
         SessionState::physical_optimizers(self)
     }
 
+    fn physical_analyzers(&self) -> &[Arc<dyn PhysicalAnalyzerRule + Send + Sync>] {
+        SessionState::physical_analyzers(self)
+    }
+
     fn statistics_registry(&self) -> Option<&StatisticsRegistry> {
         SessionState::statistics_registry(self)
     }
@@ -454,6 +465,19 @@ impl SessionState {
         self
     }
 
+    /// Add `physical_analyzer_rule` to the end of the list of
+    /// [`PhysicalAnalyzerRule`]s run before physical optimization.
+    pub fn add_physical_analyzer_rule(
+        &mut self,
+        physical_analyzer_rule: Arc<dyn PhysicalAnalyzerRule + Send + Sync>,
+    ) -> &Self {
+        Arc::make_mut(&mut self.inner)
+            .physical_analyzers
+            .rules
+            .push(physical_analyzer_rule);
+        self
+    }
+
     // the add_optimizer_rule takes an owned reference
     // it should probably be renamed to `with_optimizer_rule` to follow builder style
     // and `add_optimizer_rule` that takes &mut self added instead of this
@@ -946,6 +970,11 @@ impl SessionState {
         &self.inner.physical_optimizers.rules
     }
 
+    /// Return the physical analyzers
+    pub fn physical_analyzers(&self) -> &[Arc<dyn PhysicalAnalyzerRule + Send + Sync>] {
+        &self.inner.physical_analyzers.rules
+    }
+
     /// return the configuration options
     pub fn config_options(&self) -> &Arc<ConfigOptions> {
         self.inner.config.options()
@@ -1144,6 +1173,7 @@ pub struct SessionStateBuilder {
     #[cfg(feature = "sql")]
     type_planner: Option<Arc<dyn TypePlanner>>,
     optimizer: Option<Optimizer>,
+    physical_analyzers: Option<PhysicalAnalyzer>,
     physical_optimizers: Option<PhysicalOptimizer>,
     query_planner: Option<Arc<dyn QueryPlanner + Send + Sync>>,
     catalog_list: Option<Arc<dyn CatalogProviderList>>,
@@ -1167,6 +1197,7 @@ pub struct SessionStateBuilder {
     // fields to support convenience functions
     analyzer_rules: Option<Vec<Arc<dyn AnalyzerRule + Send + Sync>>>,
     optimizer_rules: Option<Vec<Arc<dyn OptimizerRule + Send + Sync>>>,
+    physical_analyzer_rules: Option<Vec<Arc<dyn PhysicalAnalyzerRule + Send + Sync>>>,
     physical_optimizer_rules: Option<Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>>>,
 }
 
@@ -1188,6 +1219,7 @@ impl SessionStateBuilder {
             #[cfg(feature = "sql")]
             type_planner: None,
             optimizer: None,
+            physical_analyzers: None,
             physical_optimizers: None,
             query_planner: None,
             catalog_list: None,
@@ -1211,6 +1243,7 @@ impl SessionStateBuilder {
             // fields to support convenience functions
             analyzer_rules: None,
             optimizer_rules: None,
+            physical_analyzer_rules: None,
             physical_optimizer_rules: None,
         }
     }
@@ -1250,6 +1283,7 @@ impl SessionStateBuilder {
             #[cfg(feature = "sql")]
             type_planner: existing.type_planner,
             optimizer: Some(existing.optimizer),
+            physical_analyzers: Some(existing.physical_analyzers),
             physical_optimizers: Some(existing.physical_optimizers),
             query_planner: Some(existing.query_planner),
             catalog_list: Some(existing.catalog_list),
@@ -1277,6 +1311,7 @@ impl SessionStateBuilder {
             // fields to support convenience functions
             analyzer_rules: None,
             optimizer_rules: None,
+            physical_analyzer_rules: None,
             physical_optimizer_rules: None,
         }
     }
@@ -1442,6 +1477,27 @@ impl SessionStateBuilder {
         self
     }
 
+    /// Set the [`PhysicalAnalyzerRule`]s run before physical optimization.
+    pub fn with_physical_analyzer_rules(
+        mut self,
+        physical_analyzers: Vec<Arc<dyn PhysicalAnalyzerRule + Send + Sync>>,
+    ) -> Self {
+        self.physical_analyzers = Some(PhysicalAnalyzer::with_rules(physical_analyzers));
+        self
+    }
+
+    /// Add `physical_analyzer_rule` to the end of the list of
+    /// [`PhysicalAnalyzerRule`]s run before physical optimization.
+    pub fn with_physical_analyzer_rule(
+        mut self,
+        physical_analyzer_rule: Arc<dyn PhysicalAnalyzerRule + Send + Sync>,
+    ) -> Self {
+        let mut rules = self.physical_analyzer_rules.unwrap_or_default();
+        rules.push(physical_analyzer_rule);
+        self.physical_analyzer_rules = Some(rules);
+        self
+    }
+
     /// Set the [`QueryPlanner`]
     pub fn with_query_planner(
         mut self,
@@ -1689,6 +1745,7 @@ impl SessionStateBuilder {
             #[cfg(feature = "sql")]
             type_planner,
             optimizer,
+            physical_analyzers,
             physical_optimizers,
             query_planner,
             catalog_list,
@@ -1711,6 +1768,7 @@ impl SessionStateBuilder {
             statistics_registry,
             analyzer_rules,
             optimizer_rules,
+            physical_analyzer_rules,
             physical_optimizer_rules,
         } = self;
 
@@ -1726,6 +1784,7 @@ impl SessionStateBuilder {
             #[cfg(feature = "sql")]
             type_planner,
             optimizer: optimizer.unwrap_or_default(),
+            physical_analyzers: physical_analyzers.unwrap_or_default(),
             physical_optimizers: physical_optimizers.unwrap_or_default(),
             query_planner: query_planner
                 .unwrap_or_else(|| Arc::new(DefaultQueryPlanner {})),
@@ -1871,6 +1930,14 @@ impl SessionStateBuilder {
             }
         }
 
+        if let Some(physical_analyzer_rules) = physical_analyzer_rules {
+            let physical_analyzers =
+                &mut Arc::make_mut(&mut state.inner).physical_analyzers;
+            for physical_analyzer_rule in physical_analyzer_rules {
+                physical_analyzers.rules.push(physical_analyzer_rule);
+            }
+        }
+
         if let Some(physical_optimizer_rules) = physical_optimizer_rules {
             let physical_optimizers =
                 &mut Arc::make_mut(&mut state.inner).physical_optimizers;
@@ -1919,6 +1986,11 @@ impl SessionStateBuilder {
         &mut self.physical_optimizers
     }
 
+    /// Returns the current physical_analyzers value
+    pub fn physical_analyzers(&mut self) -> &mut Option<PhysicalAnalyzer> {
+        &mut self.physical_analyzers
+    }
+
     /// Returns the current query_planner value
     pub fn query_planner(&mut self) -> &mut Option<Arc<dyn QueryPlanner + Send + Sync>> {
         &mut self.query_planner
@@ -2030,6 +2102,13 @@ impl SessionStateBuilder {
     ) -> &mut Option<Vec<Arc<dyn PhysicalOptimizerRule + Send + Sync>>> {
         &mut self.physical_optimizer_rules
     }
+
+    /// Returns the current physical_analyzer_rules value
+    pub fn physical_analyzer_rules(
+        &mut self,
+    ) -> &mut Option<Vec<Arc<dyn PhysicalAnalyzerRule + Send + Sync>>> {
+        &mut self.physical_analyzer_rules
+    }
 }
 
 impl Debug for SessionStateBuilder {
@@ -2058,6 +2137,8 @@ impl Debug for SessionStateBuilder {
             .field("analyzer", &self.analyzer)
             .field("optimizer_rules", &self.optimizer_rules)
             .field("optimizer", &self.optimizer)
+            .field("physical_analyzer_rules", &self.physical_analyzer_rules)
+            .field("physical_analyzers", &self.physical_analyzers)
             .field("physical_optimizer_rules", &self.physical_optimizer_rules)
             .field("physical_optimizers", &self.physical_optimizers)
             .field("table_functions", &self.table_functions)