diff --git a/python/Cargo.lock b/python/Cargo.lock index f2a8b5b..f1535a0 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -1778,7 +1778,7 @@ dependencies = [ [[package]] name = "geodatafusion" -version = "0.1.0-dev" +version = "0.1.0-beta.1" dependencies = [ "arrow-arith", "arrow-array", diff --git a/python/src/lib.rs b/python/src/lib.rs index e8de730..9b29c48 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -1,6 +1,7 @@ #![cfg_attr(not(test), warn(unused_crate_dependencies))] pub(crate) mod constants; +mod table_provider; mod udf; mod utils; diff --git a/python/src/table_provider/flatgeobuf.rs b/python/src/table_provider/flatgeobuf.rs new file mode 100644 index 0000000..c6e164e --- /dev/null +++ b/python/src/table_provider/flatgeobuf.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); + +#[pymethods] +impl PyFlatGeobufTableProvider { + pub fn __datafusion_table_provider__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let name = cr"datafusion_table_provider".into(); + + let provider = FFI_TableProvider::new(self.0.clone(), false, None); + + PyCapsule::new(py, provider, Some(name)) + } +} diff --git a/python/src/table_provider/mod.rs b/python/src/table_provider/mod.rs new file mode 100644 index 0000000..89d76c4 --- /dev/null +++ b/python/src/table_provider/mod.rs @@ -0,0 +1 @@ +mod flatgeobuf;