Enable flight

Author: crystalxyz

Cargo.lock

@@ -111,7 +111,7 @@ brew install roapi
 # 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
+cargo run -p datafusion-table-providers --example flight-sql --no-default-features --features flight
 ```
 
 ### ODBC

README.md

@@ -24,6 +24,7 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use crate::flight::{flight_channel, to_df_err, FlightMetadata, FlightProperties, SizeLimits};
+use crate::sql::db_connection_pool::runtime::run_async_with_tokio;
 use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_client::FlightServiceClient;
 use arrow_flight::{FlightClient, FlightEndpoint, Ticket};
@@ -190,7 +191,9 @@ async fn flight_stream(
 ) -> Result<SendableRecordBatchStream> {
     let mut errors: Vec<Box<dyn Error + Send + Sync>> = vec![];
     for loc in partition.locations.iter() {
-        let client = flight_client(loc, grpc_headers.as_ref(), &size_limits).await?;
+        let get_client = || async { flight_client(loc, grpc_headers.as_ref(), &size_limits).await };
+        let client = run_async_with_tokio(get_client).await?;
+        // let client = flight_client(loc, grpc_headers.as_ref(), &size_limits).await?;
         match try_fetch_stream(client, &partition.ticket, schema.clone()).await {
             Ok(stream) => return Ok(stream),
             Err(e) => errors.push(Box::new(e)),

core/src/flight/exec.rs

@@ -15,9 +15,10 @@ doc = false
 
 [dependencies]
 arrow = { workspace = true }
+arrow-flight = {workspace = true}
 datafusion = { workspace = true, features = ["pyarrow"] }
 datafusion-ffi = { workspace = true }
-datafusion-table-providers = { workspace = true, features = ["sqlite", "duckdb", "odbc", "mysql", "postgres"] }
+datafusion-table-providers = { workspace = true, features = ["sqlite", "duckdb", "odbc", "mysql", "postgres", "flight"] }
 pyo3 = { version = "0.23" }
 tokio = { version = "1.42", features = ["macros", "rt", "rt-multi-thread", "sync"] }
 duckdb = { workspace = true }

python/Cargo.toml

@@ -0,0 +1,14 @@
+from datafusion import SessionContext
+from datafusion_table_providers import flight
+
+ctx = SessionContext()
+pool = flight.FlightTableFactory()
+table_provider = pool.get_table("http://localhost:32010", {"flight.sql.query": "SELECT * FROM taxi"})
+table_name = "taxi_flight_table"
+ctx.register_table_provider(table_name, table_provider)
+ctx.sql(f"""
+        SELECT "VendorID", COUNT(*), SUM(passenger_count), SUM(total_amount)
+            FROM {table_name}
+            GROUP BY "VendorID"
+            ORDER BY COUNT(*) DESC
+        """).show()

python/examples/flight_demo.py

@@ -0,0 +1,36 @@
+# 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.
+"""Python interface for sqlite table provider."""
+
+from typing import Any, List
+from . import _internal
+
+class FlightTableFactory:
+    """Flight table factory."""
+
+    def __init__(self) -> None:
+        """Create a Flight table factory."""
+        self._raw = _internal.flight.RawFlightTableFactory()
+
+    def get_table(self, entry_point: str, options: dict) -> Any:
+        """Return the table provider for table.
+
+        Args:
+            entry_point: uri
+            options: table information
+        """
+        return self._raw.get_table(entry_point, options)

python/python/datafusion_table_providers/flight.py

@@ -1 +1,51 @@
+use std::sync::Arc;
 
+use datafusion::catalog::TableProvider;
+use datafusion_table_providers::flight::{sql::FlightSqlDriver, FlightDriver, FlightTableFactory};
+use pyo3::{prelude::*, types::PyDict};
+
+use crate::{
+    utils::{pydict_to_hashmap, to_pyerr, wait_for_future},
+    RawTableProvider,
+};
+
+#[pyclass(module = "datafusion_table_providers._internal.Flight")]
+struct RawFlightTableFactory {
+    factory: FlightTableFactory,
+}
+
+#[pymethods]
+impl RawFlightTableFactory {
+    #[new]
+    #[pyo3(signature = ())]
+    pub fn new() -> PyResult<Self> {
+        let driver: Arc<dyn FlightDriver> = Arc::new(FlightSqlDriver::new());
+
+        Ok(Self {
+            factory: FlightTableFactory::new(Arc::clone(&driver)),
+        })
+    }
+
+    pub fn get_table(
+        &self,
+        py: Python,
+        entry_point: &str,
+        options: &Bound<'_, PyDict>,
+    ) -> PyResult<RawTableProvider> {
+        let options = pydict_to_hashmap(options)?;
+        let table: Arc<dyn TableProvider> = Arc::new(
+            wait_for_future(py, self.factory.open_table(entry_point, options)).map_err(to_pyerr)?,
+        );
+
+        Ok(RawTableProvider {
+            table,
+            supports_pushdown_filters: true,
+        })
+    }
+}
+
+pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<RawFlightTableFactory>()?;
+
+    Ok(())
+}

python/src/flight.rs

@@ -61,5 +61,9 @@ fn _internal(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     postgres::init_module(&postgres)?;
     m.add_submodule(&postgres)?;
 
+    let flight = PyModule::new(py, "flight")?;
+    flight::init_module(&flight)?;
+    m.add_submodule(&flight)?;
+
     Ok(())
 }