feat(python, geoarrow-rust-core): Implement geometry_col (#1381)

to access a geometry column from a RecordBatch or stream

Author: kylebarron

python/Cargo.lock

@@ -1,22 +1,9 @@
 from __future__ import annotations
 
-from pathlib import Path
-from typing import Any, Sequence, Tuple, overload
+from typing import overload
 
-from arro3.core import Table
 from arro3.core.types import ArrowArrayExportable, ArrowStreamExportable
 
-try:
-    import numpy as np
-    from numpy.typing import NDArray
-except ImportError:
-    pass
-
-try:
-    import geopandas as gpd
-except ImportError:
-    pass
-
 from ._array import GeoArray as GeoArray
 from ._array_reader import GeoArrayReader as GeoArrayReader
 from ._chunked_array import GeoChunkedArray as GeoChunkedArray
@@ -51,21 +38,34 @@ from ._interop import to_wkt as to_wkt
 from ._operations import get_type_id as get_type_id
 from ._scalar import GeoScalar as GeoScalar
 
-# @overload
-# def geometry_col(input: ArrowArrayExportable) -> GeoArray: ...
-# @overload
-# def geometry_col(input: ArrowStreamExportable) -> GeoChunkedArray: ...
-# def geometry_col(
-#     input: ArrowArrayExportable | ArrowStreamExportable,
-# ) -> GeoArray | GeoChunkedArray:
-#     """Access the geometry column of a Table or RecordBatch
-
-#     Args:
-#         input: The Arrow RecordBatch or Table to extract the geometry column from.
-
-#     Returns:
-#         A geometry array or chunked array.
-#     """
+@overload
+def geometry_col(
+    input: ArrowArrayExportable,
+    *,
+    name: str | None = None,
+) -> GeoArray: ...
+@overload
+def geometry_col(
+    input: ArrowStreamExportable,
+    *,
+    name: str | None = None,
+) -> GeoArrayReader: ...
+def geometry_col(
+    input: ArrowArrayExportable | ArrowStreamExportable,
+    *,
+    name: str | None = None,
+) -> GeoArray | GeoArrayReader:
+    """Access the geometry column of a Table or RecordBatch
+
+    Args:
+        input: The Arrow RecordBatch or Table to extract the geometry column from.
+
+    Keyword Args:
+        name: The name of the geometry column to extract. If not provided, an error will be produced if there are multiple columns with GeoArrow metadata.
+
+    Returns:
+        A geometry array or chunked array.
+    """
 
 # Interop
 

python/geoarrow-core/python/geoarrow/rust/core/_rust.pyi

@@ -3,6 +3,7 @@
 mod constructors;
 mod interop;
 mod operations;
+mod table;
 
 use pyo3::exceptions::PyRuntimeWarning;
 use pyo3::intern;
@@ -118,7 +119,7 @@ fn _rust(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
 
     // Top-level table functions
 
-    // m.add_function(wrap_pyfunction!(crate::table::geometry_col, m)?)?;
+    m.add_function(wrap_pyfunction!(crate::table::geometry_col, m)?)?;
 
     // Interop
 

python/geoarrow-core/src/lib.rs

@@ -1,39 +1,90 @@
-use crate::ffi::to_python::{chunked_native_array_to_pyobject, native_array_to_pyobject};
-use crate::interop::util::pytable_to_table;
-use geoarrow::array::NativeArrayDyn;
-use geoarrow::schema::GeoSchemaExt;
-use pyo3::exceptions::PyNotImplementedError;
+use arrow_schema::Schema;
+use geoarrow_array::{GeoArrowArrayIterator, WrapArray};
+use geoarrow_schema::GeoArrowType;
+use pyo3::IntoPyObjectExt;
+use pyo3::exceptions::{PyIndexError, PyValueError};
 use pyo3::prelude::*;
+use pyo3::pybacked::PyBackedStr;
 use pyo3_arrow::input::AnyRecordBatch;
-use pyo3_geoarrow::PyGeoArrowResult;
+use pyo3_geoarrow::input::AnyGeoArray;
+use pyo3_geoarrow::{PyGeoArray, PyGeoArrayReader, PyGeoArrowResult};
 
 #[pyfunction]
-pub fn geometry_col(py: Python, input: AnyRecordBatch) -> PyGeoArrowResult<PyObject> {
-    match input {
-        AnyRecordBatch::RecordBatch(rb) => {
-            let batch = rb.into_inner();
-            let schema = batch.schema();
+#[pyo3(signature = (input, *, name = None))]
+pub fn geometry_col<'py>(
+    py: Python<'py>,
+    input: Bound<'py, PyAny>,
+    name: Option<PyBackedStr>,
+) -> PyGeoArrowResult<Bound<'py, PyAny>> {
+    // If the input is already a GeoArray, just return it
+    if let Ok(input) = input.extract::<AnyGeoArray>() {
+        match input {
+            AnyGeoArray::Array(array) => {
+                return Ok(array.into_bound_py_any(py)?);
+            }
+            AnyGeoArray::Stream(stream) => {
+                return Ok(stream.into_bound_py_any(py)?);
+            }
+        }
+    }
 
-            let geom_indices = schema.as_ref().geometry_columns();
-            let index = if geom_indices.len() == 1 {
-                geom_indices[0]
-            } else {
-                return Err(PyNotImplementedError::new_err(
-                    "Accessing from multiple geometry columns not yet supported",
-                )
-                .into());
-            };
-
-            let field = schema.field(index);
-            let array = batch.column(index).as_ref();
-            let geo_arr = NativeArrayDyn::from_arrow_array(array, field)?.into_inner();
-            native_array_to_pyobject(py, geo_arr)
+    // Otherwise, assume it's a RecordBatch or RecordBatchStream
+    let input = input.extract::<AnyRecordBatch>()?;
+
+    let schema = input.schema()?;
+
+    let (geom_index, geom_type) = if let Some(name) = name {
+        let (idx, field) = schema
+            .column_with_name(&name)
+            .ok_or(PyIndexError::new_err(format!(
+                "Column name {name} not found"
+            )))?;
+
+        let geom_type = GeoArrowType::from_arrow_field(field)?;
+        (idx, geom_type)
+    } else {
+        let geom_cols = geometry_columns(schema.as_ref());
+        if geom_cols.is_empty() {
+            return Err(PyValueError::new_err("No geometry columns found").into());
+        } else if geom_cols.len() == 1 {
+            geom_cols.into_iter().next().unwrap()
+        } else {
+            return Err(PyValueError::new_err(
+                "Multiple geometry columns: 'name' must be provided.",
+            )
+            .into());
+        }
+    };
+
+    match input {
+        AnyRecordBatch::RecordBatch(batch) => {
+            let geo_array = geom_type.wrap_array(batch.as_ref().column(geom_index).as_ref())?;
+            Ok(PyGeoArray::new(geo_array).into_bound_py_any(py)?)
         }
         AnyRecordBatch::Stream(stream) => {
-            let table = stream.into_table()?;
-            let table = pytable_to_table(table)?;
-            let chunked_geom_arr = table.geometry_column(None)?;
-            chunked_native_array_to_pyobject(py, chunked_geom_arr)
+            let reader = stream.into_reader()?;
+            let output_geo_type = geom_type.clone();
+            let iter = reader
+                .into_iter()
+                .map(move |batch| geom_type.wrap_array(batch?.column(geom_index).as_ref()));
+            let output_reader = Box::new(GeoArrowArrayIterator::new(iter, output_geo_type));
+
+            Ok(PyGeoArrayReader::new(output_reader).into_bound_py_any(py)?)
         }
     }
 }
+
+fn geometry_columns(schema: &Schema) -> Vec<(usize, GeoArrowType)> {
+    schema
+        .fields()
+        .iter()
+        .enumerate()
+        .filter_map(|(idx, field)| {
+            if let Ok(geom_type) = GeoArrowType::from_extension_field(field) {
+                Some((idx, geom_type))
+            } else {
+                None
+            }
+        })
+        .collect()
+}

python/geoarrow-core/src/table/mod.rs

@@ -0,0 +1,72 @@
+import geopandas as gpd
+import numpy as np
+import pytest
+import shapely
+from arro3.core import Table
+from geoarrow.rust.core import GeoArray, geometry_col
+
+
+def geoarrow_array():
+    geoms = shapely.points([1, 2, 3], [4, 5, 6])
+    return GeoArray.from_arrow(gpd.GeoSeries(geoms).to_arrow("geoarrow"))
+
+
+def test_batch_no_geom_cols():
+    arr = np.array([1, 2, 3])
+    # We should have simpler RecordBatch constructor
+    # https://github.com/kylebarron/arro3/issues/418
+    batch = Table.from_arrays([arr], names=["no_geom"]).to_batches()[0]
+    with pytest.raises(ValueError, match="No geometry columns found"):
+        geometry_col(batch)
+
+
+def test_batch_one_geom_col():
+    arr = geoarrow_array()
+    batch = Table.from_arrays([arr], names=["geom"]).to_batches()[0]
+    fetched_arr = geometry_col(batch)
+    assert arr == fetched_arr
+
+
+def test_batch_two_geom_cols():
+    arr = geoarrow_array()
+    batch = Table.from_arrays([arr, arr], names=["geom1", "geom2"]).to_batches()[0]
+    with pytest.raises(ValueError, match="Multiple geometry columns"):
+        geometry_col(batch)
+
+    assert geometry_col(batch, name="geom1") == arr
+    assert geometry_col(batch, name="geom2") == arr
+
+
+def test_geo_array_input():
+    arr = geoarrow_array()
+    assert arr == geometry_col(arr)
+
+
+# TODO: implement once we have easy GeoChunkedArray constructor
+# def test_geo_chunked_array_input():
+#     arr = geoarrow_array()
+#     chunked = GeoChunkedArray.from_arrays([arr, arr])
+#     assert chunked == geometry_col(chunked)
+
+
+def test_table_no_geom_cols():
+    arr = np.array([1, 2, 3])
+    table = Table.from_arrays([arr], names=["no_geom"])
+    with pytest.raises(ValueError, match="No geometry columns found"):
+        geometry_col(table)
+
+
+def test_table_one_geom_col():
+    arr = geoarrow_array()
+    table = Table.from_arrays([arr], names=["geom"])
+    assert geometry_col(table).read_next_array() == arr
+
+
+def test_table_two_geom_cols():
+    arr = geoarrow_array()
+    table = Table.from_arrays([arr, arr], names=["geom1", "geom2"])
+    with pytest.raises(ValueError, match="Multiple geometry columns"):
+        geometry_col(table)
+
+    assert geometry_col(table, name="geom1").read_next_array() == arr
+    assert geometry_col(table, name="geom2").read_next_array() == arr

python/tests/core/operations/test_geometry_col.py

@@ -2,6 +2,8 @@
 
 ## [Unreleased]
 
+- Add `PyGeoArrowArrayReader::into_geoarrow_py`.
+
 ## 0.6.1 - 2025-10-16
 
 - docs(pyo3-geoarrow): Improve docs #1377

rust/pyo3-geoarrow/CHANGELOG.md

@@ -51,17 +51,31 @@ impl PyGeoArray {
         self.0
     }
 
-    /// Export to a geoarrow.rust.core.GeoArrowArray.
+    /// Export to a geoarrow.rust.core.GeoArray.
     ///
     /// This requires that you depend on geoarrow-rust-core from your Python package.
     pub fn to_geoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let geoarrow_mod = py.import(intern!(py, "geoarrow.rust.core"))?;
+        geoarrow_mod.getattr(intern!(py, "GeoArray"))?.call_method1(
+            intern!(py, "from_arrow_pycapsule"),
+            self.__arrow_c_array__(py, None)?,
+        )
+    }
+
+    /// Export to a geoarrow.rust.core.GeoArray.
+    ///
+    /// This requires that you depend on geoarrow-rust-core from your Python package.
+    pub fn into_geoarrow_py(self, py: Python) -> PyResult<Bound<PyAny>> {
+        let geoarrow_mod = py.import(intern!(py, "geoarrow.rust.core"))?;
+        let array_capsules = to_array_pycapsules(
+            py,
+            self.0.data_type().to_field("", true).into(),
+            &self.0.to_array_ref(),
+            None,
+        )?;
         geoarrow_mod
-            .getattr(intern!(py, "GeoArrowArray"))?
-            .call_method1(
-                intern!(py, "from_arrow_pycapsule"),
-                self.__arrow_c_array__(py, None)?,
-            )
+            .getattr(intern!(py, "GeoArray"))?
+            .call_method1(intern!(py, "from_arrow_pycapsule"), array_capsules)
     }
 }