feat(geodatafusion-python): FlatGeobuf TableProvider

Author: kylebarron

python/Cargo.lock

@@ -27,6 +27,7 @@ geoarrow-cast = { path = "../rust/geoarrow-cast" }
 geoarrow-flatgeobuf = { path = "../rust/geoarrow-flatgeobuf" }
 geoarrow-schema = { path = "../rust/geoarrow-schema" }
 geodatafusion = { path = "../rust/geodatafusion" }
+geodatafusion-flatgeobuf = { path = "../rust/geodatafusion-flatgeobuf" }
 geoparquet = { path = "../rust/geoparquet" }
 geozero = "0.14"
 http-range-client = { version = "0.9.0", default-features = false }
@@ -35,6 +36,7 @@ numpy = "0.25"
 object_store = "0.12"
 parquet = "56"
 pyo3 = { version = "0.25", features = ["hashbrown", "serde", "anyhow"] }
+pyo3-async-runtimes = { version = "0.25", features = ["tokio-runtime"] }
 # https://github.com/kylebarron/arro3/pull/354
 pyo3-arrow = { git = "https://github.com/kylebarron/arro3", rev = "fda03cebe9cfbf8ad292f8070ea21c22b32af1a3" }
 pyo3-geoarrow = { path = "../rust/pyo3-geoarrow" }

python/Cargo.toml

@@ -23,5 +23,7 @@ crate-type = ["cdylib"]
 datafusion = { workspace = true }
 datafusion-ffi = { workspace = true }
 geodatafusion = { workspace = true }
+geodatafusion-flatgeobuf = { workspace = true }
 pyo3 = { workspace = true }
+pyo3-async-runtimes = { workspace = true }
 pyo3-geoarrow = { workspace = true, features = ["geozero"] }

python/geodatafusion/Cargo.toml

@@ -1,6 +1,7 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 
 pub(crate) mod constants;
+mod table_provider;
 mod udf;
 mod utils;
 
@@ -55,5 +56,10 @@ fn _rust(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
         .getattr(intern!(py, "modules"))?
         .set_item("geodatafusion.geo", geo_mod)?;
 
+    m.add_function(wrap_pyfunction!(
+        table_provider::flatgeobuf::new_flatgeobuf,
+        m
+    )?)?;
+
     Ok(())
 }

python/geodatafusion/src/lib.rs

@@ -0,0 +1,50 @@
+use datafusion::catalog::TableProvider;
+use datafusion::datasource::listing::{
+    ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
+};
+use datafusion::prelude::SessionContext;
+use datafusion_ffi::table_provider::FFI_TableProvider;
+use geodatafusion_flatgeobuf::FlatGeobufFormat;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
+use pyo3_async_runtimes::tokio::get_runtime;
+use std::sync::Arc;
+
+#[pyfunction]
+pub(crate) fn new_flatgeobuf(path: &str) -> PyFlatGeobufTableProvider {
+    let format = Arc::new(FlatGeobufFormat::default());
+
+    let options = ListingOptions::new(format).with_file_extension(".fgb");
+
+    let table_path = ListingTableUrl::parse(path).unwrap();
+
+    let state = SessionContext::new().state();
+    let runtime = get_runtime();
+    let inferred_schema =
+        runtime.block_on(async { options.infer_schema(&state, &table_path).await.unwrap() });
+
+    let config = ListingTableConfig::new(table_path)
+        .with_listing_options(options)
+        .with_schema(inferred_schema);
+
+    let table = ListingTable::try_new(config).unwrap();
+    PyFlatGeobufTableProvider(Arc::new(table))
+}
+
+#[pyclass(module = "geodatafusion", name = "FlatGeobufTableProvider", frozen)]
+pub(crate) struct PyFlatGeobufTableProvider(Arc<dyn TableProvider + Send>);
+
+#[pymethods]
+impl PyFlatGeobufTableProvider {
+    pub fn __datafusion_table_provider__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_table_provider".into();
+
+        let provider = FFI_TableProvider::new(self.0.clone(), false, None);
+
+        PyCapsule::new(py, provider, Some(name))
+    }
+}

python/geodatafusion/src/table_provider/flatgeobuf.rs

@@ -0,0 +1 @@
+pub(crate) mod flatgeobuf;

python/geodatafusion/src/table_provider/mod.rs

@@ -10,3 +10,18 @@ def test_simple():
     sql = "SELECT ST_AsText(ST_GeomFromText('POINT(1 2)'));"
     df = ctx.sql(sql)
     assert df.to_arrow_table().columns[0][0].as_py() == "POINT(1 2)"
+
+
+def test_flatgeobuf():
+    from geodatafusion import new_flatgeobuf
+
+    path = "/Users/kyle/Downloads/countries.fgb"
+    test = new_flatgeobuf(path)
+
+    ctx = SessionContext()
+    register_all(ctx)
+    ctx.register_table_provider("countries", test)
+
+    sql = "SELECT * FROM countries;"
+    df = ctx.sql(sql)
+    df.show()