datafusion_federation/sql/
mod.rs

1mod analyzer;
2pub mod ast_analyzer;
3mod executor;
4mod schema;
5mod table;
6mod table_reference;
7
8use std::{any::Any, fmt, sync::Arc, vec};
9
10use analyzer::RewriteTableScanAnalyzer;
11use async_trait::async_trait;
12use datafusion::{
13    arrow::datatypes::{Schema, SchemaRef},
14    catalog::Session,
15    common::{
16        tree_node::{Transformed, TreeNode, TreeNodeRecursion},
17        Statistics,
18    },
19    config::ConfigOptions,
20    error::{DataFusionError, Result},
21    execution::TaskContext,
22    logical_expr::{Extension, LogicalPlan},
23    optimizer::{optimizer::Optimizer, OptimizerConfig, OptimizerRule},
24    physical_expr::EquivalenceProperties,
25    physical_plan::{
26        apply_expression_roots,
27        execution_plan::{Boundedness, EmissionType},
28        filter_pushdown::{ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation},
29        metrics::MetricsSet,
30        DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, PlanProperties,
31        SendableRecordBatchStream, StatisticsArgs,
32    },
33    sql::{sqlparser::ast::Statement, unparser::Unparser},
34};
35
36pub use executor::{AstAnalyzer, LogicalOptimizer, SQLExecutor, SQLExecutorRef, SqlQueryRewriter};
37pub use schema::{MultiSchemaProvider, SQLSchemaProvider};
38pub use table::{RemoteTable, SQLTable, SQLTableSource};
39pub use table_reference::RemoteTableRef;
40
41use crate::{
42    get_table_source, schema_cast, FederatedPlanNode, FederationPlanner, FederationProvider,
43};
44
45// SQLFederationProvider provides federation to SQL DMBSs.
46#[derive(Debug)]
47pub struct SQLFederationProvider {
48    pub optimizer: Arc<Optimizer>,
49    pub executor: Arc<dyn SQLExecutor>,
50}
51
52impl SQLFederationProvider {
53    pub fn new(executor: Arc<dyn SQLExecutor>) -> Self {
54        Self {
55            optimizer: Arc::new(Optimizer::with_rules(vec![Arc::new(
56                SQLFederationOptimizerRule::new(executor.clone()),
57            )])),
58            executor,
59        }
60    }
61}
62
63impl FederationProvider for SQLFederationProvider {
64    fn name(&self) -> &str {
65        "sql_federation_provider"
66    }
67
68    fn compute_context(&self) -> Option<String> {
69        self.executor.compute_context()
70    }
71
72    fn optimizer(&self) -> Option<Arc<Optimizer>> {
73        Some(self.optimizer.clone())
74    }
75}
76
77#[derive(Debug)]
78struct SQLFederationOptimizerRule {
79    planner: Arc<SQLFederationPlanner>,
80}
81
82impl SQLFederationOptimizerRule {
83    pub fn new(executor: Arc<dyn SQLExecutor>) -> Self {
84        Self {
85            planner: Arc::new(SQLFederationPlanner::new(Arc::clone(&executor))),
86        }
87    }
88}
89
90impl OptimizerRule for SQLFederationOptimizerRule {
91    /// Try to rewrite `plan` to an optimized form, returning `Transformed::yes`
92    /// if the plan was rewritten and `Transformed::no` if it was not.
93    ///
94    /// Note: this function is only called if [`Self::supports_rewrite`] returns
95    /// true. Otherwise the Optimizer calls  [`Self::try_optimize`]
96    fn rewrite(
97        &self,
98        plan: LogicalPlan,
99        _config: &dyn OptimizerConfig,
100    ) -> Result<Transformed<LogicalPlan>> {
101        if let LogicalPlan::Extension(Extension { ref node }) = plan {
102            if node.name() == "Federated" {
103                // Avoid attempting double federation
104                return Ok(Transformed::no(plan));
105            }
106        }
107
108        let fed_plan = FederatedPlanNode::new(plan.clone(), self.planner.clone());
109        let ext_node = Extension {
110            node: Arc::new(fed_plan),
111        };
112
113        let mut plan = LogicalPlan::Extension(ext_node);
114        if let Some(mut rewriter) = self.planner.executor.logical_optimizer() {
115            plan = rewriter(plan)?;
116        }
117
118        Ok(Transformed::yes(plan))
119    }
120
121    /// A human readable name for this analyzer rule
122    fn name(&self) -> &str {
123        "federate_sql"
124    }
125
126    /// Does this rule support rewriting owned plans (rather than by reference)?
127    fn supports_rewrite(&self) -> bool {
128        true
129    }
130}
131
132#[derive(Debug)]
133pub struct SQLFederationPlanner {
134    pub executor: Arc<dyn SQLExecutor>,
135}
136
137impl SQLFederationPlanner {
138    pub fn new(executor: Arc<dyn SQLExecutor>) -> Self {
139        Self { executor }
140    }
141}
142
143#[async_trait]
144impl FederationPlanner for SQLFederationPlanner {
145    async fn plan_federation(
146        &self,
147        node: &FederatedPlanNode,
148        _session_state: &dyn Session,
149    ) -> Result<Arc<dyn ExecutionPlan>> {
150        let schema = Arc::new(node.plan().schema().as_arrow().clone());
151        let plan = node.plan().clone();
152        let statistics = self.executor.statistics(&plan).await?;
153        let input = Arc::new(VirtualExecutionPlan::new(
154            plan,
155            Arc::clone(&self.executor),
156            statistics,
157        ));
158        let schema_cast_exec = schema_cast::SchemaCastScanExec::new(input, schema);
159        Ok(Arc::new(schema_cast_exec))
160    }
161}
162
163#[derive(Debug, Clone)]
164pub struct VirtualExecutionPlan {
165    plan: LogicalPlan,
166    executor: Arc<dyn SQLExecutor>,
167    props: Arc<PlanProperties>,
168    statistics: Statistics,
169    filters: Vec<Arc<dyn PhysicalExpr>>,
170}
171
172impl VirtualExecutionPlan {
173    pub fn new(plan: LogicalPlan, executor: Arc<dyn SQLExecutor>, statistics: Statistics) -> Self {
174        let schema: Schema = plan.schema().as_arrow().clone();
175        let props = Arc::new(PlanProperties::new(
176            EquivalenceProperties::new(Arc::new(schema)),
177            Partitioning::UnknownPartitioning(1),
178            EmissionType::Incremental,
179            Boundedness::Bounded,
180        ));
181        Self {
182            plan,
183            executor,
184            props,
185            statistics,
186            filters: Vec::new(),
187        }
188    }
189
190    pub fn plan(&self) -> &LogicalPlan {
191        &self.plan
192    }
193
194    pub fn executor(&self) -> &Arc<dyn SQLExecutor> {
195        &self.executor
196    }
197
198    pub fn statistics(&self) -> &Statistics {
199        &self.statistics
200    }
201
202    fn schema(&self) -> SchemaRef {
203        let df_schema = self.plan.schema().as_arrow().clone();
204        Arc::new(df_schema)
205    }
206
207    fn final_sql(&self) -> Result<String> {
208        let plan = self.plan.clone();
209        let plan = RewriteTableScanAnalyzer::rewrite(plan)?;
210        let (logical_optimizers, ast_analyzers, sql_query_rewriters) = gather_analyzers(&plan)?;
211        let plan = apply_logical_optimizers(plan, logical_optimizers)?;
212        let ast = self.plan_to_statement(&plan)?;
213        let ast = self.rewrite_with_executor_ast_analyzer(ast)?;
214        let ast = apply_ast_analyzers(ast, ast_analyzers)?;
215        apply_sql_query_rewriters(ast.to_string(), sql_query_rewriters)
216    }
217
218    fn rewrite_with_executor_ast_analyzer(
219        &self,
220        ast: Statement,
221    ) -> Result<Statement, datafusion::error::DataFusionError> {
222        if let Some(mut analyzer) = self.executor.ast_analyzer() {
223            Ok(analyzer(ast)?)
224        } else {
225            Ok(ast)
226        }
227    }
228
229    fn plan_to_statement(&self, plan: &LogicalPlan) -> Result<Statement> {
230        Unparser::new(self.executor.dialect().as_ref()).plan_to_sql(plan)
231    }
232}
233
234fn gather_analyzers(
235    plan: &LogicalPlan,
236) -> Result<(
237    Vec<LogicalOptimizer>,
238    Vec<AstAnalyzer>,
239    Vec<SqlQueryRewriter>,
240)> {
241    let mut logical_optimizers = vec![];
242    let mut ast_analyzers = vec![];
243    let mut sql_query_rewriters = vec![];
244
245    plan.apply(|node| {
246        if let LogicalPlan::TableScan(table) = node {
247            let provider = get_table_source(&table.source)
248                .expect("caller is virtual exec so this is valid")
249                .expect("caller is virtual exec so this is valid");
250            if let Some(source) = (provider.as_ref() as &dyn Any).downcast_ref::<SQLTableSource>() {
251                if let Some(analyzer) = source.table.logical_optimizer() {
252                    logical_optimizers.push(analyzer);
253                }
254                if let Some(analyzer) = source.table.ast_analyzer() {
255                    ast_analyzers.push(analyzer);
256                }
257                if let Some(rewriter) = source.table.sql_query_rewriter() {
258                    sql_query_rewriters.push(rewriter);
259                }
260            }
261        }
262        Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue)
263    })?;
264
265    Ok((logical_optimizers, ast_analyzers, sql_query_rewriters))
266}
267
268fn apply_logical_optimizers(
269    mut plan: LogicalPlan,
270    analyzers: Vec<LogicalOptimizer>,
271) -> Result<LogicalPlan> {
272    for mut analyzer in analyzers {
273        let old_schema = plan.schema().clone();
274        plan = analyzer(plan)?;
275        let new_schema = plan.schema();
276        if &old_schema != new_schema {
277            return Err(DataFusionError::Execution(format!(
278                "Schema altered during logical analysis, expected: {}, found: {}",
279                old_schema, new_schema
280            )));
281        }
282    }
283    Ok(plan)
284}
285
286fn apply_ast_analyzers(mut statement: Statement, analyzers: Vec<AstAnalyzer>) -> Result<Statement> {
287    for mut analyzer in analyzers {
288        statement = analyzer(statement)?;
289    }
290    Ok(statement)
291}
292
293fn apply_sql_query_rewriters(
294    mut query: String,
295    rewriters: Vec<SqlQueryRewriter>,
296) -> Result<String> {
297    for mut rewriter in rewriters {
298        query = rewriter(query)?;
299    }
300    Ok(query)
301}
302
303impl DisplayAs for VirtualExecutionPlan {
304    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> std::fmt::Result {
305        write!(f, "VirtualExecutionPlan")?;
306        write!(f, " name={}", self.executor.name())?;
307        if let Some(ctx) = self.executor.compute_context() {
308            write!(f, " compute_context={ctx}")?;
309        };
310        let mut plan = match RewriteTableScanAnalyzer::rewrite(self.plan.clone()) {
311            Ok(plan) => plan,
312            Err(_) => self.plan.clone(),
313        };
314        if let Ok(statement) = self.plan_to_statement(&plan) {
315            write!(f, " base_sql={statement}")?;
316        }
317
318        let (logical_optimizers, ast_analyzers, sql_query_rewriters) = match gather_analyzers(&plan)
319        {
320            Ok(analyzers) => analyzers,
321            Err(_) => return Ok(()),
322        };
323
324        let old_plan = plan.clone();
325
326        plan = match apply_logical_optimizers(plan, logical_optimizers) {
327            Ok(plan) => plan,
328            _ => return Ok(()),
329        };
330
331        let statement = match self.plan_to_statement(&plan) {
332            Ok(statement) => statement,
333            _ => return Ok(()),
334        };
335
336        if plan != old_plan {
337            write!(f, " rewritten_logical_sql={statement}")?;
338        }
339
340        let old_statement = statement.clone();
341        let statement = match self.rewrite_with_executor_ast_analyzer(statement) {
342            Ok(statement) => statement,
343            _ => return Ok(()),
344        };
345        if old_statement != statement {
346            write!(f, " rewritten_executor_sql={statement}")?;
347        }
348
349        let old_statement = statement.clone();
350        let statement = match apply_ast_analyzers(statement, ast_analyzers) {
351            Ok(statement) => statement,
352            _ => return Ok(()),
353        };
354        if old_statement != statement {
355            write!(f, " rewritten_ast_analyzer={statement}")?;
356        }
357
358        let sql = statement.to_string();
359        let rewritten_sql = match apply_sql_query_rewriters(sql.clone(), sql_query_rewriters) {
360            Ok(sql) => sql,
361            _ => return Ok(()),
362        };
363        if sql != rewritten_sql {
364            write!(f, " rewritten_sql_query={rewritten_sql}")?;
365        }
366
367        Ok(())
368    }
369}
370
371impl ExecutionPlan for VirtualExecutionPlan {
372    fn name(&self) -> &str {
373        "sql_federation_exec"
374    }
375
376    fn schema(&self) -> SchemaRef {
377        self.schema()
378    }
379
380    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
381        vec![]
382    }
383
384    fn with_new_children(
385        self: Arc<Self>,
386        _: Vec<Arc<dyn ExecutionPlan>>,
387    ) -> Result<Arc<dyn ExecutionPlan>> {
388        Ok(self)
389    }
390
391    fn execute(
392        &self,
393        _partition: usize,
394        _context: Arc<TaskContext>,
395    ) -> Result<SendableRecordBatchStream> {
396        self.executor
397            .execute(&self.final_sql()?, self.schema(), &self.filters)
398    }
399
400    fn properties(&self) -> &Arc<PlanProperties> {
401        &self.props
402    }
403
404    fn statistics_from_inputs(
405        &self,
406        _input_stats: &[Arc<Statistics>],
407        _args: &StatisticsArgs,
408    ) -> Result<Arc<Statistics>> {
409        Ok(Arc::new(self.statistics.clone()))
410    }
411
412    fn metrics(&self) -> Option<MetricsSet> {
413        self.executor.metrics()
414    }
415
416    /// Passes the dynamic filters pushed down at runtime (e.g. by joins and TopK) to the
417    /// [`SQLExecutor`], which may use them to narrow down the query.
418    ///
419    /// The filters are never reported as pushed down: they are not part of the SQL query, and
420    /// the executor may ignore them. Static filters (those pushed down before the
421    /// [`FilterPushdownPhase::Post`] phase, such as the predicate of a `FilterExec`) are not
422    /// taken at all, so DataFusion keeps evaluating them.
423    fn handle_child_pushdown_result(
424        &self,
425        phase: FilterPushdownPhase,
426        child_pushdown_result: ChildPushdownResult,
427        _config: &ConfigOptions,
428    ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
429        if phase != FilterPushdownPhase::Post || child_pushdown_result.parent_filters.is_empty() {
430            return Ok(FilterPushdownPropagation::all_unsupported(
431                child_pushdown_result,
432            ));
433        }
434
435        let mut node = self.clone();
436        node.filters.extend(
437            child_pushdown_result
438                .parent_filters
439                .iter()
440                .map(|f| Arc::clone(&f.filter)),
441        );
442
443        Ok(
444            FilterPushdownPropagation::all_unsupported(child_pushdown_result)
445                .with_updated_node(Arc::new(node) as Arc<dyn ExecutionPlan>),
446        )
447    }
448
449    fn apply_expressions(
450        &self,
451        f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
452    ) -> Result<TreeNodeRecursion> {
453        apply_expression_roots(&self.filters, f)
454    }
455}
456
457#[cfg(test)]
458mod tests {
459    use std::any::Any;
460    use std::collections::HashSet;
461    use std::sync::atomic::{AtomicUsize, Ordering};
462    use std::sync::Arc;
463
464    use super::table::RemoteTable;
465    use super::*;
466    use crate::sql::{
467        RemoteTableRef, SQLExecutor, SQLFederationProvider, SQLTable, SQLTableSource,
468    };
469    use crate::FederatedTableProviderAdaptor;
470    use async_trait::async_trait;
471    use datafusion::arrow::datatypes::{Schema, SchemaRef};
472    use datafusion::common::tree_node::TreeNodeRecursion;
473    use datafusion::common::TableReference;
474    use datafusion::execution::SendableRecordBatchStream;
475    use datafusion::logical_expr::Operator;
476    use datafusion::physical_expr::expressions::{binary, col, lit};
477    use datafusion::physical_optimizer::{filter_pushdown::FilterPushdown, PhysicalOptimizerRule};
478    use datafusion::physical_plan::filter::FilterExec;
479    use datafusion::physical_plan::filter_pushdown::ChildFilterPushdownResult;
480    use datafusion::sql::unparser::dialect::Dialect;
481    use datafusion::sql::unparser::{self};
482    use datafusion::{
483        arrow::datatypes::{DataType, Field},
484        datasource::TableProvider,
485        execution::context::SessionContext,
486    };
487
488    #[derive(Debug, Clone)]
489    struct TestExecutor {
490        compute_context: String,
491    }
492
493    #[async_trait]
494    impl SQLExecutor for TestExecutor {
495        fn name(&self) -> &str {
496            "TestExecutor"
497        }
498
499        fn compute_context(&self) -> Option<String> {
500            Some(self.compute_context.clone())
501        }
502
503        fn dialect(&self) -> Arc<dyn Dialect> {
504            Arc::new(unparser::dialect::DefaultDialect {})
505        }
506
507        fn execute(
508            &self,
509            _query: &str,
510            _schema: SchemaRef,
511            _filters: &[Arc<dyn PhysicalExpr>],
512        ) -> Result<SendableRecordBatchStream> {
513            unimplemented!()
514        }
515
516        async fn table_names(&self) -> Result<Vec<String>> {
517            unimplemented!()
518        }
519
520        async fn get_table_schema(&self, _table_name: &str) -> Result<SchemaRef> {
521            unimplemented!()
522        }
523    }
524
525    fn get_test_table_provider(name: String, executor: TestExecutor) -> Arc<dyn TableProvider> {
526        let schema = Arc::new(Schema::new(vec![
527            Field::new("a", DataType::Int64, false),
528            Field::new("b", DataType::Utf8, false),
529            Field::new("c", DataType::Date32, false),
530        ]));
531        let table_ref = RemoteTableRef::try_from(name).unwrap();
532        let table = Arc::new(RemoteTable::new(table_ref, schema));
533        let provider = Arc::new(SQLFederationProvider::new(Arc::new(executor)));
534        let table_source = Arc::new(SQLTableSource { provider, table });
535        Arc::new(FederatedTableProviderAdaptor::new(table_source))
536    }
537
538    fn get_test_table_provider_with_table(
539        table: Arc<dyn SQLTable>,
540        executor: TestExecutor,
541    ) -> Arc<dyn TableProvider> {
542        let provider = Arc::new(SQLFederationProvider::new(Arc::new(executor)));
543        let table_source = Arc::new(SQLTableSource::new_with_table(provider, table));
544        Arc::new(FederatedTableProviderAdaptor::new(table_source))
545    }
546
547    #[derive(Debug)]
548    struct SqlRewriteTable {
549        table: RemoteTable,
550        rewrite_calls: Arc<AtomicUsize>,
551        suffix: String,
552    }
553
554    impl SqlRewriteTable {
555        fn new(
556            table_ref: RemoteTableRef,
557            schema: SchemaRef,
558            rewrite_calls: Arc<AtomicUsize>,
559            suffix: impl Into<String>,
560        ) -> Self {
561            Self {
562                table: RemoteTable::new(table_ref, schema),
563                rewrite_calls,
564                suffix: suffix.into(),
565            }
566        }
567    }
568
569    impl SQLTable for SqlRewriteTable {
570        fn as_any(&self) -> &dyn Any {
571            self
572        }
573
574        fn table_reference(&self) -> TableReference {
575            self.table.table_reference().clone()
576        }
577
578        fn schema(&self) -> SchemaRef {
579            Arc::clone(self.table.schema())
580        }
581
582        fn sql_query_rewriter(&self) -> Option<SqlQueryRewriter> {
583            let rewrite_calls = Arc::clone(&self.rewrite_calls);
584            let suffix = self.suffix.clone();
585            Some(Box::new(move |sql| {
586                rewrite_calls.fetch_add(1, Ordering::SeqCst);
587                Ok(format!("{sql} {suffix}"))
588            }))
589        }
590    }
591
592    #[tokio::test]
593    async fn basic_sql_federation_test() -> Result<(), DataFusionError> {
594        let test_executor_a = TestExecutor {
595            compute_context: "a".into(),
596        };
597
598        let test_executor_b = TestExecutor {
599            compute_context: "b".into(),
600        };
601
602        let table_a1_ref = "table_a1".to_string();
603        let table_a1 = get_test_table_provider(table_a1_ref.clone(), test_executor_a.clone());
604
605        let table_a2_ref = "table_a2".to_string();
606        let table_a2 = get_test_table_provider(table_a2_ref.clone(), test_executor_a);
607
608        let table_b1_ref = "table_b1(1)".to_string();
609        let table_b1_df_ref = "table_local_b1".to_string();
610
611        let table_b1 = get_test_table_provider(table_b1_ref.clone(), test_executor_b);
612
613        // Create a new SessionState with the optimizer rule we created above
614        let state = crate::default_session_state();
615        let ctx = SessionContext::new_with_state(state);
616
617        ctx.register_table(table_a1_ref.clone(), table_a1).unwrap();
618        ctx.register_table(table_a2_ref.clone(), table_a2).unwrap();
619        ctx.register_table(table_b1_df_ref.clone(), table_b1)
620            .unwrap();
621
622        let query = r#"
623            SELECT * FROM table_a1
624            UNION ALL
625            SELECT * FROM table_a2
626            UNION ALL
627            SELECT * FROM table_local_b1;
628        "#;
629
630        let df = ctx.sql(query).await?;
631
632        let logical_plan = df.into_optimized_plan()?;
633
634        let mut table_a1_federated = false;
635        let mut table_a2_federated = false;
636        let mut table_b1_federated = false;
637
638        let _ = logical_plan.apply(|node| {
639            if let LogicalPlan::Extension(node) = node {
640                if let Some(node) = node.node.as_any().downcast_ref::<FederatedPlanNode>() {
641                    let _ = node.plan().apply(|node| {
642                        if let LogicalPlan::TableScan(table) = node {
643                            if table.table_name.table() == table_a1_ref {
644                                table_a1_federated = true;
645                            }
646                            if table.table_name.table() == table_a2_ref {
647                                table_a2_federated = true;
648                            }
649                            // assuming table name is rewritten via analyzer
650                            if table.table_name.table() == table_b1_df_ref {
651                                table_b1_federated = true;
652                            }
653                        }
654                        Ok(TreeNodeRecursion::Continue)
655                    });
656                }
657            }
658            Ok(TreeNodeRecursion::Continue)
659        });
660
661        assert!(table_a1_federated);
662        assert!(table_a2_federated);
663        assert!(table_b1_federated);
664
665        let physical_plan = ctx.state().create_physical_plan(&logical_plan).await?;
666
667        let mut final_queries = vec![];
668
669        let _ = physical_plan.apply(|node| {
670            if node.name() == "sql_federation_exec" {
671                let node = (node.as_ref() as &dyn Any)
672                    .downcast_ref::<VirtualExecutionPlan>()
673                    .unwrap();
674
675                final_queries.push(node.final_sql()?);
676            }
677            Ok(TreeNodeRecursion::Continue)
678        });
679
680        let expected = vec![
681            "SELECT table_a1.a, table_a1.b, table_a1.c FROM table_a1",
682            "SELECT table_a2.a, table_a2.b, table_a2.c FROM table_a2",
683            "SELECT table_b1.a, table_b1.b, table_b1.c FROM table_b1(1) AS table_b1",
684        ];
685
686        assert_eq!(
687            HashSet::<&str>::from_iter(final_queries.iter().map(|x| x.as_str())),
688            HashSet::from_iter(expected)
689        );
690
691        Ok(())
692    }
693
694    #[tokio::test]
695    async fn multi_reference_sql_federation_test() -> Result<(), DataFusionError> {
696        let test_executor_a = TestExecutor {
697            compute_context: "test".into(),
698        };
699
700        let lowercase_table_ref = "default.table".to_string();
701        let lowercase_local_table_ref = "dftable".to_string();
702        let lowercase_table =
703            get_test_table_provider(lowercase_table_ref.clone(), test_executor_a.clone());
704
705        let capitalized_table_ref = "default.Table(1)".to_string();
706        let capitalized_local_table_ref = "dfview".to_string();
707        let capitalized_table =
708            get_test_table_provider(capitalized_table_ref.clone(), test_executor_a);
709
710        // Create a new SessionState with the optimizer rule we created above
711        let state = crate::default_session_state();
712        let ctx = SessionContext::new_with_state(state);
713
714        ctx.register_table(lowercase_local_table_ref.clone(), lowercase_table)
715            .unwrap();
716        ctx.register_table(capitalized_local_table_ref.clone(), capitalized_table)
717            .unwrap();
718
719        let query = r#"
720                SELECT * FROM dftable
721                UNION ALL
722                SELECT * FROM dfview;
723            "#;
724
725        let df = ctx.sql(query).await?;
726
727        let logical_plan = df.into_optimized_plan()?;
728
729        let mut lowercase_table = false;
730        let mut capitalized_table = false;
731
732        let _ = logical_plan.apply(|node| {
733            if let LogicalPlan::Extension(node) = node {
734                if let Some(node) = node.node.as_any().downcast_ref::<FederatedPlanNode>() {
735                    let _ = node.plan().apply(|node| {
736                        if let LogicalPlan::TableScan(table) = node {
737                            if table.table_name.table() == lowercase_local_table_ref {
738                                lowercase_table = true;
739                            }
740                            if table.table_name.table() == capitalized_local_table_ref {
741                                capitalized_table = true;
742                            }
743                        }
744                        Ok(TreeNodeRecursion::Continue)
745                    });
746                }
747            }
748            Ok(TreeNodeRecursion::Continue)
749        });
750
751        assert!(lowercase_table);
752        assert!(capitalized_table);
753
754        let physical_plan = ctx.state().create_physical_plan(&logical_plan).await?;
755
756        let mut final_queries = vec![];
757
758        let _ = physical_plan.apply(|node| {
759            if node.name() == "sql_federation_exec" {
760                let node = (node.as_ref() as &dyn Any)
761                    .downcast_ref::<VirtualExecutionPlan>()
762                    .unwrap();
763
764                final_queries.push(node.final_sql()?);
765            }
766            Ok(TreeNodeRecursion::Continue)
767        });
768
769        let expected = vec![
770            r#"SELECT "table".a, "table".b, "table".c FROM "default"."table" UNION ALL SELECT "Table".a, "Table".b, "Table".c FROM "default"."Table"(1) AS Table"#,
771        ];
772
773        assert_eq!(
774            HashSet::<&str>::from_iter(final_queries.iter().map(|x| x.as_str())),
775            HashSet::from_iter(expected)
776        );
777
778        Ok(())
779    }
780
781    /// EXPLAIN ANALYZE must not federate the Analyze wrapper — only the inner
782    /// query should be federated. Otherwise the SQL Unparser fails because it
783    /// cannot convert Analyze to SQL.
784    #[tokio::test]
785    async fn explain_analyze_not_federated() -> Result<(), DataFusionError> {
786        let executor = TestExecutor {
787            compute_context: "a".into(),
788        };
789
790        let table_ref = "test_table".to_string();
791        let table = get_test_table_provider(table_ref.clone(), executor);
792
793        let state = crate::default_session_state();
794        let ctx = SessionContext::new_with_state(state);
795        ctx.register_table(table_ref, table).unwrap();
796
797        let plan = ctx
798            .sql("EXPLAIN ANALYZE SELECT * FROM test_table")
799            .await?
800            .into_optimized_plan()?;
801
802        // The top-level node must be Analyze, not Federated.
803        assert!(
804            matches!(plan, LogicalPlan::Analyze(_)),
805            "Expected Analyze at root, got: {}",
806            plan.display_indent()
807        );
808
809        // The inner plan should contain a Federated extension node.
810        let mut found_federated = false;
811        plan.apply(|node| {
812            if let LogicalPlan::Extension(ext) = node {
813                if ext.node.name() == "Federated" {
814                    found_federated = true;
815                    return Ok(TreeNodeRecursion::Stop);
816                }
817            }
818            Ok(TreeNodeRecursion::Continue)
819        })?;
820        assert!(
821            found_federated,
822            "Expected a Federated node inside the Analyze plan"
823        );
824
825        // Physical planning should succeed (this is where it used to fail).
826        let physical_plan = ctx.state().create_physical_plan(&plan).await?;
827        assert_eq!(physical_plan.name(), "AnalyzeExec");
828
829        Ok(())
830    }
831
832    #[tokio::test]
833    async fn sql_query_rewriter_hook_invoked_and_rewrites_sql() -> Result<(), DataFusionError> {
834        let executor = TestExecutor {
835            compute_context: "rewrite".into(),
836        };
837        let rewrite_calls = Arc::new(AtomicUsize::new(0));
838        let table_ref = "table_with_rewriter".to_string();
839        let table = Arc::new(SqlRewriteTable::new(
840            table_ref.clone().try_into().unwrap(),
841            Arc::new(Schema::new(vec![
842                Field::new("a", DataType::Int64, false),
843                Field::new("b", DataType::Utf8, false),
844                Field::new("c", DataType::Date32, false),
845            ])),
846            Arc::clone(&rewrite_calls),
847            "/* rewritten by sql_query_rewriter */",
848        ));
849        let table_provider = get_test_table_provider_with_table(table, executor);
850
851        let state = crate::default_session_state();
852        let ctx = SessionContext::new_with_state(state);
853        ctx.register_table(table_ref.clone(), table_provider)
854            .unwrap();
855
856        let query = format!("SELECT * FROM {table_ref}");
857        let df = ctx.sql(&query).await?;
858        let logical_plan = df.into_optimized_plan()?;
859        let physical_plan = ctx.state().create_physical_plan(&logical_plan).await?;
860
861        let mut final_queries = vec![];
862        let _ = physical_plan.apply(|node| {
863            if node.name() == "sql_federation_exec" {
864                let node = (node.as_ref() as &dyn Any)
865                    .downcast_ref::<VirtualExecutionPlan>()
866                    .unwrap();
867                final_queries.push(node.final_sql()?);
868            }
869            Ok(TreeNodeRecursion::Continue)
870        });
871
872        let [final_query] = final_queries.as_slice() else {
873            panic!("expected a single federated SQL query");
874        };
875
876        assert!(final_query.ends_with("/* rewritten by sql_query_rewriter */"));
877        assert_eq!(rewrite_calls.load(Ordering::SeqCst), 1);
878
879        Ok(())
880    }
881
882    fn test_virtual_execution_plan() -> VirtualExecutionPlan {
883        let executor = TestExecutor {
884            compute_context: "a".into(),
885        };
886        let table = get_test_table_provider("table_a1".to_string(), executor.clone());
887        let plan = datafusion::logical_expr::LogicalPlanBuilder::scan(
888            "table_a1",
889            datafusion::datasource::provider_as_source(table),
890            None,
891        )
892        .unwrap()
893        .build()
894        .unwrap();
895        let statistics = Statistics::new_unknown(plan.schema().as_arrow());
896        VirtualExecutionPlan::new(plan, Arc::new(executor), statistics)
897    }
898
899    fn a_greater_than(schema: &Schema, value: i64) -> Arc<dyn PhysicalExpr> {
900        binary(col("a", schema).unwrap(), Operator::Gt, lit(value), schema).unwrap()
901    }
902
903    /// The filters pushed down into a `VirtualExecutionPlan` are not part of its SQL query, and
904    /// a `SQLExecutor` may ignore them, so a `FilterExec` on top of it must be kept.
905    #[test]
906    fn filter_exec_above_virtual_execution_plan_is_kept() {
907        let scan: Arc<dyn ExecutionPlan> = Arc::new(test_virtual_execution_plan());
908        let predicate = a_greater_than(&scan.schema(), 1);
909        let plan = Arc::new(FilterExec::try_new(predicate, scan).unwrap());
910
911        let optimized = FilterPushdown::new()
912            .optimize(plan, &ConfigOptions::default())
913            .unwrap();
914        assert!(
915            optimized.downcast_ref::<FilterExec>().is_some(),
916            "the FilterExec should be kept:\n{}",
917            datafusion::physical_plan::displayable(optimized.as_ref()).indent(true)
918        );
919    }
920
921    /// Filters pushed down in separate pushdowns, e.g. the dynamic filters of a join and of a
922    /// TopK above it, are all passed to the executor.
923    #[test]
924    fn filters_pushed_down_separately_are_all_kept() {
925        let push = |plan: &VirtualExecutionPlan, filter: &Arc<dyn PhysicalExpr>| {
926            let result = plan
927                .handle_child_pushdown_result(
928                    FilterPushdownPhase::Post,
929                    ChildPushdownResult {
930                        parent_filters: vec![ChildFilterPushdownResult {
931                            filter: Arc::clone(filter),
932                            child_results: vec![],
933                        }],
934                        self_filters: vec![],
935                    },
936                    &ConfigOptions::default(),
937                )
938                .unwrap();
939            result
940                .updated_node
941                .expect("the filter should be added to the node")
942                .downcast_ref::<VirtualExecutionPlan>()
943                .unwrap()
944                .clone()
945        };
946
947        let plan = test_virtual_execution_plan();
948        let first = a_greater_than(&plan.schema(), 1);
949        let second = a_greater_than(&plan.schema(), 2);
950        let plan = push(&push(&plan, &first), &second);
951
952        let filters: Vec<String> = plan.filters.iter().map(|f| f.to_string()).collect();
953        assert_eq!(filters, vec![first.to_string(), second.to_string()]);
954    }
955}
datafusion_federation/sql/mod.rs

datafusion_federation/sql/
mod.rs
