infer dimensions in capacity

src/algorithm/geo/minimum_rotated_rect.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use crate::array::polygon::PolygonCapacity;
 use crate::array::*;
 use crate::chunked_array::{ChunkedGeometryArray, ChunkedGeometryArrayTrait, ChunkedPolygonArray};
@@ -52,7 +54,12 @@ impl<O: OffsetSizeTrait> MinimumRotatedRect<O> for PointArray<2> {
         // Each output polygon has exactly 5 coordinates
         let coord_capacity = ring_capacity * 5;
 
-        let capacity = PolygonCapacity::new(coord_capacity, ring_capacity, geom_capacity);
+        let capacity = PolygonCapacity::new(
+            coord_capacity,
+            ring_capacity,
+            geom_capacity,
+            HashSet::from_iter([2]),
+        );
 
         let mut output_array = PolygonBuilder::with_capacity(capacity);
 
@@ -84,7 +91,12 @@ macro_rules! iter_geo_impl {
                 // Each output polygon has exactly 5 coordinates
                 let coord_capacity = ring_capacity * 5;
 
-                let capacity = PolygonCapacity::new(coord_capacity, ring_capacity, geom_capacity);
+                let capacity = PolygonCapacity::new(
+                    coord_capacity,
+                    ring_capacity,
+                    geom_capacity,
+                    HashSet::from_iter([2]),
+                );
 
                 let mut output_array = PolygonBuilder::with_capacity(capacity);
 

src/algorithm/native/cast.rs

@@ -5,6 +5,7 @@
 //! todo: have a set of "fast cast" functions, where you first try to fast cast and fall back to
 //! slower copies if necessary. Can check that the coord type of the input and output is the same.
 
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use arrow_array::OffsetSizeTrait;
@@ -77,8 +78,11 @@ impl Cast for PointArray<2> {
                 Ok(Arc::new(builder.finish()))
             }
             MultiPoint(ct, Dimension::XY) => {
-                let capacity =
-                    MultiPointCapacity::new(self.buffer_lengths(), self.buffer_lengths());
+                let capacity = MultiPointCapacity::new(
+                    self.buffer_lengths(),
+                    self.buffer_lengths(),
+                    HashSet::from_iter([2]),
+                );
                 let mut builder = MultiPointBuilder::<i32, 2>::with_capacity_and_options(
                     capacity,
                     *ct,
@@ -89,8 +93,11 @@ impl Cast for PointArray<2> {
                 Ok(Arc::new(builder.finish()))
             }
             LargeMultiPoint(ct, Dimension::XY) => {
-                let capacity =
-                    MultiPointCapacity::new(self.buffer_lengths(), self.buffer_lengths());
+                let capacity = MultiPointCapacity::new(
+                    self.buffer_lengths(),
+                    self.buffer_lengths(),
+                    HashSet::from_iter([2]),
+                );
                 let mut builder = MultiPointBuilder::<i64, 2>::with_capacity_and_options(
                     capacity,
                     *ct,
@@ -510,6 +517,7 @@ impl<O: OffsetSizeTrait> Cast for MultiLineStringArray<O, 2> {
                 let capacity = LineStringCapacity {
                     coord_capacity: existing_capacity.coord_capacity,
                     geom_capacity: existing_capacity.ring_capacity,
+                    dimensions: HashSet::from_iter([2]),
                 };
                 let mut builder = LineStringBuilder::<i32, 2>::with_capacity_and_options(
                     capacity,
@@ -530,6 +538,7 @@ impl<O: OffsetSizeTrait> Cast for MultiLineStringArray<O, 2> {
                 let capacity = LineStringCapacity {
                     coord_capacity: existing_capacity.coord_capacity,
                     geom_capacity: existing_capacity.ring_capacity,
+                    dimensions: HashSet::from_iter([2]),
                 };
                 let mut builder = LineStringBuilder::<i64, 2>::with_capacity_and_options(
                     capacity,
@@ -621,6 +630,7 @@ impl<O: OffsetSizeTrait> Cast for MultiPolygonArray<O, 2> {
                     coord_capacity: existing_capacity.coord_capacity,
                     ring_capacity: existing_capacity.ring_capacity,
                     geom_capacity: existing_capacity.polygon_capacity,
+                    dimensions: HashSet::from_iter([2]),
                 };
                 let mut builder = PolygonBuilder::<i32, 2>::with_capacity_and_options(
                     capacity,
@@ -642,6 +652,7 @@ impl<O: OffsetSizeTrait> Cast for MultiPolygonArray<O, 2> {
                     coord_capacity: existing_capacity.coord_capacity,
                     ring_capacity: existing_capacity.ring_capacity,
                     geom_capacity: existing_capacity.polygon_capacity,
+                    dimensions: HashSet::from_iter([2]),
                 };
                 let mut builder = PolygonBuilder::<i64, 2>::with_capacity_and_options(
                     capacity,

src/array/geometrycollection/array.rs

@@ -87,7 +87,7 @@ impl<O: OffsetSizeTrait, const D: usize> GeometryCollectionArray<O, D> {
     /// The number of bytes occupied by this array.
     pub fn num_bytes(&self) -> usize {
         let validity_len = self.validity().map(|v| v.buffer().len()).unwrap_or(0);
-        validity_len + self.buffer_lengths().num_bytes::<O>()
+        validity_len + self.buffer_lengths().num_bytes::<O>(D)
     }
 }
 

src/array/geometrycollection/capacity.rs

@@ -13,7 +13,7 @@ use crate::geo_traits::{
 /// [`GeometryCollectionArray`][crate::array::GeometryCollectionArray].
 ///
 /// This can be used to reduce allocations by allocating once for exactly the array size you need.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone)]
 pub struct GeometryCollectionCapacity {
     pub(crate) mixed_capacity: MixedCapacity,
     pub(crate) geom_capacity: usize,
@@ -141,10 +141,10 @@ impl GeometryCollectionCapacity {
     }
 
     /// The number of bytes an array with this capacity would occupy.
-    pub fn num_bytes<O: OffsetSizeTrait>(&self) -> usize {
+    pub fn num_bytes<O: OffsetSizeTrait>(&self, dim: usize) -> usize {
         let offsets_byte_width = if O::IS_LARGE { 8 } else { 4 };
         let num_offsets = self.geom_capacity;
-        (offsets_byte_width * num_offsets) + self.mixed_capacity.num_bytes::<O>()
+        (offsets_byte_width * num_offsets) + self.mixed_capacity.num_bytes::<O>(dim)
     }
 }
 

src/array/linestring/array.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use crate::algorithm::native::eq::offset_buffer_eq;
@@ -133,13 +133,17 @@ impl<O: OffsetSizeTrait, const D: usize> LineStringArray<O, D> {
 
     /// The lengths of each buffer contained in this array.
     pub fn buffer_lengths(&self) -> LineStringCapacity {
-        LineStringCapacity::new(self.geom_offsets.last().to_usize().unwrap(), self.len())
+        LineStringCapacity::new(
+            self.geom_offsets.last().to_usize().unwrap(),
+            self.len(),
+            HashSet::from_iter([D]),
+        )
     }
 
     /// The number of bytes occupied by this array.
     pub fn num_bytes(&self) -> usize {
         let validity_len = self.validity().map(|v| v.buffer().len()).unwrap_or(0);
-        validity_len + self.buffer_lengths().num_bytes::<O>()
+        validity_len + self.buffer_lengths().num_bytes::<O>(D)
     }
 }
 

src/array/linestring/builder.rs

@@ -309,7 +309,7 @@ impl<O: OffsetSizeTrait, const D: usize> GeometryArrayBuilder for LineStringBuil
         coord_type: CoordType,
         metadata: Arc<ArrayMetadata>,
     ) -> Self {
-        let capacity = LineStringCapacity::new(0, geom_capacity);
+        let capacity = LineStringCapacity::new(0, geom_capacity, Default::default());
         Self::with_capacity_and_options(capacity, coord_type, metadata)
     }
 

src/array/linestring/capacity.rs

@@ -1,4 +1,5 @@
-use std::ops::Add;
+use std::collections::HashSet;
+use std::ops::{Add, AddAssign};
 
 use arrow_array::OffsetSizeTrait;
 
@@ -8,24 +9,26 @@ use crate::geo_traits::{GeometryTrait, GeometryType, LineStringTrait};
 /// A counter for the buffer sizes of a [`LineStringArray`][crate::array::LineStringArray].
 ///
 /// This can be used to reduce allocations by allocating once for exactly the array size you need.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone)]
 pub struct LineStringCapacity {
     pub(crate) coord_capacity: usize,
     pub(crate) geom_capacity: usize,
+    pub(crate) dimensions: HashSet<usize>,
 }
 
 impl LineStringCapacity {
     /// Create a new capacity with known sizes.
-    pub fn new(coord_capacity: usize, geom_capacity: usize) -> Self {
+    pub fn new(coord_capacity: usize, geom_capacity: usize, dimensions: HashSet<usize>) -> Self {
         Self {
             coord_capacity,
             geom_capacity,
+            dimensions,
         }
     }
 
     /// Create a new empty capacity.
     pub fn new_empty() -> Self {
-        Self::new(0, 0)
+        Self::new(0, 0, HashSet::new())
     }
 
     /// Return `true` if the capacity is empty.
@@ -45,6 +48,7 @@ impl LineStringCapacity {
     #[inline]
     fn add_valid_line_string(&mut self, line_string: &impl LineStringTrait) {
         self.coord_capacity += line_string.num_coords();
+        self.dimensions.insert(line_string.dim());
     }
 
     #[inline]
@@ -81,10 +85,10 @@ impl LineStringCapacity {
     }
 
     /// The number of bytes an array with this capacity would occupy.
-    pub fn num_bytes<O: OffsetSizeTrait>(&self) -> usize {
+    pub fn num_bytes<O: OffsetSizeTrait>(&self, dim: usize) -> usize {
         let offsets_byte_width = if O::IS_LARGE { 8 } else { 4 };
         let num_offsets = self.geom_capacity;
-        (offsets_byte_width * num_offsets) + (self.coord_capacity * 2 * 8)
+        (offsets_byte_width * num_offsets) + (self.coord_capacity * dim * 8)
     }
 }
 
@@ -98,8 +102,16 @@ impl Add for LineStringCapacity {
     type Output = Self;
 
     fn add(self, rhs: Self) -> Self::Output {
-        let coord_capacity = self.coord_capacity + rhs.coord_capacity;
-        let geom_capacity = self.geom_capacity + rhs.geom_capacity;
-        Self::new(coord_capacity, geom_capacity)
+        let mut new = self.clone();
+        new += rhs;
+        new
+    }
+}
+
+impl AddAssign for LineStringCapacity {
+    fn add_assign(&mut self, rhs: Self) {
+        self.coord_capacity += rhs.coord_capacity();
+        self.geom_capacity += rhs.geom_capacity();
+        self.dimensions.extend(rhs.dimensions);
     }
 }

src/array/mixed/array.rs

@@ -265,7 +265,7 @@ impl<O: OffsetSizeTrait, const D: usize> MixedGeometryArray<O, D> {
 
     /// The number of bytes occupied by this array.
     pub fn num_bytes(&self) -> usize {
-        self.buffer_lengths().num_bytes::<O>()
+        self.buffer_lengths().num_bytes::<O>(D)
     }
 }
 

src/array/mixed/capacity.rs

@@ -13,7 +13,7 @@ use crate::geo_traits::*;
 /// A counter for the buffer sizes of a [`MixedGeometryArray`][crate::array::MixedGeometryArray].
 ///
 /// This can be used to reduce allocations by allocating once for exactly the array size you need.
-#[derive(Default, Debug, Clone, Copy)]
+#[derive(Default, Debug, Clone)]
 pub struct MixedCapacity {
     /// Simple: just the total number of points, nulls included
     pub(crate) point: usize,
@@ -77,28 +77,28 @@ impl MixedCapacity {
         total
     }
 
-    pub fn point_capacity(&self) -> usize {
-        self.point
+    pub fn point_capacity(&self) -> &usize {
+        &self.point
     }
 
-    pub fn line_string_capacity(&self) -> LineStringCapacity {
-        self.line_string
+    pub fn line_string_capacity(&self) -> &LineStringCapacity {
+        &self.line_string
     }
 
-    pub fn polygon_capacity(&self) -> PolygonCapacity {
-        self.polygon
+    pub fn polygon_capacity(&self) -> &PolygonCapacity {
+        &self.polygon
     }
 
-    pub fn multi_point_capacity(&self) -> MultiPointCapacity {
-        self.multi_point
+    pub fn multi_point_capacity(&self) -> &MultiPointCapacity {
+        &self.multi_point
     }
 
-    pub fn multi_line_string_capacity(&self) -> MultiLineStringCapacity {
-        self.multi_line_string
+    pub fn multi_line_string_capacity(&self) -> &MultiLineStringCapacity {
+        &self.multi_line_string
     }
 
-    pub fn multi_polygon_capacity(&self) -> MultiPolygonCapacity {
-        self.multi_polygon
+    pub fn multi_polygon_capacity(&self) -> &MultiPolygonCapacity {
+        &self.multi_polygon
     }
 
     pub fn point_compatible(&self) -> bool {
@@ -221,13 +221,13 @@ impl MixedCapacity {
     }
 
     /// The number of bytes an array with this capacity would occupy.
-    pub fn num_bytes<O: OffsetSizeTrait>(&self) -> usize {
-        let mut count = self.point * 2 * 8;
-        count += self.line_string.num_bytes::<O>();
-        count += self.polygon.num_bytes::<O>();
-        count += self.multi_point.num_bytes::<O>();
-        count += self.multi_line_string.num_bytes::<O>();
-        count += self.multi_polygon.num_bytes::<O>();
+    pub fn num_bytes<O: OffsetSizeTrait>(&self, dim: usize) -> usize {
+        let mut count = self.point * dim * 8;
+        count += self.line_string.num_bytes::<O>(dim);
+        count += self.polygon.num_bytes::<O>(dim);
+        count += self.multi_point.num_bytes::<O>(dim);
+        count += self.multi_line_string.num_bytes::<O>(dim);
+        count += self.multi_polygon.num_bytes::<O>(dim);
         count
     }
 }

src/array/mod.rs

@@ -14,7 +14,7 @@ pub use mixed::{MixedCapacity, MixedGeometryArray, MixedGeometryBuilder};
 pub use multilinestring::{MultiLineStringArray, MultiLineStringBuilder, MultiLineStringCapacity};
 pub use multipoint::{MultiPointArray, MultiPointBuilder, MultiPointCapacity};
 pub use multipolygon::{MultiPolygonArray, MultiPolygonBuilder, MultiPolygonCapacity};
-pub use point::{PointArray, PointBuilder};
+pub use point::{PointArray, PointBuilder, PointCapacity};
 pub use polygon::{PolygonArray, PolygonBuilder, PolygonCapacity};
 pub use rect::{RectArray, RectBuilder};
 

src/array/multilinestring/array.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use crate::algorithm::native::eq::offset_buffer_eq;
@@ -170,13 +170,14 @@ impl<O: OffsetSizeTrait, const D: usize> MultiLineStringArray<O, D> {
             self.ring_offsets.last().to_usize().unwrap(),
             self.geom_offsets.last().to_usize().unwrap(),
             self.len(),
+            HashSet::from_iter([D]),
         )
     }
 
     /// The number of bytes occupied by this array.
     pub fn num_bytes(&self) -> usize {
         let validity_len = self.validity().map(|v| v.buffer().len()).unwrap_or(0);
-        validity_len + self.buffer_lengths().num_bytes::<O>()
+        validity_len + self.buffer_lengths().num_bytes::<O>(D)
     }
 }
 

src/array/multilinestring/builder.rs

@@ -395,7 +395,7 @@ impl<O: OffsetSizeTrait, const D: usize> GeometryArrayBuilder for MultiLineStrin
         coord_type: CoordType,
         metadata: Arc<ArrayMetadata>,
     ) -> Self {
-        let capacity = MultiLineStringCapacity::new(0, 0, geom_capacity);
+        let capacity = MultiLineStringCapacity::new(0, 0, geom_capacity, Default::default());
         Self::with_capacity_and_options(capacity, coord_type, metadata)
     }