progress

geoarrow · Nov 14, 2024 · fa52e96 · fa52e96
1 parent 3190cee
commit fa52e96
Show file tree

Hide file tree

Showing 25 changed files with 335 additions and 136 deletions.
diff --git a/rust/geoarrow/src/algorithm/native/downcast.rs b/rust/geoarrow/src/algorithm/native/downcast.rs
@@ -144,7 +144,7 @@ impl Downcast for MultiPointArray<2> {
     fn downcast(&self, small_offsets: bool) -> Self::Output {
         // Note: this won't allow a downcast for empty MultiPoints
         if *self.geom_offsets.last() as usize == self.len() {
-            return Arc::new(PointArray::new(
+            return Arc::new(PointArray::<2>::new(
                 self.coords.clone(),
                 self.validity.clone(),
                 self.metadata(),
@@ -173,7 +173,7 @@ impl Downcast for MultiLineStringArray<2> {
 
     fn downcast(&self, small_offsets: bool) -> Self::Output {
         if *self.geom_offsets.last() as usize == self.len() {
-            return Arc::new(LineStringArray::new(
+            return Arc::new(LineStringArray::<2>::new(
                 self.coords.clone(),
                 self.ring_offsets.clone(),
                 self.validity.clone(),
@@ -203,7 +203,7 @@ impl Downcast for MultiPolygonArray<2> {
 
     fn downcast(&self, small_offsets: bool) -> Self::Output {
         if *self.geom_offsets.last() as usize == self.len() {
-            return Arc::new(PolygonArray::new(
+            return Arc::new(PolygonArray::<2>::new(
                 self.coords.clone(),
                 self.polygon_offsets.clone(),
                 self.ring_offsets.clone(),

diff --git a/rust/geoarrow/src/array/coord/combined/array.rs b/rust/geoarrow/src/array/coord/combined/array.rs
@@ -239,7 +239,8 @@ mod test {
         ));
 
         let coords2 = vec![0., 3., 1., 4., 2., 5.];
-        let buf2 = CoordBuffer::Interleaved(coords2.try_into()?);
+        let buf2 =
+            CoordBuffer::Interleaved(InterleavedCoordBuffer::new(coords2.into(), Dimension::XY));
 
         assert_eq!(buf1, buf2);
         Ok(())
@@ -253,7 +254,9 @@ mod test {
         let buf1 = CoordBuffer::Separated((x1, y1).try_into()?).slice(1, 1);
 
         let coords2 = vec![0., 3., 1., 4., 2., 5.];
-        let buf2 = CoordBuffer::Interleaved(coords2.try_into()?).slice(1, 1);
+        let buf2 =
+            CoordBuffer::Interleaved(InterleavedCoordBuffer::new(coords2.into(), Dimension::XY))
+                .slice(1, 1);
 
         assert_eq!(buf1, buf2);
         Ok(())

diff --git a/rust/geoarrow/src/array/coord/combined/builder.rs b/rust/geoarrow/src/array/coord/combined/builder.rs
@@ -85,21 +85,64 @@ impl CoordBufferBuilder {
         }
     }
 
-    // TODO: how should this handle coords that don't have the same dimension D?
-    pub fn push_coord(&mut self, point: &impl CoordTrait<T = f64>) -> Result<()> {
+    /// Push a new coord onto the end of this coordinate buffer
+    ///
+    /// ## Panics
+    ///
+    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
+    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) {
         match self {
-            CoordBufferBuilder::Interleaved(cb) => cb.push_coord(point),
-            CoordBufferBuilder::Separated(cb) => cb.push_coord(point),
+            CoordBufferBuilder::Interleaved(cb) => cb.push_coord(coord),
+            CoordBufferBuilder::Separated(cb) => cb.push_coord(coord),
         }
     }
 
-    // TODO: how should this handle coords that don't have the same dimension D?
+    /// Push a new coord onto the end of this coordinate buffer
+    ///
+    /// ## Errors
+    ///
+    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
+    pub fn try_push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> Result<()> {
+        match self {
+            CoordBufferBuilder::Interleaved(cb) => cb.try_push_coord(coord),
+            CoordBufferBuilder::Separated(cb) => cb.try_push_coord(coord),
+        }
+    }
+
+    /// Push a valid coordinate with NaN values
+    ///
+    /// Used in the case of point and rect arrays, where a `null` array value still needs to have
+    /// space allocated for it.
+    pub fn push_nan_coord(&mut self) {
+        match self {
+            CoordBufferBuilder::Interleaved(cb) => cb.push_nan_coord(),
+            CoordBufferBuilder::Separated(cb) => cb.push_nan_coord(),
+        }
+    }
+
+    /// Push a new point onto the end of this coordinate buffer
+    ///
+    /// ## Panics
+    ///
+    /// - If the added point does not have the same dimension as the coordinate buffer.
     pub fn push_point(&mut self, point: &impl PointTrait<T = f64>) {
         match self {
             CoordBufferBuilder::Interleaved(cb) => cb.push_point(point),
             CoordBufferBuilder::Separated(cb) => cb.push_point(point),
         }
     }
+
+    /// Push a new point onto the end of this coordinate buffer
+    ///
+    /// ## Errors
+    ///
+    /// - If the added point does not have the same dimension as the coordinate buffer.
+    pub fn try_push_point(&mut self, point: &impl PointTrait<T = f64>) -> Result<()> {
+        match self {
+            CoordBufferBuilder::Interleaved(cb) => cb.try_push_point(point),
+            CoordBufferBuilder::Separated(cb) => cb.try_push_point(point),
+        }
+    }
 }
 
 impl From<CoordBufferBuilder> for CoordBuffer {

diff --git a/rust/geoarrow/src/array/coord/interleaved/builder.rs b/rust/geoarrow/src/array/coord/interleaved/builder.rs
@@ -74,7 +74,21 @@ impl InterleavedCoordBufferBuilder {
         self.len() == 0
     }
 
-    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> Result<()> {
+    /// Push a new coord onto the end of this coordinate buffer
+    ///
+    /// ## Panics
+    ///
+    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
+    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) {
+        self.try_push_coord(coord).unwrap()
+    }
+
+    /// Push a new coord onto the end of this coordinate buffer
+    ///
+    /// ## Errors
+    ///
+    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
+    pub fn try_push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> Result<()> {
         // TODO: should check xyz/zym
         if coord.dim().size() != self.dim.size() {
             return Err(GeoArrowError::General(
@@ -90,14 +104,37 @@ impl InterleavedCoordBufferBuilder {
         Ok(())
     }
 
+    /// Push a valid coordinate with NaN values
+    ///
+    /// Used in the case of point and rect arrays, where a `null` array value still needs to have
+    /// space allocated for it.
+    pub fn push_nan_coord(&mut self) {
+        for _ in 0..self.dim.size() {
+            self.coords.push(f64::NAN);
+        }
+    }
+
+    /// Push a new point onto the end of this coordinate buffer
+    ///
+    /// ## Panics
+    ///
+    /// - If the added point does not have the same dimension as the coordinate buffer.
     pub fn push_point(&mut self, point: &impl PointTrait<T = f64>) {
+        self.try_push_point(point).unwrap()
+    }
+
+    /// Push a new point onto the end of this coordinate buffer
+    ///
+    /// ## Errors
+    ///
+    /// - If the added point does not have the same dimension as the coordinate buffer.
+    pub fn try_push_point(&mut self, point: &impl PointTrait<T = f64>) -> Result<()> {
         if let Some(coord) = point.coord() {
-            self.push_coord(&coord);
+            self.try_push_coord(&coord)?;
         } else {
-            for _ in 0..self.dim.size() {
-                self.coords.push(f64::NAN);
-            }
-        }
+            self.push_nan_coord();
+        };
+        Ok(())
     }
 
     pub fn from_coords<G: CoordTrait<T = f64>>(coords: &[G], dim: Dimension) -> Result<Self> {

diff --git a/rust/geoarrow/src/array/coord/separated/array.rs b/rust/geoarrow/src/array/coord/separated/array.rs
@@ -51,11 +51,6 @@ impl SeparatedCoordBuffer {
         Ok(Self { buffers, dim })
     }
 
-    /// Access the underlying coordinate buffers.
-    pub(crate) fn coords(&self) -> &[ScalarBuffer<f64>; 4] {
-        &self.buffers
-    }
-
     pub fn dim(&self) -> Dimension {
         self.dim
     }

diff --git a/rust/geoarrow/src/array/coord/separated/builder.rs b/rust/geoarrow/src/array/coord/separated/builder.rs
@@ -112,7 +112,21 @@ impl SeparatedCoordBufferBuilder {
         self.len() == 0
     }
 
-    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> Result<()> {
+    /// Push a new coord onto the end of this coordinate buffer
+    ///
+    /// ## Panics
+    ///
+    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
+    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) {
+        self.try_push_coord(coord).unwrap()
+    }
+
+    /// Push a new coord onto the end of this coordinate buffer
+    ///
+    /// ## Errors
+    ///
+    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
+    pub fn try_push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> Result<()> {
         // TODO: should check xyz/zym
         if coord.dim().size() != self.dim.size() {
             return Err(GeoArrowError::General(
@@ -128,20 +142,43 @@ impl SeparatedCoordBufferBuilder {
         Ok(())
     }
 
+    /// Push a valid coordinate with NaN values
+    ///
+    /// Used in the case of point and rect arrays, where a `null` array value still needs to have
+    /// space allocated for it.
+    pub fn push_nan_coord(&mut self) {
+        for i in 0..self.dim.size() {
+            self.buffers[i].push(f64::NAN);
+        }
+    }
+
+    /// Push a new point onto the end of this coordinate buffer
+    ///
+    /// ## Panics
+    ///
+    /// - If the added point does not have the same dimension as the coordinate buffer.
     pub fn push_point(&mut self, point: &impl PointTrait<T = f64>) {
+        self.try_push_point(point).unwrap()
+    }
+
+    /// Push a new point onto the end of this coordinate buffer
+    ///
+    /// ## Errors
+    ///
+    /// - If the added point does not have the same dimension as the coordinate buffer.
+    pub fn try_push_point(&mut self, point: &impl PointTrait<T = f64>) -> Result<()> {
         if let Some(coord) = point.coord() {
-            self.push_coord(&coord);
+            self.try_push_coord(&coord)?;
         } else {
-            for i in 0..self.dim.size() {
-                self.buffers[i].push(f64::NAN);
-            }
-        }
+            self.push_nan_coord();
+        };
+        Ok(())
     }
 
     pub fn from_coords<G: CoordTrait<T = f64>>(coords: &[G], dim: Dimension) -> Result<Self> {
         let mut buffer = SeparatedCoordBufferBuilder::with_capacity(coords.len(), dim);
         for coord in coords {
-            buffer.push_coord(coord);
+            buffer.try_push_coord(coord)?;
         }
         Ok(buffer)
     }

diff --git a/rust/geoarrow/src/array/linestring/array.rs b/rust/geoarrow/src/array/linestring/array.rs
@@ -353,7 +353,7 @@ impl<const D: usize> TryFrom<&GenericListArray<i32>> for LineStringArray<D> {
     type Error = GeoArrowError;
 
     fn try_from(value: &GenericListArray<i32>) -> Result<Self> {
-        let coords: CoordBuffer = value.values().as_ref().try_into()?;
+        let coords = CoordBuffer::from_arrow(value.values().as_ref(), D.try_into()?)?;
         let geom_offsets = value.offsets();
         let validity = value.nulls();
 
@@ -370,7 +370,7 @@ impl<const D: usize> TryFrom<&GenericListArray<i64>> for LineStringArray<D> {
     type Error = GeoArrowError;
 
     fn try_from(value: &GenericListArray<i64>) -> Result<Self> {
-        let coords: CoordBuffer = value.values().as_ref().try_into()?;
+        let coords = CoordBuffer::from_arrow(value.values().as_ref(), D.try_into()?)?;
         let geom_offsets = offsets_buffer_i64_to_i32(value.offsets())?;
         let validity = value.nulls();
 

diff --git a/rust/geoarrow/src/array/linestring/builder.rs b/rust/geoarrow/src/array/linestring/builder.rs
@@ -22,7 +22,7 @@ use std::sync::Arc;
 pub struct LineStringBuilder<const D: usize> {
     metadata: Arc<ArrayMetadata>,
 
-    pub(crate) coords: CoordBufferBuilder<D>,
+    pub(crate) coords: CoordBufferBuilder,
 
     /// Offsets into the coordinate array where each geometry starts
     pub(crate) geom_offsets: OffsetsBuilder<i32>,
@@ -52,12 +52,18 @@ impl<const D: usize> LineStringBuilder<D> {
         metadata: Arc<ArrayMetadata>,
     ) -> Self {
         let coords = match coord_type {
-            CoordType::Interleaved => CoordBufferBuilder::Interleaved(
-                InterleavedCoordBufferBuilder::with_capacity(capacity.coord_capacity()),
-            ),
-            CoordType::Separated => CoordBufferBuilder::Separated(
-                SeparatedCoordBufferBuilder::with_capacity(capacity.coord_capacity()),
-            ),
+            CoordType::Interleaved => {
+                CoordBufferBuilder::Interleaved(InterleavedCoordBufferBuilder::with_capacity(
+                    capacity.coord_capacity(),
+                    D.try_into().unwrap(),
+                ))
+            }
+            CoordType::Separated => {
+                CoordBufferBuilder::Separated(SeparatedCoordBufferBuilder::with_capacity(
+                    capacity.coord_capacity(),
+                    D.try_into().unwrap(),
+                ))
+            }
         };
         Self {
             coords,
@@ -107,7 +113,7 @@ impl<const D: usize> LineStringBuilder<D> {
     /// - The validity is not `None` and its length is different from the number of geometries
     /// - if the largest geometry offset does not match the number of coordinates
     pub fn try_new(
-        coords: CoordBufferBuilder<D>,
+        coords: CoordBufferBuilder,
         geom_offsets: OffsetsBuilder<i32>,
         validity: NullBufferBuilder,
         metadata: Arc<ArrayMetadata>,
@@ -126,13 +132,7 @@ impl<const D: usize> LineStringBuilder<D> {
     }
 
     /// Extract the low-level APIs from the [`LineStringBuilder`].
-    pub fn into_inner(
-        self,
-    ) -> (
-        CoordBufferBuilder<D>,
-        OffsetsBuilder<i32>,
-        NullBufferBuilder,
-    ) {
+    pub fn into_inner(self) -> (CoordBufferBuilder, OffsetsBuilder<i32>, NullBufferBuilder) {
         (self.coords, self.geom_offsets, self.validity)
     }
 
@@ -231,7 +231,7 @@ impl<const D: usize> LineStringBuilder<D> {
         if let Some(line_string) = value {
             let num_coords = line_string.num_coords();
             for coord in line_string.coords() {
-                self.coords.push_coord(&coord);
+                self.coords.try_push_coord(&coord)?;
             }
             self.try_push_length(num_coords)?;
         } else {
@@ -265,8 +265,8 @@ impl<const D: usize> LineStringBuilder<D> {
     /// This is marked as unsafe because care must be taken to ensure that pushing raw coordinates
     /// to the array upholds the necessary invariants of the array.
     #[inline]
-    pub unsafe fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) {
-        self.coords.push_coord(coord)
+    pub unsafe fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> Result<()> {
+        self.coords.try_push_coord(coord)
     }
 
     #[inline]

diff --git a/rust/geoarrow/src/array/multilinestring/array.rs b/rust/geoarrow/src/array/multilinestring/array.rs
@@ -400,7 +400,7 @@ impl<const D: usize> TryFrom<&GenericListArray<i32>> for MultiLineStringArray<D>
         let rings_array = rings_dyn_array.as_list::<i32>();
 
         let ring_offsets = rings_array.offsets();
-        let coords: CoordBuffer = rings_array.values().as_ref().try_into()?;
+        let coords = CoordBuffer::from_arrow(rings_array.values().as_ref(), D.try_into()?)?;
 
         Ok(Self::new(
             coords,
@@ -423,7 +423,7 @@ impl<const D: usize> TryFrom<&GenericListArray<i64>> for MultiLineStringArray<D>
         let rings_array = rings_dyn_array.as_list::<i64>();
 
         let ring_offsets = offsets_buffer_i64_to_i32(rings_array.offsets())?;
-        let coords: CoordBuffer = rings_array.values().as_ref().try_into()?;
+        let coords = CoordBuffer::from_arrow(rings_array.values().as_ref(), D.try_into()?)?;
 
         Ok(Self::new(
             coords,