pyo3_geoarrow/
array.rs

1use std::sync::Arc;
2
3use geoarrow_array::GeoArrowArray;
4use geoarrow_array::array::from_arrow_array;
5use geoarrow_cast::downcast::NativeType;
6use geoarrow_schema::{
7    BoxType, GeometryCollectionType, LineStringType, MultiLineStringType, MultiPointType,
8    MultiPolygonType, PointType, PolygonType,
9};
10use pyo3::exceptions::PyIndexError;
11use pyo3::intern;
12use pyo3::prelude::*;
13use pyo3::types::{PyCapsule, PyTuple, PyType};
14use pyo3_arrow::PyArray;
15use pyo3_arrow::ffi::to_array_pycapsules;
16
17use crate::PyCoordType;
18use crate::data_type::PyGeoType;
19use crate::error::{PyGeoArrowError, PyGeoArrowResult};
20use crate::scalar::PyGeoScalar;
21use crate::utils::text_repr::text_repr;
22
23/// Python wrapper for a GeoArrow geometry array.
24///
25/// This type wraps a Rust GeoArrow array and exposes it to Python through the Arrow C Data
26/// Interface. It supports zero-copy data exchange with Arrow-compatible Python libraries.
27#[pyclass(module = "geoarrow.rust.core", name = "GeoArray", subclass, frozen)]
28pub struct PyGeoArray(Arc<dyn GeoArrowArray>);
29
30impl PyGeoArray {
31    /// Create a new [`PyGeoArray`] from a GeoArrow array.
32    pub fn new(array: Arc<dyn GeoArrowArray>) -> Self {
33        Self(array)
34    }
35
36    /// Import from raw Arrow capsules
37    pub fn from_arrow_pycapsule(
38        schema_capsule: &Bound<PyCapsule>,
39        array_capsule: &Bound<PyCapsule>,
40    ) -> PyGeoArrowResult<Self> {
41        PyArray::from_arrow_pycapsule(schema_capsule, array_capsule)?.try_into()
42    }
43
44    /// Access the underlying GeoArrow array.
45    pub fn inner(&self) -> &Arc<dyn GeoArrowArray> {
46        &self.0
47    }
48
49    /// Consume this wrapper and return the underlying GeoArrow array.
50    pub fn into_inner(self) -> Arc<dyn GeoArrowArray> {
51        self.0
52    }
53
54    /// Export to a geoarrow.rust.core.GeoArray.
55    ///
56    /// This requires that you depend on geoarrow-rust-core from your Python package.
57    pub fn to_geoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
58        let geoarrow_mod = py.import(intern!(py, "geoarrow.rust.core"))?;
59        geoarrow_mod.getattr(intern!(py, "GeoArray"))?.call_method1(
60            intern!(py, "from_arrow_pycapsule"),
61            self.__arrow_c_array__(py, None)?,
62        )
63    }
64
65    /// Export to a geoarrow.rust.core.GeoArray.
66    ///
67    /// This requires that you depend on geoarrow-rust-core from your Python package.
68    pub fn into_geoarrow_py(self, py: Python) -> PyResult<Bound<PyAny>> {
69        let geoarrow_mod = py.import(intern!(py, "geoarrow.rust.core"))?;
70        let array_capsules = to_array_pycapsules(
71            py,
72            self.0.data_type().to_field("", true).into(),
73            &self.0.to_array_ref(),
74            None,
75        )?;
76        geoarrow_mod
77            .getattr(intern!(py, "GeoArray"))?
78            .call_method1(intern!(py, "from_arrow_pycapsule"), array_capsules)
79    }
80}
81
82#[pymethods]
83impl PyGeoArray {
84    #[new]
85    fn py_new(data: Self) -> Self {
86        data
87    }
88
89    #[pyo3(signature = (requested_schema=None))]
90    fn __arrow_c_array__<'py>(
91        &'py self,
92        py: Python<'py>,
93        requested_schema: Option<Bound<'py, PyCapsule>>,
94    ) -> PyGeoArrowResult<Bound<'py, PyTuple>> {
95        let field = Arc::new(self.0.data_type().to_field("", true));
96        let array = self.0.to_array_ref();
97        Ok(to_array_pycapsules(py, field, &array, requested_schema)?)
98    }
99
100    fn __eq__(&self, other: &Bound<PyAny>) -> bool {
101        // Do extraction within body because `__eq__` should never raise an exception.
102        if let Ok(other) = other.extract::<Self>() {
103            self.0.data_type() == other.0.data_type()
104                && self.0.to_array_ref() == other.0.to_array_ref()
105        } else {
106            false
107        }
108    }
109
110    // #[getter]
111    // fn __geo_interface__<'py>(&'py self, py: Python<'py>) -> PyGeoArrowResult<Bound<'py, PyAny>> {
112    //     // Note: We create a Table out of this array so that each row can be its own Feature in a
113    //     // FeatureCollection
114
115    //     let field = self.0.extension_field();
116    //     let geometry = self.0.to_array_ref();
117    //     let schema = Arc::new(Schema::new(vec![field]));
118    //     let batch = RecordBatch::try_new(schema.clone(), vec![geometry])?;
119
120    //     let mut table = geoarrow::table::Table::try_new(vec![batch], schema)?;
121    //     let json_string = table.to_json().map_err(GeoArrowError::GeozeroError)?;
122
123    //     let json_mod = py.import(intern!(py, "json"))?;
124    //     let args = (json_string,);
125    //     Ok(json_mod.call_method1(intern!(py, "loads"), args)?)
126    // }
127
128    fn __getitem__(&self, i: isize) -> PyGeoArrowResult<PyGeoScalar> {
129        // Handle negative indexes from the end
130        let i = if i < 0 {
131            let i = self.0.len() as isize + i;
132            if i < 0 {
133                return Err(PyIndexError::new_err("Index out of range").into());
134            }
135            i as usize
136        } else {
137            i as usize
138        };
139        if i >= self.0.len() {
140            return Err(PyIndexError::new_err("Index out of range").into());
141        }
142
143        PyGeoScalar::try_new(self.0.slice(i, 1))
144    }
145
146    fn __len__(&self) -> usize {
147        self.0.len()
148    }
149
150    fn __repr__(&self) -> String {
151        format!("GeoArray({})", text_repr(&self.0.data_type()))
152    }
153
154    #[classmethod]
155    fn from_arrow(_cls: &Bound<PyType>, data: Self) -> Self {
156        data
157    }
158
159    #[classmethod]
160    #[pyo3(name = "from_arrow_pycapsule")]
161    fn from_arrow_pycapsule_py(
162        _cls: &Bound<PyType>,
163        schema_capsule: &Bound<PyCapsule>,
164        array_capsule: &Bound<PyCapsule>,
165    ) -> PyGeoArrowResult<Self> {
166        Self::from_arrow_pycapsule(schema_capsule, array_capsule)
167    }
168
169    #[getter]
170    fn null_count(&self) -> usize {
171        self.0.logical_null_count()
172    }
173
174    #[pyo3(signature = (to_type, /))]
175    fn cast(&self, to_type: PyGeoType) -> PyGeoArrowResult<Self> {
176        let casted = geoarrow_cast::cast::cast(self.0.as_ref(), &to_type.into_inner())?;
177        Ok(Self(casted))
178    }
179
180    #[pyo3(
181        signature = (*, coord_type = PyCoordType::Separated),
182        text_signature = "(*, coord_type='separated')"
183    )]
184    fn downcast(&self, coord_type: PyCoordType) -> PyGeoArrowResult<Self> {
185        if let Some((native_type, dim)) =
186            geoarrow_cast::downcast::infer_downcast_type(std::iter::once(self.0.as_ref()))?
187        {
188            let metadata = self.0.data_type().metadata().clone();
189            let coord_type = coord_type.into();
190            let to_type = match native_type {
191                NativeType::Point => PointType::new(dim, metadata)
192                    .with_coord_type(coord_type)
193                    .into(),
194                NativeType::LineString => LineStringType::new(dim, metadata)
195                    .with_coord_type(coord_type)
196                    .into(),
197                NativeType::Polygon => PolygonType::new(dim, metadata)
198                    .with_coord_type(coord_type)
199                    .into(),
200                NativeType::MultiPoint => MultiPointType::new(dim, metadata)
201                    .with_coord_type(coord_type)
202                    .into(),
203                NativeType::MultiLineString => MultiLineStringType::new(dim, metadata)
204                    .with_coord_type(coord_type)
205                    .into(),
206                NativeType::MultiPolygon => MultiPolygonType::new(dim, metadata)
207                    .with_coord_type(coord_type)
208                    .into(),
209                NativeType::GeometryCollection => GeometryCollectionType::new(dim, metadata)
210                    .with_coord_type(coord_type)
211                    .into(),
212                NativeType::Rect => BoxType::new(dim, metadata).into(),
213            };
214            self.cast(PyGeoType::new(to_type))
215        } else {
216            Ok(Self::new(self.0.clone()))
217        }
218    }
219
220    #[getter]
221    fn r#type(&self) -> PyGeoType {
222        self.0.data_type().into()
223    }
224}
225
226impl From<Arc<dyn GeoArrowArray>> for PyGeoArray {
227    fn from(value: Arc<dyn GeoArrowArray>) -> Self {
228        Self(value)
229    }
230}
231
232impl From<PyGeoArray> for Arc<dyn GeoArrowArray> {
233    fn from(value: PyGeoArray) -> Self {
234        value.0
235    }
236}
237
238impl<'py> FromPyObject<'_, 'py> for PyGeoArray {
239    type Error = PyErr;
240
241    fn extract(ob: Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {
242        Ok(ob.extract::<PyArray>()?.try_into()?)
243    }
244}
245
246impl TryFrom<PyArray> for PyGeoArray {
247    type Error = PyGeoArrowError;
248
249    fn try_from(value: PyArray) -> Result<Self, Self::Error> {
250        let (array, field) = value.into_inner();
251        let geo_arr = from_arrow_array(&array, &field)?;
252        Ok(Self(geo_arr))
253    }
254}
pyo3_geoarrow/array.rs

pyo3_geoarrow/
array.rs
