fix: realign explain changes to spiceai-52

Author: lukekim

Cargo.toml

@@ -1,7 +1,7 @@
 [workspace]
 resolver = "2"
 
-members = ["datafusion-federation"]
+members = ["datafusion-federation", "integration-test"]
 
 [workspace.package]
 version = "0.4.2"

datafusion-federation/Cargo.toml

@@ -26,11 +26,12 @@ async-trait.workspace = true
 datafusion.workspace = true
 async-stream.workspace = true
 arrow-json.workspace = true
+tracing = "0.1"
 
 [dev-dependencies]
 tokio.workspace = true
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
-tracing = "0.1.40"
+insta = { version = "1.42.0", features = ["filters"] }
 
 [[example]]
 name = "df-csv"

datafusion-federation/examples/df-csv-advanced.rs

@@ -6,12 +6,12 @@ use datafusion::{
     execution::{
         context::SessionContext, options::CsvReadOptions, session_state::SessionStateBuilder,
     },
-    optimizer::Optimizer,
+    optimizer::Analyzer,
 };
 
 use datafusion_federation::{
     sql::{MultiSchemaProvider, SQLFederationProvider, SQLSchemaProvider},
-    FederatedQueryPlanner, FederationOptimizerRule,
+    FederatedQueryPlanner, FederationAnalyzerRule,
 };
 
 use shared::{overwrite_default_schema, MockPostgresExecutor, MockSqliteExecutor};
@@ -115,15 +115,15 @@ async fn main() {
     /////////////////////
     //  Main(local)  DB
     /////////////////////
-    // Get the default optimizer rules
-    let mut rules = Optimizer::new().rules;
+    // Get the default analyzer rules
+    let mut rules = Analyzer::new().rules;
 
-    // Create a new federation optimizer rule and add it to the default rules
-    rules.push(Arc::new(FederationOptimizerRule::new()));
+    // Create a new federation analyzer rule and add it to the default rules
+    rules.push(Arc::new(FederationAnalyzerRule::new()));
 
-    // Create a new SessionState with the optimizer rule we created above
+    // Create a new SessionState with the analyzer rule we created above
     let state = SessionStateBuilder::new()
-        .with_optimizer_rules(rules)
+        .with_analyzer_rules(rules)
         .with_query_planner(Arc::new(FederatedQueryPlanner::new()))
         .build();
 

…-federation/src/optimizer/scan_result.rs …n-federation/src/analyzer/scan_result.rsdatafusion-federation/src/optimizer/scan_result.rs renamed to datafusion-federation/src/analyzer/scan_result.rs datafusion-federation/src/optimizer/scan_result.rs renamed to datafusion-federation/src/analyzer/scan_result.rs

@@ -1,35 +1,37 @@
-use datafusion::common::tree_node::TreeNodeRecursion;
-
 use crate::FederationProviderRef;
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::error::{DataFusionError, Result};
 
 /// Used to track if all sources, including tableScan, plan inputs and
 /// expressions, represents an un-ambiguous, none or a sole' [`crate::FederationProvider`].
 pub enum ScanResult {
+    /// [`LogicalPlan`] subtree has no [`FederationProvider`].
     None,
+
+    /// There is a single [`FederationProvider`] in the [`LogicalPlan`] subtree.
     Distinct(FederationProviderRef),
+
+    /// The [`LogicalPlan`] subtree has several different [`FederationProvider`]s.
     Ambiguous,
 }
 
 impl ScanResult {
-    pub fn merge(&mut self, other: Self) {
+    pub fn merge(&mut self, other: impl Into<Self>) {
+        let other = other.into();
         match (&self, &other) {
             (_, ScanResult::None) => {}
             (ScanResult::None, _) => *self = other,
             (ScanResult::Ambiguous, _) | (_, ScanResult::Ambiguous) => {
-                *self = ScanResult::Ambiguous
+                *self = ScanResult::Ambiguous;
             }
             (ScanResult::Distinct(provider), ScanResult::Distinct(other_provider)) => {
                 if provider != other_provider {
-                    *self = ScanResult::Ambiguous
+                    *self = ScanResult::Ambiguous;
                 }
             }
         }
     }
 
-    pub fn add(&mut self, provider: Option<FederationProviderRef>) {
-        self.merge(ScanResult::from(provider))
-    }
-
     pub fn is_ambiguous(&self) -> bool {
         matches!(self, ScanResult::Ambiguous)
     }
@@ -41,11 +43,13 @@ impl ScanResult {
         !self.is_none()
     }
 
-    pub fn unwrap(self) -> Option<FederationProviderRef> {
+    pub fn unwrap(self) -> Result<Option<FederationProviderRef>> {
         match self {
-            ScanResult::None => None,
-            ScanResult::Distinct(provider) => Some(provider),
-            ScanResult::Ambiguous => panic!("called `ScanResult::unwrap()` on a `Ambiguous` value"),
+            ScanResult::None => Ok(None),
+            ScanResult::Distinct(provider) => Ok(Some(provider)),
+            ScanResult::Ambiguous => Err(DataFusionError::External(
+                "called `ScanResult::unwrap()` on a `Ambiguous` value".into(),
+            )),
         }
     }
 
@@ -77,6 +81,19 @@ impl PartialEq<Option<FederationProviderRef>> for ScanResult {
     }
 }
 
+impl PartialEq for ScanResult {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (ScanResult::None, ScanResult::None) => true,
+            (ScanResult::Distinct(provider1), ScanResult::Distinct(provider2)) => {
+                provider1 == provider2
+            }
+            (ScanResult::Ambiguous, ScanResult::Ambiguous) => true,
+            _ => false,
+        }
+    }
+}
+
 impl std::fmt::Debug for ScanResult {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

datafusion-federation/src/lib.rs

@@ -1,4 +1,4 @@
-mod optimizer;
+mod analyzer;
 mod plan_node;
 pub mod schema_cast;
 #[cfg(feature = "sql")]
@@ -13,58 +13,73 @@ use std::{
 
 use datafusion::{
     execution::session_state::{SessionState, SessionStateBuilder},
-    optimizer::{optimizer::Optimizer, OptimizerRule},
+    logical_expr::LogicalPlan,
+    optimizer::{
+        analyzer::{
+            resolve_grouping_function::ResolveGroupingFunction, type_coercion::TypeCoercion,
+        },
+        Analyzer, AnalyzerRule,
+    },
 };
 
-pub use optimizer::{get_table_source, FederationOptimizerRule};
+pub use analyzer::{get_table_source, FederationAnalyzerRule};
 pub use plan_node::{
     FederatedPlanNode, FederatedPlanner, FederatedQueryPlanner, FederationPlanner,
 };
 pub use table_provider::{FederatedTableProviderAdaptor, FederatedTableSource};
 
 pub fn default_session_state() -> SessionState {
-    let rules = default_optimizer_rules();
+    let rules = default_analyzer_rules();
     SessionStateBuilder::new()
-        .with_optimizer_rules(rules)
+        .with_analyzer_rules(rules)
         .with_query_planner(Arc::new(FederatedQueryPlanner::new()))
         .with_default_features()
         .build()
 }
 
-pub fn default_optimizer_rules() -> Vec<Arc<dyn OptimizerRule + Send + Sync>> {
-    // Get the default optimizer
-    let df_default = Optimizer::new();
-    let mut default_rules = df_default.rules;
-
-    // Insert the FederationOptimizerRule after the ScalarSubqueryToJoin.
-    // This ensures ScalarSubquery are replaced before we try to federate.
-    let Some(pos) = default_rules
-        .iter()
-        .position(|x| x.name() == "scalar_subquery_to_join")
-    else {
-        panic!("Could not locate ScalarSubqueryToJoin");
-    };
-
-    // TODO: check if we should allow other optimizers to run before the federation rule.
+/// datafusion-federation customizes the order of the analyzer rules, since some of them are only relevant when `DataFusion` is executing the query,
+/// as opposed to when underlying federated query engines will execute the query.
+///
+/// This list should be kept in sync with the default rules in `Analyzer::new()`, but with the federation analyzer rule added.
+pub fn default_analyzer_rules() -> Vec<Arc<dyn AnalyzerRule + Send + Sync>> {
+    vec![
+        Arc::new(FederationAnalyzerRule::new()),
+        // The rest of these rules are run after the federation analyzer since they only affect internal DataFusion execution.
+        Arc::new(ResolveGroupingFunction::new()),
+        Arc::new(TypeCoercion::new()),
+    ]
+}
 
-    let federation_rule = Arc::new(FederationOptimizerRule::new());
-    default_rules.insert(pos + 1, federation_rule);
+pub type FederationProviderRef = Arc<dyn FederationProvider>;
 
-    default_rules
+impl From<Arc<Analyzer>> for FederationAnalyzerForLogicalPlan {
+    fn from(value: Arc<Analyzer>) -> Self {
+        Self::With(value)
+    }
 }
 
-pub type FederationProviderRef = Arc<dyn FederationProvider>;
-pub trait FederationProvider: Send + Sync {
+pub trait FederationProvider: Send + Sync + std::fmt::Debug {
     // Returns the name of the provider, used for comparison.
     fn name(&self) -> &str;
 
     // Returns the compute context in which this federation provider
     // will execute a query. For example: database instance & catalog.
     fn compute_context(&self) -> Option<String>;
 
-    // Returns an optimizer that can cut out part of the plan
-    // to federate it.
-    fn optimizer(&self) -> Option<Arc<Optimizer>>;
+    // Returns an analyzer that can cut out, and federate part of the [`LogicalPlan`].
+    //
+    // Returns:
+    //   - [`None`] if the provider cannot federate any plan (e.g. [`crate::analyzer::NopFederationProvider`]).
+    //   - Some(FederationAnalyzerForLogicalPlan::Unable) if the provider cannot federate this specific [`LogicalPlan`].
+    fn analyzer(&self, plan: &LogicalPlan) -> Option<FederationAnalyzerForLogicalPlan>;
+}
+
+/// [`LogicalPlan`] specific federation [`Analyzer`] from a [`FederationProvider`].
+#[derive(Debug)]
+pub enum FederationAnalyzerForLogicalPlan {
+    /// The [`FederationProvider`] cannot federate the [`LogicalPlan`].
+    Unable,
+    With(Arc<Analyzer>),
 }
 
 impl fmt::Display for dyn FederationProvider {
@@ -88,3 +103,29 @@ impl Hash for dyn FederationProvider {
 }
 
 impl Eq for dyn FederationProvider {}
+
+#[cfg(test)]
+mod tests {
+    use datafusion::optimizer::Analyzer;
+
+    /// Verifies that the default analyzer rules are in the expected order.
+    ///
+    /// If this test fails, `DataFusion` has modified the default analyzer rules and `get_analyzer_rules()` should be updated.
+    #[test]
+    fn test_verify_default_analyzer_rules() {
+        let default_rules = Analyzer::new().rules;
+        assert_eq!(
+            default_rules.len(),
+            2,
+            "Default analyzer rules have changed"
+        );
+        let expected_rule_names = vec!["resolve_grouping_function", "type_coercion"];
+        for (rule, expected_name) in default_rules.iter().zip(expected_rule_names.into_iter()) {
+            assert_eq!(
+                expected_name,
+                rule.name(),
+                "Default analyzer rule order has changed"
+            );
+        }
+    }
+}