Zero-copy as_chunked

geoarrow · Mar 1, 2024 · 6db2031 · 6db2031
1 parent 6dd81b6
commit 6db2031
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 10 deletions.
diff --git a/python/core/src/algorithm/geo/frechet_distance.rs b/python/core/src/algorithm/geo/frechet_distance.rs
@@ -4,6 +4,7 @@ use crate::error::PyGeoArrowResult;
 use crate::ffi::from_python::input::AnyGeometryBroadcastInput;
 use crate::ffi::from_python::AnyGeometryInput;
 use geoarrow::algorithm::geo::{FrechetDistance, FrechetDistanceLineString};
+use geoarrow::algorithm::native::as_chunked_geometry_array;
 use geoarrow::array::{AsChunkedGeometryArray, AsGeometryArray};
 use geoarrow::datatypes::GeoDataType;
 use geoarrow::io::geo::geometry_to_geo;
@@ -16,30 +17,44 @@ pub fn frechet_distance(
     other: AnyGeometryBroadcastInput,
 ) -> PyGeoArrowResult<PyObject> {
     match (input, other) {
-        (AnyGeometryInput::Array(left), AnyGeometryBroadcastInput::Array(right)) => {
-            let result = FrechetDistance::frechet_distance(&left.as_ref(), &right.as_ref())?;
-            let result = Float64Array::from(result);
-            Python::with_gil(|py| Ok(result.into_py(py)))
-        }
-        (AnyGeometryInput::Chunked(left), AnyGeometryBroadcastInput::Chunked(right)) => {
-            let result = FrechetDistance::frechet_distance(&left.as_ref(), &right.as_ref())?;
-            let result = ChunkedFloat64Array::from(result);
-            Python::with_gil(|py| Ok(result.into_py(py)))
-        }
         (AnyGeometryInput::Array(left), AnyGeometryBroadcastInput::Scalar(right)) => {
             let scalar = geo::LineString::try_from(geometry_to_geo(&right.0))
                 .map_err(|_| PyValueError::new_err("Expected type LineString"))?;
             let result = FrechetDistanceLineString::frechet_distance(&left.as_ref(), &scalar)?;
             let result = Float64Array::from(result);
             Python::with_gil(|py| Ok(result.into_py(py)))
         }
+        (AnyGeometryInput::Array(left), AnyGeometryBroadcastInput::Array(right)) => {
+            let result = FrechetDistance::frechet_distance(&left.as_ref(), &right.as_ref())?;
+            let result = Float64Array::from(result);
+            Python::with_gil(|py| Ok(result.into_py(py)))
+        }
+        // TODO: Unknown whether this should be supported. I like "array in, array out".
+        // (AnyGeometryInput::Array(left), AnyGeometryBroadcastInput::Chunked(right)) => {
+        //     let left_chunked = as_chunked_geometry_array(&left.as_ref(), &right.chunk_lengths())?;
+        //     let result = FrechetDistance::frechet_distance(&left_chunked.as_ref(), &right.as_ref())?;
+        //     let result = Float64Array::from(result);
+        //     Python::with_gil(|py| Ok(result.into_py(py)))
+        // }
         (AnyGeometryInput::Chunked(left), AnyGeometryBroadcastInput::Scalar(right)) => {
             let scalar = geo::LineString::try_from(geometry_to_geo(&right.0))
                 .map_err(|_| PyValueError::new_err("Expected type LineString"))?;
             let result = FrechetDistanceLineString::frechet_distance(&left.as_ref(), &scalar)?;
             let result = ChunkedFloat64Array::from(result);
             Python::with_gil(|py| Ok(result.into_py(py)))
         }
+        (AnyGeometryInput::Chunked(left), AnyGeometryBroadcastInput::Array(right)) => {
+            let right_chunked = as_chunked_geometry_array(right.as_ref(), &left.chunk_lengths())?;
+            let result =
+                FrechetDistance::frechet_distance(&left.as_ref(), &right_chunked.as_ref())?;
+            let result = ChunkedFloat64Array::from(result);
+            Python::with_gil(|py| Ok(result.into_py(py)))
+        }
+        (AnyGeometryInput::Chunked(left), AnyGeometryBroadcastInput::Chunked(right)) => {
+            let result = FrechetDistance::frechet_distance(&left.as_ref(), &right.as_ref())?;
+            let result = ChunkedFloat64Array::from(result);
+            Python::with_gil(|py| Ok(result.into_py(py)))
+        }
         _ => Err(PyValueError::new_err("Unsupported input types.").into()),
     }
 }

diff --git a/src/algorithm/native/as_chunked.rs b/src/algorithm/native/as_chunked.rs
@@ -0,0 +1,28 @@
+use std::sync::Arc;
+
+use crate::chunked_array::{from_arrow_chunks, ChunkedGeometryArrayTrait};
+use crate::error::Result;
+use crate::GeometryArrayTrait;
+
+// TODO: don't go through Arc<dyn Array>
+// Update geometry array trait to put slice on the main trait
+// Put slice() on each individual array directly, and delegate to it from geom trait
+pub fn as_chunked_geometry_array(
+    array: &dyn GeometryArrayTrait,
+    chunk_lengths: &[usize],
+) -> Result<Arc<dyn ChunkedGeometryArrayTrait>> {
+    assert_eq!(array.len(), chunk_lengths.iter().sum::<usize>());
+
+    let mut new_chunks = Vec::with_capacity(chunk_lengths.len());
+    let mut offset = 0;
+    for length in chunk_lengths {
+        new_chunks.push(array.to_array_ref());
+        offset += length;
+    }
+
+    let array_refs = new_chunks
+        .iter()
+        .map(|arr| arr.as_ref())
+        .collect::<Vec<_>>();
+    from_arrow_chunks(array_refs.as_slice(), array.extension_field().as_ref())
+}
diff --git a/src/algorithm/native/mod.rs b/src/algorithm/native/mod.rs
@@ -3,6 +3,7 @@
 //! Where possible, operations on scalars are implemented in terms of [geometry
 //! traits](../../geo_traits).
 
+mod as_chunked;
 mod binary;
 pub mod bounding_rect;
 mod cast;
@@ -17,6 +18,7 @@ mod total_bounds;
 pub(crate) mod type_id;
 mod unary;
 
+pub use as_chunked::as_chunked_geometry_array;
 pub use binary::Binary;
 pub use cast::Cast;
 pub use concatenate::Concatenate;

diff --git a/src/chunked_array/chunked_array.rs b/src/chunked_array/chunked_array.rs
@@ -232,6 +232,8 @@ pub trait ChunkedGeometryArrayTrait: std::fmt::Debug + Send + Sync {
 
     /// The number of chunks in this chunked array.
     fn num_chunks(&self) -> usize;
+
+    fn chunk_lengths(&self) -> Vec<usize>;
 }
 
 impl ChunkedGeometryArrayTrait for ChunkedPointArray {
@@ -256,6 +258,10 @@ impl ChunkedGeometryArrayTrait for ChunkedPointArray {
     fn num_chunks(&self) -> usize {
         self.chunks.len()
     }
+
+    fn chunk_lengths(&self) -> Vec<usize> {
+        self.chunks.iter().map(|chunk| chunk.len()).collect()
+    }
 }
 
 macro_rules! impl_trait {
@@ -282,6 +288,10 @@ macro_rules! impl_trait {
             fn num_chunks(&self) -> usize {
                 self.chunks.len()
             }
+
+            fn chunk_lengths(&self) -> Vec<usize> {
+                self.chunks.iter().map(|chunk| chunk.len()).collect()
+            }
         }
     };
 }
@@ -317,6 +327,10 @@ impl ChunkedGeometryArrayTrait for ChunkedRectArray {
     fn num_chunks(&self) -> usize {
         self.chunks.len()
     }
+
+    fn chunk_lengths(&self) -> Vec<usize> {
+        self.chunks.iter().map(|chunk| chunk.len()).collect()
+    }
 }
 
 /// Construct