Repository: stac-utils/stac-rs (language: python)

Issue:
Include duckdb in Python bindings
This will enable DuckDB search from Python


Discussion / hints:

Fix patch:
diff --git a/python/Cargo.toml b/python/Cargo.toml
index a1f926db..77ef6f67 100644
--- a/python/Cargo.toml
+++ b/python/Cargo.toml
@@ -13,6 +13,9 @@ name = "stacrs"
 crate-type = ["cdylib"]
 
 [dependencies]
+duckdb = { workspace = true, features = [
+    "bundled",
+] } # we don't use it directly, but we need to ensure it's bundled
 geojson = { workspace = true }
 pyo3 = { workspace = true, features = ["extension-module"] }
 pythonize = { workspace = true }
@@ -25,4 +28,5 @@ stac = { workspace = true, features = [
     "validate-blocking",
 ] }
 stac-api = { workspace = true, features = ["client"] }
+stac-duckdb = { workspace = true }
 tokio = { workspace = true, features = ["rt"] }
diff --git a/python/data/extended-item.parquet b/python/data/extended-item.parquet
new file mode 100644
index 00000000..5ef4689e
Binary files /dev/null and b/python/data/extended-item.parquet differ
diff --git a/python/src/error.rs b/python/src/error.rs
index 395648be..7c27ea4a 100644
--- a/python/src/error.rs
+++ b/python/src/error.rs
@@ -16,6 +16,12 @@ impl From<stac_api::Error> for Error {
     }
 }
 
+impl From<stac_duckdb::Error> for Error {
+    fn from(value: stac_duckdb::Error) -> Self {
+        Error(value.to_string())
+    }
+}
+
 impl From<geojson::Error> for Error {
     fn from(value: geojson::Error) -> Self {
         Error(value.to_string())
diff --git a/python/src/lib.rs b/python/src/lib.rs
index 524eb996..12f69490 100644
--- a/python/src/lib.rs
+++ b/python/src/lib.rs
@@ -7,6 +7,7 @@ mod search;
 mod validate;
 mod write;
 
+use duckdb as _;
 use error::Error;
 use pyo3::prelude::*;
 
diff --git a/python/src/search.rs b/python/src/search.rs
index 4e016788..d5bf630f 100644
--- a/python/src/search.rs
+++ b/python/src/search.rs
@@ -6,6 +6,7 @@ use pyo3::{
 use serde::de::DeserializeOwned;
 use stac::Format;
 use stac_api::{BlockingClient, Fields, Item, ItemCollection, Items, Search};
+use stac_duckdb::Client;
 use std::str::FromStr;
 use tokio::runtime::Builder;
 
@@ -40,6 +41,9 @@ use tokio::runtime::Builder;
 ///         will be interpreted as cql2-text, dictionaries as cql2-json.
 ///     query (dict[str, Any] | None): Additional filtering based on properties.
 ///         It is recommended to use filter instead, if possible.
+///     use_duckdb (bool | None): Query with DuckDB. If None and the href has a
+///         'parquet' or 'geoparquet' extension, will be set to True. Defaults
+///         to None.
 ///
 /// Returns:
 ///     list[dict[str, Any]]: A list of the returned STAC items.
@@ -53,7 +57,7 @@ use tokio::runtime::Builder;
 ///     ...     max_items=1,
 ///     ... )
 #[pyfunction]
-#[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None))]
+#[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, use_duckdb=None))]
 pub fn search<'py>(
     py: Python<'py>,
     href: String,
@@ -69,6 +73,7 @@ pub fn search<'py>(
     sortby: Option<StringOrList>,
     filter: Option<StringOrDict>,
     query: Option<Py<PyDict>>,
+    use_duckdb: Option<bool>,
 ) -> PyResult<Bound<'py, PyList>> {
     let items = search_items(
         href,
@@ -84,6 +89,7 @@ pub fn search<'py>(
         sortby,
         filter,
         query,
+        use_duckdb,
     )?;
     pythonize::pythonize(py, &items)
         .map_err(PyErr::from)
@@ -126,6 +132,9 @@ pub fn search<'py>(
 ///     format (str | None): The output format. If none, will be inferred from
 ///         the outfile extension, and if that fails will fall back to compact JSON.
 ///     options (list[tuple[str, str]] | None): Configuration values to pass to the object store backend.
+///     use_duckdb (bool | None): Query with DuckDB. If None and the href has a
+///         'parquet' or 'geoparquet' extension, will be set to True. Defaults
+///         to None.
 ///
 /// Returns:
 ///     list[dict[str, Any]]: A list of the returned STAC items.
@@ -139,7 +148,7 @@ pub fn search<'py>(
 ///     ...     max_items=1,
 ///     ... )
 #[pyfunction]
-#[pyo3(signature = (outfile, href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, format=None, options=None))]
+#[pyo3(signature = (outfile, href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, format=None, options=None, use_duckdb=None))]
 pub fn search_to(
     outfile: String,
     href: String,
@@ -157,6 +166,7 @@ pub fn search_to(
     query: Option<Py<PyDict>>,
     format: Option<String>,
     options: Option<Vec<(String, String)>>,
+    use_duckdb: Option<bool>,
 ) -> PyResult<usize> {
     let items = search_items(
         href,
@@ -172,6 +182,7 @@ pub fn search_to(
         sortby,
         filter,
         query,
+        use_duckdb,
     )?;
     let format = format
         .map(|s| s.parse())
@@ -206,8 +217,8 @@ fn search_items(
     sortby: Option<StringOrList>,
     filter: Option<StringOrDict>,
     query: Option<Py<PyDict>>,
+    use_duckdb: Option<bool>,
 ) -> PyResult<Vec<Item>> {
-    let client = BlockingClient::new(&href).map_err(Error::from)?;
     let mut fields = Fields::default();
     if let Some(include) = include {
         fields.include = include.into();
@@ -225,7 +236,7 @@ fn search_items(
             .map(|q| pythonize::depythonize(&q.into_bound(py)))
             .transpose()
     })?;
-    let search = Search {
+    let mut search = Search {
         intersects: intersects.map(|i| i.into()).transpose()?,
         ids: ids.map(|ids| ids.into()),
         collections: collections.map(|c| c.into()),
@@ -243,18 +254,33 @@ fn search_items(
             ..Default::default()
         },
     };
-    let items = client.search(search).map_err(Error::from)?;
-    if let Some(max_items) = max_items {
-        items
-            .take(max_items)
-            .collect::<Result<_, _>>()
+    if use_duckdb
+        .unwrap_or_else(|| matches!(Format::infer_from_href(&href), Some(Format::Geoparquet(_))))
+    {
+        if let Some(max_items) = max_items {
+            search.items.limit = Some(max_items.try_into()?);
+        }
+        let client = Client::from_href(href).map_err(Error::from)?;
+        client
+            .search_to_json(search)
+            .map(|item_collection| item_collection.items)
             .map_err(Error::from)
             .map_err(PyErr::from)
     } else {
-        items
-            .collect::<Result<_, _>>()
-            .map_err(Error::from)
-            .map_err(PyErr::from)
+        let client = BlockingClient::new(&href).map_err(Error::from)?;
+        let items = client.search(search).map_err(Error::from)?;
+        if let Some(max_items) = max_items {
+            items
+                .take(max_items)
+                .collect::<Result<_, _>>()
+                .map_err(Error::from)
+                .map_err(PyErr::from)
+        } else {
+            items
+                .collect::<Result<_, _>>()
+                .map_err(Error::from)
+                .map_err(PyErr::from)
+        }
     }
 }
 


Test patch:
diff --git a/python/tests/test_search.py b/python/tests/test_search.py
index b4d57599..ba2f07a8 100644
--- a/python/tests/test_search.py
+++ b/python/tests/test_search.py
@@ -44,3 +44,8 @@ def test_search_to_geoparquet(tmp_path: Path) -> None:
     table = pyarrow.parquet.read_table(tmp_path / "out.parquet")
     items = list(stac_geoparquet.arrow.stac_table_to_items(table))
     assert len(items) == 1
+
+
+def test_search_geoparquet(data: Path) -> None:
+    items = stacrs.search(str(data / "extended-item.parquet"))
+    assert len(items) == 1