wip

Author: phillipleblanc

Cargo.toml

@@ -21,6 +21,6 @@ datafusion-federation = { path = "./datafusion-federation", version = "0.4.2" }
 futures = "0.3.31"
 tokio = { version = "1.41", features = ["full"] }
 
-[patch.crates-io]
-duckdb = { git = "https://github.com/spiceai/duckdb-rs.git", rev = "2e24b958e44ec7419290249e27a15f1a19703fff" }
+#[patch.crates-io]
+#duckdb = { git = "https://github.com/spiceai/duckdb-rs.git", rev = "2e24b958e44ec7419290249e27a15f1a19703fff" }
 

datafusion-federation/src/lib.rs

@@ -5,7 +5,6 @@ pub mod schema_cast;
 #[cfg(feature = "sql")]
 pub mod sql;
 mod table_provider;
-pub mod table_reference;
 
 use std::{
     fmt,

datafusion-federation/src/optimize.rs

@@ -120,5 +120,5 @@ fn optimize_plan_node(
         return Ok(Transformed::no(plan));
     }
 
-    return rule.rewrite(plan, config);
+    rule.rewrite(plan, config)
 }

datafusion-federation/src/sql/analyzer.rs

@@ -1,7 +1,7 @@
 use std::{collections::HashMap, sync::Arc};
 
 use datafusion::{
-    common::{Column, Spans},
+    common::Column,
     logical_expr::{
         expr::{
             AggregateFunction, AggregateFunctionParams, Alias, Exists, InList, InSubquery,
@@ -16,7 +16,7 @@ use datafusion::{
 
 use crate::get_table_source;
 
-use super::SQLTableSource;
+use super::{table_reference::MultiPartTableReference, SQLTableSource};
 
 type Result<T> = std::result::Result<T, datafusion::error::DataFusionError>;
 
@@ -34,7 +34,7 @@ impl RewriteTableScanAnalyzer {
 /// Rewrite table scans to use the original federated table name.
 fn rewrite_table_scans(
     plan: &LogicalPlan,
-    known_rewrites: &mut HashMap<TableReference, TableReference>,
+    known_rewrites: &mut HashMap<MultiPartTableReference, TableReference>,
 ) -> Result<LogicalPlan> {
     if plan.inputs().is_empty() {
         if let LogicalPlan::TableScan(table_scan) = plan {
@@ -192,7 +192,6 @@ fn rewrite_table_scans_in_expr(
             Ok(Expr::ScalarSubquery(Subquery {
                 subquery: Arc::new(new_subquery),
                 outer_ref_columns,
-                spans: Spans::new(),
             }))
         }
         Expr::BinaryExpr(binary_expr) => {
@@ -466,7 +465,6 @@ fn rewrite_table_scans_in_expr(
             let subquery = Subquery {
                 subquery: Arc::new(subquery_plan),
                 outer_ref_columns,
-                spans: Spans::new(),
             };
             Ok(Expr::Exists(Exists::new(subquery, exists.negated)))
         }
@@ -482,7 +480,6 @@ fn rewrite_table_scans_in_expr(
             let subquery = Subquery {
                 subquery: Arc::new(subquery_plan),
                 outer_ref_columns,
-                spans: Spans::new(),
             };
             Ok(Expr::InSubquery(InSubquery::new(
                 Box::new(expr),
@@ -570,6 +567,7 @@ fn rewrite_table_scans_in_expr(
 #[cfg(test)]
 mod tests {
     use crate::sql::table::SQLTable;
+    use crate::sql::table_reference::MultiPartTableReference;
     use crate::sql::{RemoteTableRef, SQLExecutor, SQLFederationProvider, SQLTableSource};
     use crate::FederatedTableProviderAdaptor;
     use async_trait::async_trait;
@@ -634,8 +632,8 @@ mod tests {
     }
 
     impl SQLTable for TestTable {
-        fn table_reference(&self) -> TableReference {
-            TableReference::from(&self.name)
+        fn table_reference(&self) -> MultiPartTableReference {
+            MultiPartTableReference::from(&self.name)
         }
 
         fn schema(&self) -> datafusion::arrow::datatypes::SchemaRef {

datafusion-federation/src/sql/ast_analyzer.rs

@@ -1,14 +1,10 @@
 use std::ops::ControlFlow;
 
-use datafusion::sql::{
-    sqlparser::ast::{
-        FunctionArg, Ident, ObjectName, Statement, TableAlias, TableFactor, TableFunctionArgs,
-        VisitMut, VisitorMut,
-    },
-    TableReference,
+use datafusion::sql::sqlparser::ast::{
+    FunctionArg, Ident, Statement, TableAlias, TableFactor, TableFunctionArgs, VisitMut, VisitorMut,
 };
 
-use super::AstAnalyzer;
+use super::{table_reference::MultiPartTableReference, AstAnalyzer};
 
 pub fn replace_table_args_analyzer(mut visitor: TableArgReplace) -> AstAnalyzer {
     let x = move |mut statement: Statement| {
@@ -38,12 +34,12 @@ pub fn replace_table_args_analyzer(mut visitor: TableArgReplace) -> AstAnalyzer
 /// ```
 #[derive(Debug, Clone, PartialEq, Eq, Default)]
 pub struct TableArgReplace {
-    pub tables: Vec<(TableReference, TableFunctionArgs)>,
+    pub tables: Vec<(MultiPartTableReference, TableFunctionArgs)>,
 }
 
 impl TableArgReplace {
     /// Constructs a new `TableArgReplace` instance.
-    pub fn new(tables: Vec<(TableReference, Vec<FunctionArg>)>) -> Self {
+    pub fn new(tables: Vec<(MultiPartTableReference, Vec<FunctionArg>)>) -> Self {
         Self {
             tables: tables
                 .into_iter()
@@ -61,7 +57,7 @@ impl TableArgReplace {
     }
 
     /// Adds a new table argument replacement.
-    pub fn with(mut self, table: TableReference, args: Vec<FunctionArg>) -> Self {
+    pub fn with(mut self, table: MultiPartTableReference, args: Vec<FunctionArg>) -> Self {
         self.tables.push((
             table,
             TableFunctionArgs {
@@ -88,7 +84,8 @@ impl VisitorMut for TableArgReplace {
             name, args, alias, ..
         } = table_factor
         {
-            let name_as_tableref = name_to_table_reference(name);
+            let name = &*name;
+            let name_as_tableref = name.into();
             if let Some((table, arg)) = self
                 .tables
                 .iter()
@@ -106,25 +103,3 @@ impl VisitorMut for TableArgReplace {
         ControlFlow::Continue(())
     }
 }
-
-fn name_to_table_reference(name: &ObjectName) -> TableReference {
-    let first = name
-        .0
-        .first()
-        .map(|n| n.as_ident().expect("expected Ident").value.to_string());
-    let second = name
-        .0
-        .get(1)
-        .map(|n| n.as_ident().expect("expected Ident").value.to_string());
-    let third = name
-        .0
-        .get(2)
-        .map(|n| n.as_ident().expect("expected Ident").value.to_string());
-
-    match (first, second, third) {
-        (Some(first), Some(second), Some(third)) => TableReference::full(first, second, third),
-        (Some(first), Some(second), None) => TableReference::partial(first, second),
-        (Some(first), None, None) => TableReference::bare(first),
-        _ => panic!("Invalid table name"),
-    }
-}

datafusion-federation/src/sql/mod.rs

@@ -12,10 +12,7 @@ use analyzer::RewriteTableScanAnalyzer;
 use async_trait::async_trait;
 use datafusion::{
     arrow::datatypes::{Schema, SchemaRef},
-    common::{
-        tree_node::{Transformed, TreeNode},
-        HashMap,
-    },
+    common::tree_node::TreeNode,
     config::ConfigOptions,
     error::{DataFusionError, Result},
     execution::{context::SessionState, TaskContext},
@@ -36,9 +33,7 @@ pub use table::{RemoteTable, SQLTableSource};
 pub use table_reference::RemoteTableRef;
 
 use crate::{
-    schema_cast,
-    table_reference::{MultiPartTableReference, MultiTableReference},
-    FederatedPlanNode, FederationPlanner, FederationProvider,
+    get_table_source, schema_cast, FederatedPlanNode, FederationPlanner, FederationProvider,
 };
 
 // SQLFederationProvider provides federation to SQL DMBSs.
@@ -106,7 +101,7 @@ impl AnalyzerRule for SQLFederationAnalyzerRule {
             plan = rewriter(plan)?;
         }
 
-        Ok(Transformed::yes(plan))
+        Ok(plan)
     }
 
     /// A human readable name for this analyzer rule
@@ -172,36 +167,32 @@ impl VirtualExecutionPlan {
     }
 
     // TODO merge
-    fn sql(&self) -> Result<String> {
-        // Find all table scans, recover the SQLTableSource, find the remote table name and replace the name of the TableScan table.
-        let mut known_rewrites = HashMap::new();
-        let subquery_uses_partial_path = self.executor.subquery_use_partial_path();
-        let rewritten_plan = rewrite::plan::rewrite_table_scans(
-            &self.plan,
-            &mut known_rewrites,
-            subquery_uses_partial_path,
-            &mut None,
-        )?;
-        let mut ast = self.plan_to_sql(&rewritten_plan)?;
-
-        // If there are any MultiPartTableReference, rewrite the AST to use the original table names.
-        let multi_table_reference_rewrites = known_rewrites
-            .into_iter()
-            .filter_map(|(table_ref, rewrite)| match rewrite {
-                MultiPartTableReference::Multi(rewrite) => Some((table_ref, rewrite)),
-                _ => None,
-            })
-            .collect::<HashMap<TableReference, MultiTableReference>>();
-        if !multi_table_reference_rewrites.is_empty() {
-            rewrite::ast::rewrite_multi_part_statement(&mut ast, &multi_table_reference_rewrites);
-        }
-
-        if let Some(analyzer) = self.executor.ast_analyzer() {
-            ast = analyzer(ast)?;
-        }
-
-        Ok(format!("{ast}"))
-    }
+    // fn sql(&self) -> Result<String> {
+    //     // Find all table scans, recover the SQLTableSource, find the remote table name and replace the name of the TableScan table.
+    //     let mut known_rewrites = HashMap::new();
+    //     let subquery_uses_partial_path = self.executor.subquery_use_partial_path();
+    //     let rewritten_plan = rewrite::plan::rewrite_table_scans(
+    //         &self.plan,
+    //         &mut known_rewrites,
+    //         subquery_uses_partial_path,
+    //         &mut None,
+    //     )?;
+    //     let mut ast = self.plan_to_statement(&rewritten_plan)?;
+
+    //     // If there are any MultiPartTableReference, rewrite the AST to use the original table names.
+    //     let multi_table_reference_rewrites = known_rewrites
+    //         .into_iter()
+    //         .filter_map(|(table_ref, rewrite)| match rewrite {
+    //             MultiPartTableReference::Multi(rewrite) => Some((table_ref, rewrite)),
+    //             _ => None,
+    //         })
+    //         .collect::<HashMap<TableReference, MultiTableReference>>();
+    //     if !multi_table_reference_rewrites.is_empty() {
+    //         rewrite::ast::rewrite_multi_part_statement(&mut ast, &multi_table_reference_rewrites);
+    //     }
+
+    //     Ok(format!("{ast}"))
+    // }
 
     fn final_sql(&self) -> Result<String> {
         let plan = self.plan.clone();

datafusion-federation/src/sql/rewrite/ast.rs

@@ -14,6 +14,7 @@ use datafusion::{
 /// Rewrites table references in a SQL AST to use the original federated table names.
 /// This is similar to rewrite_table_scans but operates on the sqlparser AST instead
 /// of DataFusion logical plans.
+#[allow(dead_code)]
 pub(crate) fn rewrite_multi_part_statement(
     statement: &mut ast::Statement,
     known_rewrites: &HashMap<TableReference, MultiTableReference>,
@@ -27,6 +28,7 @@ pub(crate) fn rewrite_multi_part_statement(
     }
 }
 
+#[allow(dead_code)]
 fn rewrite_multi_part_table_with_joins(
     table_with_joins: &mut Vec<TableWithJoins>,
     known_rewrites: &HashMap<ObjectName, MultiTableReference>,

datafusion-federation/src/sql/rewrite/plan.rs

@@ -176,6 +176,7 @@ fn collect_known_rewrites_from_expr(
     }
 }
 
+#[allow(dead_code)]
 /// Rewrite table scans to use the original federated table name.
 pub(crate) fn rewrite_table_scans(
     plan: &LogicalPlan,
@@ -558,6 +559,7 @@ fn rewrite_column_name_in_expr(
     }
 }
 
+#[allow(dead_code)]
 fn rewrite_table_scans_in_expr(
     expr: Expr,
     known_rewrites: &HashMap<TableReference, MultiPartTableReference>,
@@ -1249,7 +1251,7 @@ fn rewrite_table_scans_in_expr(
 
 #[cfg(test)]
 mod tests {
-    use crate::FederatedTableProviderAdaptor;
+    use crate::{sql::RemoteTableRef, FederatedTableProviderAdaptor};
     use async_trait::async_trait;
     use datafusion::{
         arrow::datatypes::{DataType, Field, Schema, SchemaRef},
@@ -1318,14 +1320,11 @@ mod tests {
                 false,
             ),
         ]));
-        let table_source = Arc::new(
-            SQLTableSource::new_with_schema(
-                sql_federation_provider,
-                "remote_table".to_string(),
-                schema,
-            )
-            .expect("to have a valid SQLTableSource"),
-        );
+        let table_source = Arc::new(SQLTableSource::new_with_schema(
+            sql_federation_provider,
+            RemoteTableRef::try_from("remote_table").expect("valid table ref"),
+            schema,
+        ));
         Arc::new(FederatedTableProviderAdaptor::new(table_source))
     }
 
@@ -1342,14 +1341,12 @@ mod tests {
                 false,
             ),
         ]));
-        let table_source = Arc::new(
-            SQLTableSource::new_with_schema(
-                sql_federation_provider,
-                "remote_db.remote_schema.remote_table".to_string(),
-                schema,
-            )
-            .expect("to have a valid SQLTableSource"),
-        );
+        let table_source = Arc::new(SQLTableSource::new_with_schema(
+            sql_federation_provider,
+            RemoteTableRef::try_from("remote_db.remote_schema.remote_table")
+                .expect("valid table ref"),
+            schema,
+        ));
         Arc::new(FederatedTableProviderAdaptor::new(table_source))
     }
 
@@ -1679,7 +1676,7 @@ mod tests {
 
 #[cfg(test)]
 mod collect_rewrites_tests {
-    use crate::sql::{SQLExecutor, SQLFederationProvider, SQLTableSource};
+    use crate::sql::{RemoteTableRef, SQLExecutor, SQLFederationProvider, SQLTableSource};
 
     use super::*;
     use crate::FederatedTableProviderAdaptor;
@@ -1735,14 +1732,11 @@ mod collect_rewrites_tests {
 
         let sql_federation_provider =
             Arc::new(SQLFederationProvider::new(Arc::new(TestSQLExecutor {})));
-        let table_source = Arc::new(
-            SQLTableSource::new_with_schema(
-                sql_federation_provider,
-                "remote_table".to_string(),
-                schema.clone(),
-            )
-            .expect("to have a valid SQLTableSource"),
-        );
+        let table_source = Arc::new(SQLTableSource::new_with_schema(
+            sql_federation_provider,
+            RemoteTableRef::try_from("remote_table").expect("valid table ref"),
+            schema.clone(),
+        ));
         let source = Arc::new(DefaultTableSource::new(Arc::new(
             FederatedTableProviderAdaptor::new(table_source),
         )));