Merge pull request #61 from spiceai/jeadie/25-11-12/avoid-filter-federate

Avoid 'federation_sql_push_down_filter' correctly

Author: Jeadie

datafusion-federation/src/analyzer/mod.rs

@@ -1,7 +1,7 @@
 mod scan_result;
 
-use crate::FederationProvider;
 use crate::{FederatedTableProviderAdaptor, FederatedTableSource, FederationProviderRef};
+use crate::{FederationAnalyzerForLogicalPlan, FederationProvider};
 use datafusion::logical_expr::{col, expr::InSubquery, LogicalPlanBuilder};
 use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion;
 use datafusion::optimizer::push_down_filter::PushDownFilter;
@@ -243,12 +243,12 @@ impl FederationAnalyzerRule {
                     // The largest sub-plan is higher up.
                     return Ok((None, ScanResult::Distinct(provider)));
                 }
-                (true, Some(analyzer)) => {
+                (true, Some(FederationAnalyzerForLogicalPlan::With(analyzer))) => {
                     // If this is the root plan node; federate the entire plan
                     let optimized = analyzer.execute_and_check(plan.clone(), config, |_, _| {})?;
                     return Ok((Some(optimized), ScanResult::None));
                 }
-                (_, None) => {
+                (_, None | Some(FederationAnalyzerForLogicalPlan::Unable)) => {
                     // Provider CAN'T federate this specific plan shape
                     // Fall through to try federating children instead
                 }
@@ -282,7 +282,9 @@ impl FederationAnalyzerRule {
                     return Ok(original_input);
                 };
 
-                let Some(analyzer) = provider.analyzer(&original_input) else {
+                let Some(FederationAnalyzerForLogicalPlan::With(analyzer)) =
+                    provider.analyzer(&original_input)
+                else {
                     // Either provider has no analyzer, or cannot federate [`LogicalPlan`].
                     return Ok(original_input);
                 };
@@ -413,26 +415,19 @@ impl FederationAnalyzerRule {
 #[derive(Debug)]
 pub(crate) struct NopFederationProvider {}
 
-pub static NOP_NAME: &str = "nop";
-
 impl FederationProvider for NopFederationProvider {
     fn name(&self) -> &str {
-        NOP_NAME
+        "nop"
     }
 
     fn compute_context(&self) -> Option<String> {
         None
     }
 
-    fn analyzer(&self, _plan: &LogicalPlan) -> Option<Arc<datafusion::optimizer::Analyzer>> {
+    fn analyzer(&self, _plan: &LogicalPlan) -> Option<FederationAnalyzerForLogicalPlan> {
         None
     }
 }
-impl NopFederationProvider {
-    pub fn is_nop(provider: Arc<dyn FederationProvider>) -> bool {
-        provider.name() == NOP_NAME
-    }
-}
 
 /// Recursively find the [`FederationProvider`] for all [`TableReference`] instances in the plan.
 /// This is used to resolve the federation providers for [`Expr::OuterReferenceColumn`].
@@ -474,7 +469,7 @@ fn wrap_projection(plan: LogicalPlan) -> Result<LogicalPlan> {
 fn contains_federated_table(plan: &LogicalPlan) -> Result<bool> {
     let federated_table_exists = plan.exists(|x| {
         if let (Some(provider), _) = get_leaf_provider(x)? {
-            return Ok(!NopFederationProvider::is_nop(provider));
+            return Ok(provider.analyzer(plan).is_some());
         }
         Ok(false)
     })?;

datafusion-federation/src/lib.rs

@@ -52,6 +52,12 @@ pub fn default_analyzer_rules() -> Vec<Arc<dyn AnalyzerRule + Send + Sync>> {
 
 pub type FederationProviderRef = Arc<dyn FederationProvider>;
 
+impl From<Arc<Analyzer>> for FederationAnalyzerForLogicalPlan {
+    fn from(value: Arc<Analyzer>) -> Self {
+        Self::With(value)
+    }
+}
+
 pub trait FederationProvider: Send + Sync + std::fmt::Debug {
     // Returns the name of the provider, used for comparison.
     fn name(&self) -> &str;
@@ -62,10 +68,18 @@ pub trait FederationProvider: Send + Sync + std::fmt::Debug {
 
     // Returns an analyzer that can cut out, and federate part of the [`LogicalPlan`].
     //
-    // Returns [`None`] if either the provider cannot federate any plan
-    // (e.g. [`crate::analyzer::NopFederationProvider`]), or cannot federate
-    //  this specific [`LogicalPlan`].
-    fn analyzer(&self, plan: &LogicalPlan) -> Option<Arc<Analyzer>>;
+    // 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 {

datafusion-federation/src/sql/mod.rs

@@ -35,8 +35,8 @@ pub use table::{RemoteTable, SQLTable, SQLTableSource};
 pub use table_reference::{MultiPartTableReference, RemoteTableRef};
 
 use crate::{
-    get_table_source, schema_cast, FederatedPlanNode, FederationAnalyzerRule, FederationPlanner,
-    FederationProvider,
+    get_table_source, schema_cast, FederatedPlanNode, FederationAnalyzerForLogicalPlan,
+    FederationAnalyzerRule, FederationPlanner, FederationProvider,
 };
 
 /// Returns a federation analyzer rule that is optimized for SQL federation.
@@ -75,11 +75,11 @@ impl FederationProvider for SQLFederationProvider {
         self.executor.compute_context()
     }
 
-    fn analyzer(&self, plan: &LogicalPlan) -> Option<Arc<Analyzer>> {
+    fn analyzer(&self, plan: &LogicalPlan) -> Option<FederationAnalyzerForLogicalPlan> {
         if self.executor.can_execute_plan(plan) {
-            Some(Arc::clone(&self.analyzer))
+            Some(Arc::clone(&self.analyzer).into())
         } else {
-            None
+            Some(FederationAnalyzerForLogicalPlan::Unable)
         }
     }
 }
@@ -595,7 +595,7 @@ mod tests {
         // pushed down since it will be optimised into `Filter: table_a1.a > Int64(0)`.
         insta::assert_snapshot!(ctx
             .sql(
-                r#"SELECT non_federate, b, c FROM (SELECT a AS 'non_federate', b, c FROM table_a1) WHERE non_federate > 0"#,
+                r#"SELECT a as non_federate, b, c FROM (SELECT a, b, c FROM table_a1) WHERE a > 0"#,
             )
             .await?
             .into_optimized_plan()?