Python: Convert WKT and EWKB to geoarrow (#369)

geoarrow · Jan 1, 2024 · 62fef84 · 62fef84
1 parent 4f5f617
commit 62fef84
Show file tree

Hide file tree

Showing 25 changed files with 346 additions and 97 deletions.
diff --git a/python/core/python/geoarrow/rust/core/rust.pyi b/python/core/python/geoarrow/rust/core/rust.pyi
@@ -167,7 +167,11 @@ class MixedGeometryArray:
     def geodesic_perimeter(self) -> Float64Array: ...
     def is_empty(self) -> BooleanArray: ...
     @classmethod
+    def from_ewkb(cls, array: _ArrowArrayExportable) -> Self: ...
+    @classmethod
     def from_wkb(cls, array: _ArrowArrayExportable) -> Self: ...
+    @classmethod
+    def from_wkt(cls, array: _ArrowArrayExportable) -> Self: ...
     def to_wkb(self) -> WKBArray: ...
 
 class GeometryCollectionArray:
@@ -188,7 +192,11 @@ class GeometryCollectionArray:
     def geodesic_perimeter(self) -> Float64Array: ...
     def is_empty(self) -> BooleanArray: ...
     @classmethod
+    def from_ewkb(cls, array: _ArrowArrayExportable) -> Self: ...
+    @classmethod
     def from_wkb(cls, array: _ArrowArrayExportable) -> Self: ...
+    @classmethod
+    def from_wkt(cls, array: _ArrowArrayExportable) -> Self: ...
     def to_wkb(self) -> WKBArray: ...
 
 class WKBArray:
@@ -526,6 +534,18 @@ def centroid(input: _ArrowArrayExportable) -> PointArray: ...
 def convex_hull(input: _ArrowArrayExportable) -> PolygonArray: ...
 
 # I/O
+def from_ewkb(
+    input: _ArrowArrayExportable,
+) -> (
+    PointArray
+    | LineStringArray
+    | PolygonArray
+    | MultiPointArray
+    | MultiLineStringArray
+    | MultiPolygonArray
+    | MixedGeometryArray
+    | GeometryCollectionArray
+): ...
 def from_wkb(
     input: _ArrowArrayExportable,
 ) -> (
@@ -538,6 +558,18 @@ def from_wkb(
     | MixedGeometryArray
     | GeometryCollectionArray
 ): ...
+def from_wkt(
+    input: _ArrowArrayExportable,
+) -> (
+    PointArray
+    | LineStringArray
+    | PolygonArray
+    | MultiPointArray
+    | MultiLineStringArray
+    | MultiPolygonArray
+    | MixedGeometryArray
+    | GeometryCollectionArray
+): ...
 def to_wkb(input: _ArrowArrayExportable) -> WKBArray: ...
 def read_csv(path: str, geometry_column_name: str) -> GeoTable: ...
 def read_flatgeobuf(path: str) -> GeoTable: ...

diff --git a/python/core/src/io/ewkb.rs b/python/core/src/io/ewkb.rs
@@ -0,0 +1,76 @@
+use std::sync::Arc;
+
+use geoarrow::array::{from_arrow_array, AsGeometryArray, CoordType};
+use geoarrow::datatypes::GeoDataType;
+use geoarrow::io::geozero::FromEWKB;
+use geoarrow::GeometryArrayTrait;
+use pyo3::exceptions::PyTypeError;
+use pyo3::prelude::*;
+use pyo3::types::PyType;
+
+use crate::array::*;
+use crate::ffi::from_python::import_arrow_c_array;
+use crate::ffi::to_python::geometry_array_to_pyobject;
+
+/// Convert an Arrow BinaryArray from EWKB to its GeoArrow-native counterpart.
+#[pyfunction]
+pub fn from_ewkb(ob: &PyAny) -> PyResult<PyObject> {
+    let (array, field) = import_arrow_c_array(ob)?;
+    // TODO: need to improve crate's error handling
+    let array = from_arrow_array(&array, &field).unwrap();
+    let ref_array = array.as_ref();
+    let geo_array: Arc<dyn GeometryArrayTrait> = match array.data_type() {
+        GeoDataType::WKB => {
+            FromEWKB::from_ewkb(ref_array.as_wkb(), CoordType::Interleaved).unwrap()
+        }
+        GeoDataType::LargeWKB => {
+            FromEWKB::from_ewkb(ref_array.as_large_wkb(), CoordType::Interleaved).unwrap()
+        }
+        other => {
+            return Err(PyTypeError::new_err(format!(
+                "Unexpected array type {:?}",
+                other
+            )))
+        }
+    };
+    Python::with_gil(|py| geometry_array_to_pyobject(py, geo_array))
+}
+
+macro_rules! impl_from_ewkb {
+    ($py_array:ty, $geoarrow_array:ty) => {
+        #[pymethods]
+        impl $py_array {
+            /// Parse from EWKB
+            #[classmethod]
+            pub fn from_ewkb(_cls: &PyType, ob: &PyAny) -> PyResult<$py_array> {
+                let (array, field) = import_arrow_c_array(ob)?;
+                let array = from_arrow_array(&array, &field).unwrap();
+                let ref_array = array.as_ref();
+                match array.data_type() {
+                    GeoDataType::WKB => Ok(<$geoarrow_array>::from_ewkb(
+                        ref_array.as_wkb(),
+                        CoordType::Interleaved,
+                    )
+                    .unwrap()
+                    .into()),
+                    GeoDataType::LargeWKB => Ok(<$geoarrow_array>::from_ewkb(
+                        ref_array.as_large_wkb(),
+                        CoordType::Interleaved,
+                    )
+                    .unwrap()
+                    .into()),
+                    other => Err(PyTypeError::new_err(format!(
+                        "Unexpected array type {:?}",
+                        other
+                    ))),
+                }
+            }
+        }
+    };
+}
+
+impl_from_ewkb!(MixedGeometryArray, geoarrow::array::MixedGeometryArray<i32>);
+impl_from_ewkb!(
+    GeometryCollectionArray,
+    geoarrow::array::GeometryCollectionArray<i32>
+);
diff --git a/python/core/src/io/mod.rs b/python/core/src/io/mod.rs
@@ -1,4 +1,6 @@
 pub mod csv;
+pub mod ewkb;
 pub mod flatgeobuf;
 pub mod geojson;
 pub mod wkb;
+pub mod wkt;
diff --git a/python/core/src/io/wkt.rs b/python/core/src/io/wkt.rs
@@ -0,0 +1,72 @@
+use std::sync::Arc;
+
+use arrow::datatypes::DataType;
+use arrow_array::cast::AsArray;
+use geoarrow::array::CoordType;
+use geoarrow::io::geozero::FromWKT;
+use geoarrow::GeometryArrayTrait;
+use pyo3::exceptions::PyTypeError;
+use pyo3::prelude::*;
+use pyo3::types::PyType;
+
+use crate::array::*;
+use crate::ffi::from_python::import_arrow_c_array;
+use crate::ffi::to_python::geometry_array_to_pyobject;
+
+/// Convert an Arrow StringArray from WKT to its GeoArrow-native counterpart.
+#[pyfunction]
+pub fn from_wkt(ob: &PyAny) -> PyResult<PyObject> {
+    let (array, _field) = import_arrow_c_array(ob)?;
+    let geo_array: Arc<dyn GeometryArrayTrait> = match array.data_type() {
+        DataType::Utf8 => {
+            FromWKT::from_wkt(array.as_string::<i32>(), CoordType::Interleaved).unwrap()
+        }
+        DataType::LargeUtf8 => {
+            FromWKT::from_wkt(array.as_string::<i64>(), CoordType::Interleaved).unwrap()
+        }
+        other => {
+            return Err(PyTypeError::new_err(format!(
+                "Unexpected array type {:?}",
+                other
+            )))
+        }
+    };
+    Python::with_gil(|py| geometry_array_to_pyobject(py, geo_array))
+}
+
+macro_rules! impl_from_wkt {
+    ($py_array:ty, $geoarrow_array:ty) => {
+        #[pymethods]
+        impl $py_array {
+            /// Parse from WKT
+            #[classmethod]
+            pub fn from_wkt(_cls: &PyType, ob: &PyAny) -> PyResult<$py_array> {
+                let (array, _field) = import_arrow_c_array(ob)?;
+                match array.data_type() {
+                    DataType::Utf8 => Ok(<$geoarrow_array>::from_wkt(
+                        array.as_string::<i32>(),
+                        CoordType::Interleaved,
+                    )
+                    .unwrap()
+                    .into()),
+                    DataType::LargeUtf8 => Ok(<$geoarrow_array>::from_wkt(
+                        array.as_string::<i64>(),
+                        CoordType::Interleaved,
+                    )
+                    .unwrap()
+                    .into()),
+                    other => Err(PyTypeError::new_err(format!(
+                        "Unexpected array type {:?}",
+                        other
+                    ))),
+                }
+            }
+        }
+    };
+}
+
+impl_from_wkt!(MixedGeometryArray, geoarrow::array::MixedGeometryArray<i32>);
+impl_from_wkt!(
+    GeometryCollectionArray,
+    geoarrow::array::GeometryCollectionArray<i32>
+);
diff --git a/src/algorithm/geo/area.rs b/src/algorithm/geo/area.rs
@@ -1,6 +1,6 @@
 use crate::algorithm::geo::utils::zeroes;
 use crate::array::*;
-use crate::chunked_array::{chunked_try_map, ChunkedArray, ChunkedGeometryArray};
+use crate::chunked_array::{ChunkedArray, ChunkedGeometryArray};
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -180,11 +180,13 @@ impl<G: GeometryArrayTrait> Area for ChunkedGeometryArray<G> {
     type Output = Result<ChunkedArray<Float64Array>>;
 
     fn signed_area(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| chunk.as_ref().signed_area())?.try_into()
+        self.try_map(|chunk| chunk.as_ref().signed_area())?
+            .try_into()
     }
 
     fn unsigned_area(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| chunk.as_ref().unsigned_area())?.try_into()
+        self.try_map(|chunk| chunk.as_ref().unsigned_area())?
+            .try_into()
     }
 }
 

diff --git a/src/algorithm/geo/bounding_rect.rs b/src/algorithm/geo/bounding_rect.rs
@@ -1,5 +1,5 @@
 use crate::array::*;
-use crate::chunked_array::{chunked_try_map, ChunkedGeometryArray};
+use crate::chunked_array::ChunkedGeometryArray;
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -117,6 +117,7 @@ impl<G: GeometryArrayTrait> BoundingRect for ChunkedGeometryArray<G> {
     type Output = Result<ChunkedGeometryArray<RectArray>>;
 
     fn bounding_rect(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| chunk.as_ref().bounding_rect())?.try_into()
+        self.try_map(|chunk| chunk.as_ref().bounding_rect())?
+            .try_into()
     }
 }
diff --git a/src/algorithm/geo/center.rs b/src/algorithm/geo/center.rs
@@ -1,5 +1,5 @@
 use crate::array::*;
-use crate::chunked_array::{chunked_try_map, ChunkedGeometryArray};
+use crate::chunked_array::ChunkedGeometryArray;
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -91,6 +91,6 @@ impl<G: GeometryArrayTrait> Center for ChunkedGeometryArray<G> {
     type Output = Result<ChunkedGeometryArray<PointArray>>;
 
     fn center(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| chunk.as_ref().center())?.try_into()
+        self.try_map(|chunk| chunk.as_ref().center())?.try_into()
     }
 }
diff --git a/src/algorithm/geo/centroid.rs b/src/algorithm/geo/centroid.rs
@@ -1,5 +1,5 @@
 use crate::array::*;
-use crate::chunked_array::{chunked_try_map, ChunkedGeometryArray};
+use crate::chunked_array::ChunkedGeometryArray;
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -137,6 +137,6 @@ impl<G: GeometryArrayTrait> Centroid for ChunkedGeometryArray<G> {
     type Output = Result<ChunkedGeometryArray<PointArray>>;
 
     fn centroid(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| chunk.as_ref().centroid())?.try_into()
+        self.try_map(|chunk| chunk.as_ref().centroid())?.try_into()
     }
 }
diff --git a/src/algorithm/geo/chaikin_smoothing.rs b/src/algorithm/geo/chaikin_smoothing.rs
@@ -46,7 +46,7 @@ macro_rules! impl_chunked {
     ($chunked_array:ty) => {
         impl<O: OffsetSizeTrait> ChaikinSmoothing for $chunked_array {
             fn chaikin_smoothing(&self, n_iterations: u32) -> Self {
-                chunked_map(self, |chunk| chunk.chaikin_smoothing(n_iterations.into()))
+                self.map(|chunk| chunk.chaikin_smoothing(n_iterations.into()))
                     .try_into()
                     .unwrap()
             }

diff --git a/src/algorithm/geo/convex_hull.rs b/src/algorithm/geo/convex_hull.rs
@@ -1,5 +1,5 @@
 use crate::array::*;
-use crate::chunked_array::{chunked_try_map, ChunkedGeometryArray};
+use crate::chunked_array::ChunkedGeometryArray;
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -129,7 +129,8 @@ impl<O: OffsetSizeTrait, G: GeometryArrayTrait> ConvexHull<O> for ChunkedGeometr
     type Output = Result<ChunkedGeometryArray<PolygonArray<O>>>;
 
     fn convex_hull(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| chunk.as_ref().convex_hull())?.try_into()
+        self.try_map(|chunk| chunk.as_ref().convex_hull())?
+            .try_into()
     }
 }
 

diff --git a/src/algorithm/geo/densify.rs b/src/algorithm/geo/densify.rs
@@ -86,7 +86,7 @@ macro_rules! impl_chunked {
             type Output = $struct_name;
 
             fn densify(&self, max_distance: f64) -> Self::Output {
-                chunked_map(self, |chunk| chunk.densify(max_distance))
+                self.map(|chunk| chunk.densify(max_distance))
                     .try_into()
                     .unwrap()
             }

diff --git a/src/algorithm/geo/dimensions.rs b/src/algorithm/geo/dimensions.rs
@@ -1,5 +1,5 @@
 use crate::array::*;
-use crate::chunked_array::{chunked_try_map, ChunkedArray, ChunkedGeometryArray};
+use crate::chunked_array::{ChunkedArray, ChunkedGeometryArray};
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -125,6 +125,7 @@ impl<G: GeometryArrayTrait> HasDimensions for ChunkedGeometryArray<G> {
     type Output = Result<ChunkedArray<BooleanArray>>;
 
     fn is_empty(&self) -> Self::Output {
-        chunked_try_map(self, |chunk| HasDimensions::is_empty(&chunk.as_ref()))?.try_into()
+        self.try_map(|chunk| HasDimensions::is_empty(&chunk.as_ref()))?
+            .try_into()
     }
 }
diff --git a/src/algorithm/geo/euclidean_length.rs b/src/algorithm/geo/euclidean_length.rs
@@ -1,6 +1,6 @@
 use crate::algorithm::geo::utils::zeroes;
 use crate::array::*;
-use crate::chunked_array::{chunked_map, ChunkedArray, ChunkedGeometryArray};
+use crate::chunked_array::{ChunkedArray, ChunkedGeometryArray};
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
 use crate::GeometryArrayTrait;
@@ -130,7 +130,7 @@ macro_rules! chunked_impl {
             type Output = Result<ChunkedArray<Float64Array>>;
 
             fn euclidean_length(&self) -> Self::Output {
-                chunked_map(self, |chunk| chunk.euclidean_length()).try_into()
+                self.map(|chunk| chunk.euclidean_length()).try_into()
             }
         }
     };