Table provider factory for Arrow Flight and Flight SQL (#76)

Co-authored-by: Phillip LeBlanc <[email protected]>

Author: ccciudatu

.gitignore

@@ -17,4 +17,6 @@ Cargo.lock
 
 .vscode/settings.json
 
-.DS_Store
+.DS_Store
+
+.idea

Cargo.toml

@@ -6,16 +6,25 @@ edition = "2021"
 repository = "https://github.com/datafusion-contrib/datafusion-table-providers"
 
 [dependencies]
-arrow = "52.0.0"
+arrow = "52.2.0"
+arrow-array = { version = "52.2.0", optional = true }
+arrow-flight = { version = "52.2.0", optional = true, features = ["flight-sql-experimental", "tls"] }
+arrow-schema = { version = "52.2.0", optional = true, features = ["serde"] }
 arrow-json = "52.2.0"
 async-stream = { version = "0.3.5", optional = true }
 async-trait = "0.1.80"
 num-bigint = "0.4.4"
+base64 = { version = "0.22.1", optional = true }
+bytes = { version = "1.7.1", optional = true }
 bigdecimal = "0.4.5"
 bigdecimal_0_3_0 = { package = "bigdecimal", version = "0.3.0" }
 byteorder = "1.5.0"
 chrono = "0.4.38"
 datafusion = "41.0.0"
+datafusion-expr = { version = "41.0.0", optional = true }
+datafusion-physical-expr = { version = "41.0.0", optional = true }
+datafusion-physical-plan = { version = "41.0.0", optional = true }
+datafusion-proto = { version = "41.0.0", optional = true }
 duckdb = { version = "1", features = [
   "bundled",
   "r2d2",
@@ -26,10 +35,12 @@ duckdb = { version = "1", features = [
 fallible-iterator = "0.3.0"
 futures = "0.3.30"
 mysql_async = { version = "0.34.1", features = ["native-tls-tls", "chrono"], optional = true }
+prost = { version = "0.12" , optional = true } # pinned for arrow-flight compat
 r2d2 = { version = "0.8.10", optional = true }
 rusqlite = { version = "0.31.0", optional = true }
 sea-query = { version = "0.31.0", features = ["backend-sqlite", "backend-postgres", "postgres-array", "with-rust_decimal", "with-bigdecimal", "with-time", "with-chrono"] }
 secrecy = "0.8.0"
+serde = { version = "1.0.209", optional = true }
 serde_json = "1.0.124"
 snafu = "0.8.3"
 time = "0.3.36"
@@ -45,6 +56,7 @@ trust-dns-resolver = "0.23.2"
 url = "2.5.1"
 pem = { version = "3.0.4", optional = true }
 tokio-rusqlite = { version = "0.5.1", optional = true }
+tonic = { version = "0.11", optional = true } # pinned for arrow-flight compat
 datafusion-federation = "0.1"
 datafusion-federation-sql = { git = "https://github.com/spiceai/datafusion-federation.git", rev = "21f07bec7284bcbff2bf4e570008290b66e3dc6f" }
 itertools = "0.13.0"
@@ -60,12 +72,27 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
 test-log = { version = "0.2.16", features = ["trace"] }
 rstest = "0.22.0"
 geozero = { version = "0.13.0", features = ["with-wkb"] }
+tokio-stream = { version = "0.1.15", features = ["net"] }
 
 [features]
 mysql = ["dep:mysql_async", "dep:async-stream"]
 postgres = ["dep:tokio-postgres", "dep:uuid", "dep:postgres-native-tls", "dep:bb8", "dep:bb8-postgres", "dep:native-tls", "dep:pem", "dep:async-stream"]
 sqlite = ["dep:rusqlite", "dep:tokio-rusqlite"]
 duckdb = ["dep:duckdb", "dep:r2d2", "dep:uuid"]
+flight = [
+  "dep:arrow-array",
+  "dep:arrow-flight",
+  "dep:arrow-schema",
+  "dep:base64",
+  "dep:bytes",
+  "dep:datafusion-expr",
+  "dep:datafusion-physical-expr",
+  "dep:datafusion-physical-plan",
+  "dep:datafusion-proto",
+  "dep:prost",
+  "dep:serde",
+  "dep:tonic",
+]
 duckdb-federation = ["duckdb"]
 sqlite-federation = ["sqlite"]
 

README.md

@@ -12,6 +12,7 @@ Many of the table providers in this repo are for querying data from other databa
 - MySQL
 - SQLite
 - DuckDB
+- Flight SQL
 
 ## Examples
 
@@ -79,3 +80,13 @@ EOF
 ```bash
 cargo run --example mysql --features mysql
 ```
+
+### Flight SQL
+```bash
+brew install roapi
+# or
+#cargo install --locked --git https://github.com/roapi/roapi --branch main --bins roapi
+roapi -t taxi=https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-01.parquet &
+
+cargo run --example flight-sql --features flight
+```

examples/flight-sql.rs

@@ -0,0 +1,74 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::prelude::SessionContext;
+use datafusion_table_providers::flight::sql::{FlightSqlDriver, QUERY};
+use datafusion_table_providers::flight::FlightTableFactory;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+/// Prerequisites:
+/// ```
+/// $ brew install roapi
+/// $ roapi -t taxi=https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2024-01.parquet
+/// ```
+#[tokio::main]
+async fn main() -> datafusion::common::Result<()> {
+    let ctx = SessionContext::new();
+    let flight_sql = FlightTableFactory::new(Arc::new(FlightSqlDriver::default()));
+    let table = flight_sql
+        .open_table(
+            "http://localhost:32010",
+            HashMap::from([(QUERY.into(), "SELECT * FROM taxi".into())]),
+        )
+        .await?;
+    ctx.register_table("trip_data", Arc::new(table))?;
+    ctx.sql("select * from trip_data limit 10")
+        .await?
+        .show()
+        .await?;
+
+    // The same table created using DDL
+    ctx.state_ref()
+        .write()
+        .table_factories_mut()
+        .insert("FLIGHT_SQL".into(), Arc::new(flight_sql));
+    let _ = ctx
+        .sql(
+            r#"
+            CREATE EXTERNAL TABLE trip_data2 STORED AS FLIGHT_SQL
+            LOCATION 'http://localhost:32010'
+            OPTIONS (
+                'flight.sql.query' 'SELECT * FROM taxi'
+            )
+        "#,
+        )
+        .await?;
+
+    let df = ctx
+        .sql(
+            r#"
+            SELECT "VendorID", COUNT(*), SUM(passenger_count), SUM(total_amount)
+            FROM trip_data2
+            GROUP BY "VendorID"
+            ORDER BY COUNT(*) DESC
+        "#,
+        )
+        .await?;
+    df.clone().explain(false, false)?.show().await?;
+    df.show().await
+}