feat(python): Constructor for GeoChunkedArray (#1383)

Closes #1382

Author: kylebarron

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

@@ -1,4 +1,9 @@
-from arro3.core.types import ArrowStreamExportable
+from typing import Sequence, overload
+from arro3.core.types import (
+    ArrowStreamExportable,
+    ArrowArrayExportable,
+    ArrowSchemaExportable,
+)
 from geoarrow.rust.core._scalar import GeoScalar
 
 from ._array import GeoArray
@@ -12,6 +17,29 @@ class GeoChunkedArray:
     This class is used to handle chunked arrays in GeoArrow, which can be
     composed of multiple chunks of data.
     """
+    @overload
+    def __init__(
+        self, arrays: ArrowArrayExportable | ArrowStreamExportable, type: None = None
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        arrays: Sequence[ArrowArrayExportable],
+        type: ArrowSchemaExportable | None = None,
+    ) -> None: ...
+    def __init__(
+        self,
+        arrays: ArrowArrayExportable
+        | ArrowStreamExportable
+        | Sequence[ArrowArrayExportable],
+        type: ArrowSchemaExportable | None = None,
+    ) -> None:
+        """Construct a new GeoChunkedArray.
+
+        Args:
+            arrays: _description_
+            type: _description_. Defaults to None.
+        """
 
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
         """

python/tests/core/operations/test_geometry_col.py

@@ -3,7 +3,7 @@
 import pytest
 import shapely
 from arro3.core import Table
-from geoarrow.rust.core import GeoArray, geometry_col
+from geoarrow.rust.core import GeoArray, GeoChunkedArray, geometry_col
 
 
 def geoarrow_array():
@@ -42,11 +42,10 @@ def test_geo_array_input():
     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_geo_chunked_array_input():
+    arr = geoarrow_array()
+    chunked = GeoChunkedArray([arr, arr])
+    assert chunked == geometry_col(chunked)
 
 
 def test_table_no_geom_cols():

python/tests/core/operations/test_type_id.py

@@ -91,6 +91,6 @@ def test_points_chunked():
     geoms2 = shapely.points([10, 20, 30], [40, 50, 60])
     arr1 = GeoArray.from_arrow(gpd.GeoSeries(geoms1).to_arrow("geoarrow"))
     arr2 = GeoArray.from_arrow(gpd.GeoSeries(geoms2).to_arrow("geoarrow"))
-    ca = GeoChunkedArray.from_arrow(ChunkedArray([arr1, arr2]))
+    ca = GeoChunkedArray([arr1, arr2])
     out = get_type_id(ca).read_all()
     assert (np.asarray(out) == 1).all()

python/tests/core/test_chunked_array.py

@@ -14,7 +14,7 @@ def test_eq():
     geoms2 = shapely.points([10, 20, 30], [40, 50, 60])
     arr1 = GeoArray.from_arrow(gpd.GeoSeries(geoms1).to_arrow("geoarrow"))
     arr2 = GeoArray.from_arrow(gpd.GeoSeries(geoms2).to_arrow("geoarrow"))
-    ca = GeoChunkedArray.from_arrow(ChunkedArray([arr1, arr2]))
+    ca = GeoChunkedArray([arr1, arr2])
 
     assert ca == ca
 
@@ -30,7 +30,7 @@ def test_getitem():
     gdf = gpd.read_file(geodatasets.get_path("ny.bb"))
     arr1 = GeoArray.from_arrow(gdf.geometry.iloc[:2].to_arrow("geoarrow"))
     arr2 = GeoArray.from_arrow(gdf.geometry.iloc[2:].to_arrow("geoarrow"))
-    ca = GeoChunkedArray.from_arrow(ChunkedArray([arr1, arr2]))
+    ca = GeoChunkedArray([arr1, arr2])
 
     for i in range(len(ca)):
         assert shapely.geometry.shape(ca[i]).equals(gdf.geometry.iloc[i])  # type: ignore
@@ -41,7 +41,7 @@ def test_repr():
     geoms2 = shapely.points([10, 20, 30], [40, 50, 60])
     arr1 = GeoArray.from_arrow(gpd.GeoSeries(geoms1).to_arrow("geoarrow"))
     arr2 = GeoArray.from_arrow(gpd.GeoSeries(geoms2).to_arrow("geoarrow"))
-    ca = GeoChunkedArray.from_arrow(ChunkedArray([arr1, arr2]))
+    ca = GeoChunkedArray([arr1, arr2])
     assert (
         repr(ca) == 'GeoChunkedArray(Point(dimension="XY", coord_type="interleaved"))'
     )
@@ -51,7 +51,7 @@ def test_downcast():
     coords = np.array([[1, 4], [2, 5], [3, 6]], dtype=np.float64)
 
     point_arr = points(coords)
-    point_ca = GeoChunkedArray.from_arrow(ChunkedArray([point_arr]))
+    point_ca = GeoChunkedArray([point_arr])
     geometry_array = point_ca.cast(geometry())
     point_ca2 = geometry_array.downcast(coord_type="interleaved")
     assert point_ca == point_ca2
@@ -62,12 +62,22 @@ def test_downcast_with_crs():
 
     crs = "EPSG:4326"
     point_arr = points(coords, crs=crs)
-    point_ca = GeoChunkedArray.from_arrow(ChunkedArray([point_arr]))
+    point_ca = GeoChunkedArray([point_arr])
     geometry_array = point_ca.cast(geometry(crs=crs))
     point_ca2 = geometry_array.downcast(coord_type="interleaved")
     assert point_ca == point_ca2
 
 
+def constructor_existing_chunked_array():
+    geoms1 = shapely.points([1, 2, 3], [4, 5, 6])
+    geoms2 = shapely.points([10, 20, 30], [40, 50, 60])
+    arr1 = GeoArray.from_arrow(gpd.GeoSeries(geoms1).to_arrow("geoarrow"))
+    arr2 = GeoArray.from_arrow(gpd.GeoSeries(geoms2).to_arrow("geoarrow"))
+    ca = ChunkedArray([arr1, arr2])
+    geo_ca = GeoChunkedArray(ca)
+    assert geo_ca.chunk(0) == arr1
+
+
 class CustomException(Exception):
     pass
 

rust/pyo3-geoarrow/src/chunked_array.rs

@@ -18,6 +18,7 @@ use pyo3_arrow::{PyArrayReader, PyChunkedArray};
 
 use crate::data_type::PyGeoType;
 use crate::error::{PyGeoArrowError, PyGeoArrowResult};
+use crate::input::AnyGeoArray;
 use crate::scalar::PyGeoScalar;
 use crate::utils::text_repr::text_repr;
 use crate::{PyCoordType, PyGeoArray};
@@ -93,6 +94,46 @@ impl PyGeoChunkedArray {
 
 #[pymethods]
 impl PyGeoChunkedArray {
+    #[new]
+    #[pyo3(signature = (arrays, r#type=None))]
+    fn init(
+        py: Python,
+        arrays: &Bound<PyAny>,
+        r#type: Option<PyGeoType>,
+    ) -> PyGeoArrowResult<Self> {
+        if arrays.hasattr(intern!(py, "__arrow_c_array__"))?
+            || arrays.hasattr(intern!(py, "__arrow_c_stream__"))?
+        {
+            Ok(arrays.extract::<AnyGeoArray>()?.into_chunked_array()?)
+        } else if let Ok(geo_arrays) = arrays.extract::<Vec<PyGeoArray>>() {
+            let geo_arrays = geo_arrays
+                .into_iter()
+                .map(|arr| arr.into_inner())
+                .collect::<Vec<_>>();
+
+            if !geo_arrays
+                .windows(2)
+                .all(|w| w[0].data_type() == w[1].data_type())
+            {
+                return Err(PyTypeError::new_err(
+                    "Cannot create a ChunkedArray with differing data types.",
+                )
+                .into());
+            }
+
+            let geo_type = r#type
+                .map(|py_data_type| py_data_type.into_inner())
+                .unwrap_or_else(|| geo_arrays[0].data_type());
+
+            Ok(Self::try_new(geo_arrays, geo_type)?)
+        } else {
+            Err(
+                PyTypeError::new_err("Expected ChunkedArray-like input or sequence of arrays.")
+                    .into(),
+            )
+        }
+    }
+
     #[pyo3(signature = (requested_schema=None))]
     fn __arrow_c_stream__<'py>(
         &self,