Fix for all TPCH queries

Author: phillipleblanc

datafusion-federation/src/analyzer/mod.rs

@@ -5,7 +5,7 @@ use crate::{FederatedTableProviderAdaptor, FederatedTableSource, FederationProvi
 use datafusion::logical_expr::{col, expr::InSubquery, LogicalPlanBuilder};
 use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion;
 use datafusion::optimizer::push_down_filter::PushDownFilter;
-use datafusion::optimizer::{Optimizer, OptimizerContext};
+use datafusion::optimizer::{Optimizer, OptimizerContext, OptimizerRule};
 use datafusion::{
     common::tree_node::{Transformed, TreeNode, TreeNodeRecursion},
     config::ConfigOptions,
@@ -19,10 +19,6 @@ use scan_result::ScanResult;
 use std::collections::HashMap;
 use std::sync::Arc;
 
-pub use optimize_projections::OptimizeProjections;
-
-mod optimize_projections;
-
 /// An analyzer rule to identifying sub-plans to federate
 ///
 /// The analyzer logic walks over the plan, look for the largest subtrees that only have
@@ -65,11 +61,7 @@ impl AnalyzerRule for FederationAnalyzerRule {
 impl Default for FederationAnalyzerRule {
     fn default() -> Self {
         Self {
-            optimizer: Optimizer::with_rules(vec![
-                Arc::new(EliminateNestedUnion::new()),
-                Arc::new(PushDownFilter::new()),
-                Arc::new(OptimizeProjections::new()),
-            ]),
+            optimizer: Optimizer::with_rules(Self::default_optimizer_rules()),
         }
     }
 }
@@ -79,6 +71,13 @@ impl FederationAnalyzerRule {
         Self::default()
     }
 
+    pub fn default_optimizer_rules() -> Vec<Arc<dyn OptimizerRule + Send + Sync>> {
+        vec![
+            Arc::new(EliminateNestedUnion::new()),
+            Arc::new(PushDownFilter::new()),
+        ]
+    }
+
     /// Override the default optimizer with custom rules
     pub fn with_optimizer(mut self, optimizer: Optimizer) -> Self {
         self.optimizer = optimizer;

datafusion-federation/src/sql/mod.rs

@@ -1,6 +1,7 @@
 mod analyzer;
 pub mod ast_analyzer;
 mod executor;
+pub mod optimizer;
 mod schema;
 mod table;
 mod table_reference;
@@ -17,7 +18,7 @@ use datafusion::{
     error::{DataFusionError, Result},
     execution::{context::SessionState, TaskContext},
     logical_expr::{Extension, LogicalPlan},
-    optimizer::{Analyzer, AnalyzerRule},
+    optimizer::{eliminate_nested_union::EliminateNestedUnion, Analyzer, AnalyzerRule, Optimizer},
     physical_expr::EquivalenceProperties,
     physical_plan::{
         execution_plan::{Boundedness, EmissionType},
@@ -26,16 +27,27 @@ use datafusion::{
     },
     sql::{sqlparser::ast::Statement, unparser::Unparser},
 };
+use optimizer::{OptimizeProjectionsFederation, PushDownFilterFederation};
 
 pub use executor::{LogicalOptimizer, SQLExecutor, SQLExecutorRef};
 pub use schema::{MultiSchemaProvider, SQLSchemaProvider};
 pub use table::{RemoteTable, SQLTableSource};
 pub use table_reference::RemoteTableRef;
 
 use crate::{
-    get_table_source, schema_cast, FederatedPlanNode, FederationPlanner, FederationProvider,
+    get_table_source, schema_cast, FederatedPlanNode, FederationAnalyzerRule, FederationPlanner,
+    FederationProvider,
 };
 
+/// Returns a federation analyzer rule that is optimized for SQL federation.
+pub fn federation_analyzer_rule() -> FederationAnalyzerRule {
+    FederationAnalyzerRule::new().with_optimizer(Optimizer::with_rules(vec![
+        Arc::new(EliminateNestedUnion::new()),
+        Arc::new(PushDownFilterFederation::new()),
+        Arc::new(OptimizeProjectionsFederation::new()),
+    ]))
+}
+
 // SQLFederationProvider provides federation to SQL DMBSs.
 #[derive(Debug)]
 pub struct SQLFederationProvider {

datafusion-federation/src/sql/optimizer/mod.rs

@@ -0,0 +1,7 @@
+//! A collection of optimizer rules that benefit SQL federation.
+
+mod push_down_filter;
+pub use push_down_filter::PushDownFilterFederation;
+
+mod optimize_projections;
+pub use optimize_projections::OptimizeProjectionsFederation;

…src/analyzer/optimize_projections/mod.rs …ql/optimizer/optimize_projections/mod.rsdatafusion-federation/src/analyzer/optimize_projections/mod.rs renamed to datafusion-federation/src/sql/optimizer/optimize_projections/mod.rs datafusion-federation/src/analyzer/optimize_projections/mod.rs renamed to datafusion-federation/src/sql/optimizer/optimize_projections/mod.rs

@@ -45,16 +45,16 @@ type Result<T, E = DataFusionError> = result::Result<T, E>;
 mod required_indices;
 
 #[derive(Default, Debug)]
-pub struct OptimizeProjections {}
+pub struct OptimizeProjectionsFederation {}
 
-impl OptimizeProjections {
+impl OptimizeProjectionsFederation {
     #[must_use]
     pub fn new() -> Self {
         Self {}
     }
 }
 
-impl OptimizerRule for OptimizeProjections {
+impl OptimizerRule for OptimizeProjectionsFederation {
     fn name(&self) -> &str {
         "federation_optimize_projections"
     }
@@ -947,7 +947,11 @@ mod tests {
     fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
 
     fn assert_optimized_plan_equal(plan: LogicalPlan, expected: &str) -> Result<()> {
-        assert_optimized_plan_eq(Arc::new(super::OptimizeProjections::new()), plan, expected)
+        assert_optimized_plan_eq(
+            Arc::new(super::OptimizeProjectionsFederation::new()),
+            plan,
+            expected,
+        )
     }
 
     pub fn assert_optimized_plan_eq(
@@ -967,7 +971,8 @@ mod tests {
     }
 
     fn optimize(plan: LogicalPlan) -> Result<LogicalPlan> {
-        let optimizer = Optimizer::with_rules(vec![Arc::new(super::OptimizeProjections::new())]);
+        let optimizer =
+            Optimizer::with_rules(vec![Arc::new(super::OptimizeProjectionsFederation::new())]);
         let optimized_plan = optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
         Ok(optimized_plan)
     }

…optimize_projections/required_indices.rs …optimize_projections/required_indices.rsdatafusion-federation/src/analyzer/optimize_projections/required_indices.rs renamed to datafusion-federation/src/sql/optimizer/optimize_projections/required_indices.rs datafusion-federation/src/analyzer/optimize_projections/required_indices.rs renamed to datafusion-federation/src/sql/optimizer/optimize_projections/required_indices.rs

@@ -0,0 +1,64 @@
+use datafusion::{
+    common::{error::Result, tree_node::Transformed},
+    logical_expr::LogicalPlan,
+    optimizer::{push_down_filter::PushDownFilter, ApplyOrder, OptimizerConfig, OptimizerRule},
+};
+
+/// A wrapper around DataFusion's [`PushDownFilter`](https://github.com/apache/datafusion/blob/main/datafusion/optimizer/src/push_down_filter.rs) rule.
+///
+/// This wrapper prevents running the rule in a way that would break the SQL unparser, i.e. pushing a filter beyond a SubqueryAlias.
+#[derive(Default, Debug)]
+pub struct PushDownFilterFederation {
+    inner: PushDownFilter,
+}
+
+impl PushDownFilterFederation {
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+
+impl OptimizerRule for PushDownFilterFederation {
+    fn name(&self) -> &str {
+        "federation_sql_push_down_filter"
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        self.inner.apply_order()
+    }
+
+    fn supports_rewrite(&self) -> bool {
+        self.inner.supports_rewrite()
+    }
+
+    fn rewrite(
+        &self,
+        plan: LogicalPlan,
+        config: &dyn OptimizerConfig,
+    ) -> Result<Transformed<LogicalPlan>> {
+        if !should_run_rule_for_node(&plan) {
+            return Ok(Transformed::no(plan));
+        }
+
+        self.inner.rewrite(plan, config)
+    }
+}
+
+fn should_run_rule_for_node(node: &LogicalPlan) -> bool {
+    if let LogicalPlan::Filter(x) = node {
+        // Applying the `push_down_filter_rule` to certain nodes like `SubqueryAlias`, `Aggregate`, and `CrossJoin`
+        // can cause issues during unparsing, thus the optimization is only applied to nodes that are currently supported.
+        matches!(
+            x.input.as_ref(),
+            LogicalPlan::Join(_)
+                | LogicalPlan::TableScan(_)
+                | LogicalPlan::Projection(_)
+                | LogicalPlan::Filter(_)
+                | LogicalPlan::Distinct(_)
+                | LogicalPlan::Sort(_)
+        )
+    } else {
+        true
+    }
+}