Use geometry traits in conversion from vector of geoms (#260)

geoarrow · Nov 21, 2023 · 32184c7 · 32184c7
1 parent 8b1b15e
commit 32184c7
Show file tree

Hide file tree

Showing 21 changed files with 275 additions and 225 deletions.
diff --git a/src/array/coord/combined/mutable.rs b/src/array/coord/combined/mutable.rs
@@ -1,4 +1,5 @@
 use crate::array::{CoordBuffer, MutableInterleavedCoordBuffer, MutableSeparatedCoordBuffer};
+use crate::geo_traits::{CoordTrait, PointTrait};
 
 #[derive(Debug, Clone)]
 pub enum MutableCoordBuffer {
@@ -60,7 +61,11 @@ impl MutableCoordBuffer {
         }
     }
 
-    pub fn push_coord(&mut self, coord: geo::Coord) {
+    pub fn push_point(&mut self, coord: impl PointTrait<T = f64>) {
+        self.push_xy(coord.x(), coord.y())
+    }
+
+    pub fn push_coord(&mut self, coord: impl CoordTrait<T = f64>) {
         match self {
             MutableCoordBuffer::Interleaved(cb) => cb.push_coord(coord),
             MutableCoordBuffer::Separated(cb) => cb.push_coord(coord),

diff --git a/src/array/coord/interleaved/array.rs b/src/array/coord/interleaved/array.rs
@@ -1,7 +1,8 @@
 use std::sync::Arc;
 
-use crate::array::CoordType;
+use crate::array::{CoordType, MutableInterleavedCoordBuffer};
 use crate::error::{GeoArrowError, Result};
+use crate::geo_traits::CoordTrait;
 use crate::scalar::InterleavedCoord;
 use crate::trait_::{GeoArrayAccessor, IntoArrow};
 use crate::GeometryArrayTrait;
@@ -177,6 +178,13 @@ impl TryFrom<Vec<f64>> for InterleavedCoordBuffer {
     }
 }
 
+impl<G: CoordTrait<T = f64>> From<Vec<G>> for InterleavedCoordBuffer {
+    fn from(other: Vec<G>) -> Self {
+        let mut_arr: MutableInterleavedCoordBuffer = other.into();
+        mut_arr.into()
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

diff --git a/src/array/coord/interleaved/mutable.rs b/src/array/coord/interleaved/mutable.rs
@@ -94,3 +94,13 @@ impl From<MutableInterleavedCoordBuffer> for InterleavedCoordBuffer {
         InterleavedCoordBuffer::new(value.coords.into())
     }
 }
+
+impl<G: CoordTrait<T = f64>> From<Vec<G>> for MutableInterleavedCoordBuffer {
+    fn from(value: Vec<G>) -> Self {
+        let mut buffer = MutableInterleavedCoordBuffer::with_capacity(value.len());
+        for coord in value {
+            buffer.push_coord(coord);
+        }
+        buffer
+    }
+}
diff --git a/src/array/coord/separated/array.rs b/src/array/coord/separated/array.rs
@@ -5,8 +5,9 @@ use arrow_array::{Array, Float64Array, StructArray};
 use arrow_buffer::{NullBuffer, ScalarBuffer};
 use arrow_schema::{DataType, Field};
 
-use crate::array::CoordType;
+use crate::array::{CoordType, MutableSeparatedCoordBuffer};
 use crate::error::{GeoArrowError, Result};
+use crate::geo_traits::CoordTrait;
 use crate::scalar::SeparatedCoord;
 use crate::trait_::{GeoArrayAccessor, IntoArrow};
 use crate::GeometryArrayTrait;
@@ -182,6 +183,13 @@ impl TryFrom<(Vec<f64>, Vec<f64>)> for SeparatedCoordBuffer {
     }
 }
 
+impl<G: CoordTrait<T = f64>> From<Vec<G>> for SeparatedCoordBuffer {
+    fn from(other: Vec<G>) -> Self {
+        let mut_arr: MutableSeparatedCoordBuffer = other.into();
+        mut_arr.into()
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

diff --git a/src/array/coord/separated/mutable.rs b/src/array/coord/separated/mutable.rs
@@ -1,4 +1,5 @@
 use crate::array::SeparatedCoordBuffer;
+use crate::geo_traits::CoordTrait;
 
 #[derive(Debug, Clone)]
 pub struct MutableSeparatedCoordBuffer {
@@ -68,9 +69,9 @@ impl MutableSeparatedCoordBuffer {
         self.y[i] = coord.y;
     }
 
-    pub fn push_coord(&mut self, coord: geo::Coord) {
-        self.x.push(coord.x);
-        self.y.push(coord.y);
+    pub fn push_coord(&mut self, coord: impl CoordTrait<T = f64>) {
+        self.x.push(coord.x());
+        self.y.push(coord.y());
     }
 
     pub fn set_xy(&mut self, i: usize, x: f64, y: f64) {
@@ -103,3 +104,13 @@ impl From<MutableSeparatedCoordBuffer> for SeparatedCoordBuffer {
         SeparatedCoordBuffer::new(value.x.into(), value.y.into())
     }
 }
+
+impl<G: CoordTrait<T = f64>> From<Vec<G>> for MutableSeparatedCoordBuffer {
+    fn from(value: Vec<G>) -> Self {
+        let mut buffer = MutableSeparatedCoordBuffer::with_capacity(value.len());
+        for coord in value {
+            buffer.push_coord(coord);
+        }
+        buffer
+    }
+}
diff --git a/src/array/linestring/array.rs b/src/array/linestring/array.rs
@@ -7,6 +7,7 @@ use crate::array::zip_validity::ZipValidity;
 use crate::array::{CoordBuffer, CoordType, MultiPointArray, WKBArray};
 use crate::datatypes::GeoDataType;
 use crate::error::{GeoArrowError, Result};
+use crate::geo_traits::LineStringTrait;
 use crate::scalar::LineString;
 use crate::trait_::{GeoArrayAccessor, IntoArrow};
 use crate::util::{owned_slice_offsets, owned_slice_validity};
@@ -355,33 +356,33 @@ impl TryFrom<&dyn Array> for LineStringArray<i64> {
     }
 }
 
-impl<O: OffsetSizeTrait> From<Vec<Option<geo::LineString>>> for LineStringArray<O> {
-    fn from(other: Vec<Option<geo::LineString>>) -> Self {
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<Vec<Option<G>>> for LineStringArray<O> {
+    fn from(other: Vec<Option<G>>) -> Self {
         let mut_arr: MutableLineStringArray<O> = other.into();
         mut_arr.into()
     }
 }
 
-impl<O: OffsetSizeTrait> From<Vec<geo::LineString>> for LineStringArray<O> {
-    fn from(other: Vec<geo::LineString>) -> Self {
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<Vec<G>> for LineStringArray<O> {
+    fn from(other: Vec<G>) -> Self {
         let mut_arr: MutableLineStringArray<O> = other.into();
         mut_arr.into()
     }
 }
 
-impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, Option<geo::LineString>>>
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<bumpalo::collections::Vec<'_, Option<G>>>
     for LineStringArray<O>
 {
-    fn from(other: bumpalo::collections::Vec<'_, Option<geo::LineString>>) -> Self {
+    fn from(other: bumpalo::collections::Vec<'_, Option<G>>) -> Self {
         let mut_arr: MutableLineStringArray<O> = other.into();
         mut_arr.into()
     }
 }
 
-impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, geo::LineString>>
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<bumpalo::collections::Vec<'_, G>>
     for LineStringArray<O>
 {
-    fn from(other: bumpalo::collections::Vec<'_, geo::LineString>) -> Self {
+    fn from(other: bumpalo::collections::Vec<'_, G>) -> Self {
         let mut_arr: MutableLineStringArray<O> = other.into();
         mut_arr.into()
     }

diff --git a/src/array/linestring/mutable.rs b/src/array/linestring/mutable.rs
@@ -5,11 +5,10 @@ use crate::array::{
     WKBArray,
 };
 use crate::error::{GeoArrowError, Result};
-use crate::geo_traits::{CoordTrait, LineStringTrait};
+use crate::geo_traits::LineStringTrait;
 use crate::io::wkb::reader::linestring::WKBLineString;
 use crate::scalar::WKB;
 use crate::trait_::IntoArrow;
-use crate::GeometryArrayTrait;
 use arrow_array::{Array, GenericListArray, OffsetSizeTrait};
 use arrow_buffer::NullBufferBuilder;
 use std::convert::From;
@@ -118,7 +117,7 @@ impl<O: OffsetSizeTrait> MutableLineStringArray<O> {
             let num_coords = line_string.num_coords();
             for coord_idx in 0..num_coords {
                 let coord = line_string.coord(coord_idx).unwrap();
-                self.coords.push_xy(coord.x(), coord.y());
+                self.coords.push_coord(coord);
             }
             self.try_push_length(num_coords)?;
         } else {
@@ -186,8 +185,8 @@ impl<O: OffsetSizeTrait> From<MutableLineStringArray<O>> for GenericListArray<O>
     }
 }
 
-pub(crate) fn first_pass(
-    geoms: impl Iterator<Item = Option<impl LineStringTrait>>,
+pub(crate) fn first_pass<'a>(
+    geoms: impl Iterator<Item = Option<&'a (impl LineStringTrait + 'a)>>,
     geoms_length: usize,
 ) -> (usize, usize) {
     let mut coord_capacity = 0;
@@ -200,52 +199,63 @@ pub(crate) fn first_pass(
     (coord_capacity, geom_capacity)
 }
 
-pub(crate) fn second_pass<O: OffsetSizeTrait>(
-    geoms: impl Iterator<Item = Option<impl LineStringTrait<T = f64>>>,
+pub(crate) fn second_pass<'a, O: OffsetSizeTrait>(
+    geoms: impl Iterator<Item = Option<&'a (impl LineStringTrait<T = f64> + 'a)>>,
     coord_capacity: usize,
     geom_capacity: usize,
 ) -> MutableLineStringArray<O> {
     let mut array = MutableLineStringArray::with_capacities(coord_capacity, geom_capacity);
 
     geoms
         .into_iter()
-        .try_for_each(|maybe_multi_point| array.push_line_string(maybe_multi_point.as_ref()))
+        .try_for_each(|maybe_multi_point| array.push_line_string(maybe_multi_point))
         .unwrap();
 
     array
 }
 
-impl<O: OffsetSizeTrait> From<Vec<geo::LineString>> for MutableLineStringArray<O> {
-    fn from(geoms: Vec<geo::LineString>) -> Self {
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<Vec<G>> for MutableLineStringArray<O> {
+    fn from(geoms: Vec<G>) -> Self {
         let (coord_capacity, geom_capacity) = first_pass(geoms.iter().map(Some), geoms.len());
-        second_pass(geoms.into_iter().map(Some), coord_capacity, geom_capacity)
+        second_pass(geoms.iter().map(Some), coord_capacity, geom_capacity)
     }
 }
 
-impl<O: OffsetSizeTrait> From<Vec<Option<geo::LineString>>> for MutableLineStringArray<O> {
-    fn from(geoms: Vec<Option<geo::LineString>>) -> Self {
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<Vec<Option<G>>>
+    for MutableLineStringArray<O>
+{
+    fn from(geoms: Vec<Option<G>>) -> Self {
+        let geoms_len = geoms.len();
         let (coord_capacity, geom_capacity) =
-            first_pass(geoms.iter().map(|x| x.as_ref()), geoms.len());
-        second_pass(geoms.into_iter(), coord_capacity, geom_capacity)
+            first_pass(geoms.iter().map(|x| x.as_ref()), geoms_len);
+        second_pass(
+            geoms.iter().map(|x| x.as_ref()),
+            coord_capacity,
+            geom_capacity,
+        )
     }
 }
 
-impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, geo::LineString>>
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<bumpalo::collections::Vec<'_, G>>
     for MutableLineStringArray<O>
 {
-    fn from(geoms: bumpalo::collections::Vec<'_, geo::LineString>) -> Self {
+    fn from(geoms: bumpalo::collections::Vec<'_, G>) -> Self {
         let (coord_capacity, geom_capacity) = first_pass(geoms.iter().map(Some), geoms.len());
-        second_pass(geoms.into_iter().map(Some), coord_capacity, geom_capacity)
+        second_pass(geoms.iter().map(Some), coord_capacity, geom_capacity)
     }
 }
 
-impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, Option<geo::LineString>>>
+impl<O: OffsetSizeTrait, G: LineStringTrait<T = f64>> From<bumpalo::collections::Vec<'_, Option<G>>>
     for MutableLineStringArray<O>
 {
-    fn from(geoms: bumpalo::collections::Vec<'_, Option<geo::LineString>>) -> Self {
+    fn from(geoms: bumpalo::collections::Vec<'_, Option<G>>) -> Self {
         let (coord_capacity, geom_capacity) =
             first_pass(geoms.iter().map(|x| x.as_ref()), geoms.len());
-        second_pass(geoms.into_iter(), coord_capacity, geom_capacity)
+        second_pass(
+            geoms.iter().map(|x| x.as_ref()),
+            coord_capacity,
+            geom_capacity,
+        )
     }
 }
 
@@ -262,13 +272,7 @@ impl<O: OffsetSizeTrait> TryFrom<WKBArray<O>> for MutableLineStringArray<O> {
                     .map(|wkb| wkb.to_wkb_object().into_line_string())
             })
             .collect();
-        let (coord_capacity, geom_capacity) =
-            first_pass(wkb_objects2.iter().map(|item| item.as_ref()), value.len());
-        Ok(second_pass(
-            wkb_objects2.iter().map(|item| item.as_ref()),
-            coord_capacity,
-            geom_capacity,
-        ))
+        Ok(wkb_objects2.into())
     }
 }
 

diff --git a/src/array/multilinestring/array.rs b/src/array/multilinestring/array.rs
@@ -8,6 +8,7 @@ use crate::array::zip_validity::ZipValidity;
 use crate::array::{CoordBuffer, CoordType, LineStringArray, PolygonArray, WKBArray};
 use crate::datatypes::GeoDataType;
 use crate::error::GeoArrowError;
+use crate::geo_traits::MultiLineStringTrait;
 use crate::scalar::MultiLineString;
 use crate::trait_::{GeoArrayAccessor, IntoArrow};
 use crate::util::{owned_slice_offsets, owned_slice_validity};
@@ -404,33 +405,37 @@ impl TryFrom<&dyn Array> for MultiLineStringArray<i64> {
     }
 }
 
-impl<O: OffsetSizeTrait> From<Vec<Option<geo::MultiLineString>>> for MultiLineStringArray<O> {
-    fn from(other: Vec<Option<geo::MultiLineString>>) -> Self {
+impl<O: OffsetSizeTrait, G: MultiLineStringTrait<T = f64>> From<Vec<Option<G>>>
+    for MultiLineStringArray<O>
+{
+    fn from(other: Vec<Option<G>>) -> Self {
         let mut_arr: MutableMultiLineStringArray<O> = other.into();
         mut_arr.into()
     }
 }
 
-impl<O: OffsetSizeTrait> From<Vec<geo::MultiLineString>> for MultiLineStringArray<O> {
-    fn from(other: Vec<geo::MultiLineString>) -> Self {
+impl<O: OffsetSizeTrait, G: MultiLineStringTrait<T = f64>> From<Vec<G>>
+    for MultiLineStringArray<O>
+{
+    fn from(other: Vec<G>) -> Self {
         let mut_arr: MutableMultiLineStringArray<O> = other.into();
         mut_arr.into()
     }
 }
 
-impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, Option<geo::MultiLineString>>>
-    for MultiLineStringArray<O>
+impl<O: OffsetSizeTrait, G: MultiLineStringTrait<T = f64>>
+    From<bumpalo::collections::Vec<'_, Option<G>>> for MultiLineStringArray<O>
 {
-    fn from(other: bumpalo::collections::Vec<'_, Option<geo::MultiLineString>>) -> Self {
+    fn from(other: bumpalo::collections::Vec<'_, Option<G>>) -> Self {
         let mut_arr: MutableMultiLineStringArray<O> = other.into();
         mut_arr.into()
     }
 }
 
-impl<O: OffsetSizeTrait> From<bumpalo::collections::Vec<'_, geo::MultiLineString>>
+impl<O: OffsetSizeTrait, G: MultiLineStringTrait<T = f64>> From<bumpalo::collections::Vec<'_, G>>
     for MultiLineStringArray<O>
 {
-    fn from(other: bumpalo::collections::Vec<'_, geo::MultiLineString>) -> Self {
+    fn from(other: bumpalo::collections::Vec<'_, G>) -> Self {
         let mut_arr: MutableMultiLineStringArray<O> = other.into();
         mut_arr.into()
     }