feat: Upstream Spice AI changes to datafusion-federation (#51)

Co-authored-by: sgrebnov <[email protected]> Co-authored-by: yfu <[email protected]> Co-authored-by: Qianqian <[email protected]> Co-authored-by: Michiel De Backker <[email protected]> Co-authored-by: Suriya Kandaswamy <[email protected]> Co-authored-by: Suriya Kandaswamy <[email protected]> Co-authored-by: hozan23 <[email protected]> Co-authored-by: hozan23 <[email protected]>
datafusion-contrib · Sep 2, 2024 · af09abe · af09abe
1 parent c8fa466
commit af09abe
Show file tree

Hide file tree

Showing 17 changed files with 2,264 additions and 31 deletions.
diff --git a/.gitignore b/.gitignore
@@ -3,3 +3,4 @@
 /node_modules
 package-lock.json
 package.json
+.DS_Store
diff --git a/.vscode/launch.json b/.vscode/launch.json
@@ -0,0 +1,213 @@
+{
+  // Use IntelliSense to learn about possible attributes.
+  // Hover to view descriptions of existing attributes.
+  // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in library 'datafusion_federation'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--lib",
+          "--package=datafusion-federation"
+        ],
+        "filter": {
+          "name": "datafusion_federation",
+          "kind": "lib"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug example 'sqlite'",
+      "cargo": {
+        "args": [
+          "build",
+          "--example=sqlite",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "sqlite",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in example 'sqlite'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--example=sqlite",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "sqlite",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug example 'flight-sql'",
+      "cargo": {
+        "args": [
+          "build",
+          "--example=flight-sql",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "flight-sql",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in example 'flight-sql'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--example=flight-sql",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "flight-sql",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug example 'postgres-partial'",
+      "cargo": {
+        "args": [
+          "build",
+          "--example=postgres-partial",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "postgres-partial",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in example 'postgres-partial'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--example=postgres-partial",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "postgres-partial",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug example 'sqlite-partial'",
+      "cargo": {
+        "args": [
+          "build",
+          "--example=sqlite-partial",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "sqlite-partial",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in example 'sqlite-partial'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--example=sqlite-partial",
+          "--package=datafusion-federation-examples"
+        ],
+        "filter": {
+          "name": "sqlite-partial",
+          "kind": "example"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in library 'datafusion_federation_flight_sql'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--lib",
+          "--package=datafusion-federation-flight-sql"
+        ],
+        "filter": {
+          "name": "datafusion_federation_flight_sql",
+          "kind": "lib"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    },
+    {
+      "type": "lldb",
+      "request": "launch",
+      "name": "Debug unit tests in library 'datafusion_federation_sql'",
+      "cargo": {
+        "args": [
+          "test",
+          "--no-run",
+          "--lib",
+          "--package=datafusion-federation-sql"
+        ],
+        "filter": {
+          "name": "datafusion_federation_sql",
+          "kind": "lib"
+        }
+      },
+      "args": [],
+      "cwd": "${workspaceFolder}"
+    }
+  ]
+}
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,12 +8,15 @@ members = [
 ]
 
 [workspace.package]
-version = "0.1.3"
+version = "0.1.6"
 edition = "2021"
-license = "MIT"
+license = "Apache-2.0"
 readme = "README.md"
 
 [workspace.dependencies]
 async-trait = "0.1.81"
+async-stream = "0.3.5"
+futures = "0.3.30"
 datafusion = "41.0.0"
 datafusion-substrait = "41.0.0"
+arrow-json = "52.2.0"
diff --git a/datafusion-federation/Cargo.toml b/datafusion-federation/Cargo.toml
@@ -17,16 +17,20 @@ all-features = true
 no-default-features = true
 
 [features]
-sql = ["futures"]
+sql = []
 
 [dependencies]
+futures.workspace = true
 async-trait.workspace = true
 datafusion.workspace = true
+async-stream.workspace = true
+arrow-json.workspace = true
 
-futures = { version = "0.3.30", optional = true }
 
 [dev-dependencies]
 tokio = { version = "1.39.3", features = ["full"] }
+tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+tracing = "0.1.40"
 
 [[example]]
 name = "df-csv"

diff --git a/datafusion-federation/examples/df-csv.rs b/datafusion-federation/examples/df-csv.rs
@@ -10,7 +10,7 @@ use datafusion::{
         options::CsvReadOptions,
     },
     physical_plan::{stream::RecordBatchStreamAdapter, SendableRecordBatchStream},
-    sql::sqlparser::dialect::{Dialect, GenericDialect},
+    sql::unparser::dialect::{DefaultDialect, Dialect},
 };
 use datafusion_federation::sql::{SQLExecutor, SQLFederationProvider, SQLSchemaProvider};
 use futures::TryStreamExt;
@@ -106,7 +106,7 @@ impl SQLExecutor for InMemorySQLExecutor {
     }
 
     fn dialect(&self) -> Arc<dyn Dialect> {
-        Arc::new(GenericDialect {})
+        Arc::new(DefaultDialect {})
     }
 }
 

diff --git a/datafusion-federation/src/lib.rs b/datafusion-federation/src/lib.rs
@@ -16,9 +16,15 @@ use datafusion::{
 };
 
 pub use optimizer::{get_table_source, FederationOptimizerRule};
-pub use plan_node::{FederatedPlanNode, FederatedQueryPlanner, FederationPlanner};
+pub use plan_node::{
+    FederatedPlanNode, FederatedPlanner, FederatedQueryPlanner, FederationPlanner,
+};
 pub use table_provider::{FederatedTableProviderAdaptor, FederatedTableSource};
 
+// TODO clean up this
+// TODO move schema_cast.rs to schema_cast directory
+pub mod schema_cast;
+
 pub fn default_session_state() -> SessionState {
     let rules = default_optimizer_rules();
     SessionStateBuilder::new()

diff --git a/datafusion-federation/src/plan_node.rs b/datafusion-federation/src/plan_node.rs
@@ -54,7 +54,7 @@ impl UserDefinedLogicalNodeCore for FederatedPlanNode {
     }
 
     fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "Federated\n {:?}", self.plan)
+        write!(f, "Federated\n {}", self.plan)
     }
 
     fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
@@ -126,7 +126,7 @@ impl Hash for FederatedPlanNode {
 }
 
 #[derive(Default)]
-struct FederatedPlanner {}
+pub struct FederatedPlanner {}
 
 impl FederatedPlanner {
     pub fn new() -> Self {
@@ -152,7 +152,7 @@ impl ExtensionPlanner for FederatedPlanner {
                 ));
             }
 
-            let fed_planner = fed_node.planner.clone();
+            let fed_planner = Arc::clone(&fed_node.planner);
             let exec_plan = fed_planner.plan_federation(fed_node, session_state).await?;
             return Ok(Some(exec_plan));
         }