geoarrow_array/capacity/
geometrycollection.rs

1use std::ops::Add;
2
3use geo_traits::{
4    GeometryCollectionTrait, GeometryTrait, GeometryType, LineStringTrait, MultiLineStringTrait,
5    MultiPointTrait, MultiPolygonTrait, PointTrait, PolygonTrait,
6};
7use geoarrow_schema::Dimension;
8use geoarrow_schema::error::GeoArrowResult;
9use wkt::WktNum;
10
11use crate::builder::geo_trait_wrappers::{LineWrapper, RectWrapper, TriangleWrapper};
12use crate::capacity::MixedCapacity;
13
14/// A counter for the buffer sizes of a
15/// [`GeometryCollectionArray`][crate::array::GeometryCollectionArray].
16///
17/// This can be used to reduce allocations by allocating once for exactly the array size you need.
18#[derive(Debug, Clone, Copy)]
19pub struct GeometryCollectionCapacity {
20    pub(crate) mixed_capacity: MixedCapacity,
21    pub(crate) geom_capacity: usize,
22}
23
24impl GeometryCollectionCapacity {
25    /// Create a new capacity with known sizes.
26    pub fn new(mixed_capacity: MixedCapacity, geom_capacity: usize) -> Self {
27        Self {
28            mixed_capacity,
29            geom_capacity,
30        }
31    }
32
33    /// Create a new empty capacity.
34    pub fn new_empty() -> Self {
35        Self::new(MixedCapacity::new_empty(), 0)
36    }
37
38    /// Return `true` if the capacity is empty.
39    pub fn is_empty(&self) -> bool {
40        self.mixed_capacity.is_empty() && self.geom_capacity == 0
41    }
42
43    /// The geometry offset buffer capacity
44    pub fn geom_capacity(&self) -> usize {
45        self.geom_capacity
46    }
47
48    #[inline]
49    fn add_valid_point(&mut self, _geom: &impl PointTrait) {
50        self.mixed_capacity.add_point();
51    }
52
53    #[inline]
54    fn add_valid_line_string(&mut self, geom: &impl LineStringTrait) {
55        self.mixed_capacity.add_line_string(geom);
56    }
57
58    #[inline]
59    fn add_valid_polygon(&mut self, geom: &impl PolygonTrait) {
60        self.mixed_capacity.add_polygon(geom);
61    }
62
63    #[inline]
64    fn add_valid_multi_point(&mut self, geom: &impl MultiPointTrait) {
65        self.mixed_capacity.add_multi_point(geom);
66    }
67
68    #[inline]
69    fn add_valid_multi_line_string(&mut self, geom: &impl MultiLineStringTrait) {
70        self.mixed_capacity.add_multi_line_string(geom);
71    }
72
73    #[inline]
74    fn add_valid_multi_polygon(&mut self, geom: &impl MultiPolygonTrait) {
75        self.mixed_capacity.add_multi_polygon(geom);
76    }
77
78    #[inline]
79    fn add_valid_geometry_collection<T: WktNum>(
80        &mut self,
81        geom: &impl GeometryCollectionTrait<T = T>,
82    ) -> GeoArrowResult<()> {
83        for g in geom.geometries() {
84            self.mixed_capacity.add_geometry(&g)?
85        }
86        Ok(())
87    }
88
89    /// Add a Geometry to this capacity counter.
90    #[inline]
91    pub fn add_geometry<T: WktNum>(
92        &mut self,
93        geom: Option<&impl GeometryTrait<T = T>>,
94    ) -> GeoArrowResult<()> {
95        use GeometryType::*;
96        if let Some(geom) = geom {
97            match geom.as_type() {
98                Point(p) => self.add_valid_point(p),
99                LineString(p) => self.add_valid_line_string(p),
100                Polygon(p) => self.add_valid_polygon(p),
101                MultiPoint(p) => self.add_valid_multi_point(p),
102                MultiLineString(p) => self.add_valid_multi_line_string(p),
103                MultiPolygon(p) => self.add_valid_multi_polygon(p),
104                GeometryCollection(p) => self.add_valid_geometry_collection(p)?,
105                Rect(r) => self.add_valid_polygon(&RectWrapper::try_new(r)?),
106                Triangle(tri) => self.add_valid_polygon(&TriangleWrapper(tri)),
107                Line(l) => self.add_valid_line_string(&LineWrapper(l)),
108            }
109        };
110        Ok(())
111    }
112
113    /// Add a GeometryCollection to this capacity counter.
114    #[inline]
115    pub fn add_geometry_collection<'a, T: WktNum>(
116        &mut self,
117        geom: Option<&'a (impl GeometryCollectionTrait<T = T> + 'a)>,
118    ) -> GeoArrowResult<()> {
119        if let Some(geom) = geom {
120            self.add_valid_geometry_collection(geom)?;
121        }
122        self.geom_capacity += 1;
123        Ok(())
124    }
125
126    /// Create a capacity counter from an iterator of GeometryCollections.
127    pub fn from_geometry_collections<'a, T: WktNum>(
128        geoms: impl Iterator<Item = Option<&'a (impl GeometryCollectionTrait<T = T> + 'a)>>,
129    ) -> GeoArrowResult<Self> {
130        let mut counter = Self::new_empty();
131        for maybe_geom in geoms.into_iter() {
132            counter.add_geometry_collection(maybe_geom)?;
133        }
134        Ok(counter)
135    }
136
137    /// Create a capacity counter from an iterator of Geometries.
138    pub fn from_geometries<'a, T: WktNum>(
139        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = T> + 'a)>>,
140    ) -> GeoArrowResult<Self> {
141        let mut counter = Self::new_empty();
142        for maybe_geom in geoms.into_iter() {
143            counter.add_geometry(maybe_geom)?;
144        }
145        Ok(counter)
146    }
147
148    /// The number of bytes an array with this capacity would occupy.
149    pub fn num_bytes(&self, dim: Dimension) -> usize {
150        let offsets_byte_width = 4;
151        let num_offsets = self.geom_capacity;
152        (offsets_byte_width * num_offsets) + self.mixed_capacity.num_bytes(dim)
153    }
154}
155
156impl Default for GeometryCollectionCapacity {
157    fn default() -> Self {
158        Self::new_empty()
159    }
160}
161
162impl Add for GeometryCollectionCapacity {
163    type Output = Self;
164
165    fn add(self, rhs: Self) -> Self::Output {
166        let mixed_capacity = self.mixed_capacity + rhs.mixed_capacity;
167        let geom_capacity = self.geom_capacity + rhs.geom_capacity;
168
169        Self::new(mixed_capacity, geom_capacity)
170    }
171}
geoarrow_array/capacity/geometrycollection.rs

geoarrow_array/capacity/
geometrycollection.rs
