geoarrow_array/capacity/
geometry.rs

1use std::ops::AddAssign;
2
3use geo_traits::*;
4use geoarrow_schema::Dimension;
5use geoarrow_schema::error::GeoArrowResult;
6use wkt::WktNum;
7
8use crate::array::DimensionIndex;
9use crate::builder::geo_trait_wrappers::{LineWrapper, RectWrapper, TriangleWrapper};
10use crate::capacity::{
11    GeometryCollectionCapacity, LineStringCapacity, MultiLineStringCapacity, MultiPointCapacity,
12    MultiPolygonCapacity, PolygonCapacity,
13};
14
15/// A counter for the buffer sizes of a [`GeometryArray`][crate::array::GeometryArray].
16///
17/// This can be used to reduce allocations by allocating once for exactly the array size you need.
18#[derive(Default, Debug, Clone, Copy)]
19pub struct GeometryCapacity {
20    /// The number of null geometries. Ideally the builder will assign these to any array that has
21    /// already been allocated. Otherwise we don't know where to assign them.
22    nulls: usize,
23
24    /// Simple: just the total number of points, nulls included
25    pub(crate) points: [usize; 4],
26    /// An array of [LineStringCapacity], ordered XY, XYZ, XYM, XYZM
27    pub(crate) line_strings: [LineStringCapacity; 4],
28    pub(crate) polygons: [PolygonCapacity; 4],
29    pub(crate) mpoints: [MultiPointCapacity; 4],
30    pub(crate) mline_strings: [MultiLineStringCapacity; 4],
31    pub(crate) mpolygons: [MultiPolygonCapacity; 4],
32    pub(crate) gcs: [GeometryCollectionCapacity; 4],
33
34    /// Whether to prefer multi or single arrays for new geometries.
35    prefer_multi: bool,
36}
37
38impl GeometryCapacity {
39    /// Create a new capacity with known sizes.
40    ///
41    /// Note that the ordering within each array must be XY, XYZ, XYM, XYZM.
42    #[allow(clippy::too_many_arguments)]
43    pub fn new(
44        nulls: usize,
45        points: [usize; 4],
46        line_strings: [LineStringCapacity; 4],
47        polygons: [PolygonCapacity; 4],
48        mpoints: [MultiPointCapacity; 4],
49        mline_strings: [MultiLineStringCapacity; 4],
50        mpolygons: [MultiPolygonCapacity; 4],
51        gcs: [GeometryCollectionCapacity; 4],
52    ) -> Self {
53        Self {
54            nulls,
55            points,
56            line_strings,
57            polygons,
58            mpoints,
59            mline_strings,
60            mpolygons,
61            gcs,
62            prefer_multi: false,
63        }
64    }
65
66    /// Create a new empty capacity.
67    pub fn new_empty() -> Self {
68        Default::default()
69    }
70
71    /// Set whether this capacity counter should prefer allocating "single-type" geometries like
72    /// Point/LineString/Polygon in the arrays of their "Multi" counterparts.
73    pub fn with_prefer_multi(mut self, prefer_multi: bool) -> Self {
74        self.prefer_multi = prefer_multi;
75        self
76    }
77
78    /// Return `true` if the capacity is empty.
79    pub fn is_empty(&self) -> bool {
80        if self.points.iter().any(|c| *c > 0) {
81            return false;
82        }
83
84        if self.line_strings.iter().any(|c| !c.is_empty()) {
85            return false;
86        }
87
88        if self.polygons.iter().any(|c| !c.is_empty()) {
89            return false;
90        }
91
92        if self.mpoints.iter().any(|c| !c.is_empty()) {
93            return false;
94        }
95
96        if self.mline_strings.iter().any(|c| !c.is_empty()) {
97            return false;
98        }
99
100        if self.mpolygons.iter().any(|c| !c.is_empty()) {
101            return false;
102        }
103
104        if self.gcs.iter().any(|c| !c.is_empty()) {
105            return false;
106        }
107
108        true
109    }
110
111    /// The total number of geometries across all geometry types.
112    pub fn total_num_geoms(&self) -> usize {
113        let mut total = 0;
114
115        self.points.iter().for_each(|c| {
116            total += c;
117        });
118        self.line_strings.iter().for_each(|c| {
119            total += c.geom_capacity();
120        });
121        self.polygons.iter().for_each(|c| {
122            total += c.geom_capacity();
123        });
124        self.mpoints.iter().for_each(|c| {
125            total += c.geom_capacity();
126        });
127        self.mline_strings.iter().for_each(|c| {
128            total += c.geom_capacity();
129        });
130        self.mpolygons.iter().for_each(|c| {
131            total += c.geom_capacity();
132        });
133        self.gcs.iter().for_each(|c| {
134            total += c.geom_capacity();
135        });
136
137        total
138    }
139
140    /// Access point capacity
141    pub fn point(&self, dim: Dimension) -> usize {
142        self.points[dim.order()]
143    }
144
145    /// Access LineString capacity
146    pub fn line_string(&self, dim: Dimension) -> LineStringCapacity {
147        self.line_strings[dim.order()]
148    }
149
150    /// Access Polygon capacity
151    pub fn polygon(&self, dim: Dimension) -> PolygonCapacity {
152        self.polygons[dim.order()]
153    }
154
155    /// Access MultiPoint capacity
156    pub fn multi_point(&self, dim: Dimension) -> MultiPointCapacity {
157        self.mpoints[dim.order()]
158    }
159
160    /// Access point capacity
161    pub fn multi_line_string(&self, dim: Dimension) -> MultiLineStringCapacity {
162        self.mline_strings[dim.order()]
163    }
164
165    /// Access point capacity
166    pub fn multi_polygon(&self, dim: Dimension) -> MultiPolygonCapacity {
167        self.mpolygons[dim.order()]
168    }
169
170    /// Access GeometryCollection capacity
171    pub fn geometry_collection(&self, dim: Dimension) -> GeometryCollectionCapacity {
172        self.gcs[dim.order()]
173    }
174
175    /// Add the capacity of the given Point
176    #[inline]
177    fn add_point(&mut self, point: Option<&impl PointTrait>) -> GeoArrowResult<()> {
178        if let Some(point) = point {
179            let dim = Dimension::try_from(point.dim())?;
180            if self.prefer_multi {
181                self.mpoints[dim.order()].add_point(Some(point));
182            } else {
183                self.points[dim.order()] += 1;
184            }
185        } else {
186            self.nulls += 1;
187        }
188        Ok(())
189    }
190
191    /// Add the capacity of the given LineString
192    #[inline]
193    fn add_line_string(
194        &mut self,
195        line_string: Option<&impl LineStringTrait>,
196    ) -> GeoArrowResult<()> {
197        if let Some(line_string) = line_string {
198            let dim = Dimension::try_from(line_string.dim())?;
199            if self.prefer_multi {
200                self.mline_strings[dim.order()].add_line_string(Some(line_string));
201            } else {
202                self.line_strings[dim.order()].add_line_string(Some(line_string));
203            }
204        } else {
205            self.nulls += 1;
206        }
207        Ok(())
208    }
209
210    /// Add the capacity of the given Polygon
211    #[inline]
212    fn add_polygon(&mut self, polygon: Option<&impl PolygonTrait>) -> GeoArrowResult<()> {
213        if let Some(polygon) = polygon {
214            let dim = Dimension::try_from(polygon.dim())?;
215            if self.prefer_multi {
216                self.mpolygons[dim.order()].add_polygon(Some(polygon));
217            } else {
218                self.polygons[dim.order()].add_polygon(Some(polygon));
219            }
220        } else {
221            self.nulls += 1;
222        }
223        Ok(())
224    }
225
226    /// Add the capacity of the given MultiPoint
227    #[inline]
228    fn add_multi_point(
229        &mut self,
230        multi_point: Option<&impl MultiPointTrait>,
231    ) -> GeoArrowResult<()> {
232        if let Some(multi_point) = multi_point {
233            self.multi_point(multi_point.dim().try_into()?)
234                .add_multi_point(Some(multi_point));
235        } else {
236            self.nulls += 1;
237        }
238        Ok(())
239    }
240
241    /// Add the capacity of the given MultiLineString
242    #[inline]
243    fn add_multi_line_string(
244        &mut self,
245        multi_line_string: Option<&impl MultiLineStringTrait>,
246    ) -> GeoArrowResult<()> {
247        if let Some(multi_line_string) = multi_line_string {
248            self.multi_line_string(multi_line_string.dim().try_into()?)
249                .add_multi_line_string(Some(multi_line_string));
250        } else {
251            self.nulls += 1;
252        }
253        Ok(())
254    }
255
256    /// Add the capacity of the given MultiPolygon
257    #[inline]
258    fn add_multi_polygon(
259        &mut self,
260        multi_polygon: Option<&impl MultiPolygonTrait>,
261    ) -> GeoArrowResult<()> {
262        if let Some(multi_polygon) = multi_polygon {
263            self.multi_polygon(multi_polygon.dim().try_into()?)
264                .add_multi_polygon(Some(multi_polygon));
265        } else {
266            self.nulls += 1;
267        }
268        Ok(())
269    }
270
271    /// Add the capacity of the given Geometry
272    #[inline]
273    pub fn add_geometry<T: WktNum>(
274        &mut self,
275        geom: Option<&impl GeometryTrait<T = T>>,
276    ) -> GeoArrowResult<()> {
277        use geo_traits::GeometryType;
278
279        if let Some(geom) = geom {
280            match geom.as_type() {
281                GeometryType::Point(g) => self.add_point(Some(g)),
282                GeometryType::LineString(g) => self.add_line_string(Some(g)),
283                GeometryType::Polygon(g) => self.add_polygon(Some(g)),
284                GeometryType::MultiPoint(p) => self.add_multi_point(Some(p)),
285                GeometryType::MultiLineString(p) => self.add_multi_line_string(Some(p)),
286                GeometryType::MultiPolygon(p) => self.add_multi_polygon(Some(p)),
287                GeometryType::GeometryCollection(p) => self.add_geometry_collection(Some(p)),
288                GeometryType::Rect(r) => self.add_polygon(Some(&RectWrapper::try_new(r)?)),
289                GeometryType::Triangle(tri) => self.add_polygon(Some(&TriangleWrapper(tri))),
290                GeometryType::Line(l) => self.add_line_string(Some(&LineWrapper(l))),
291            }?;
292        } else {
293            self.nulls += 1;
294        }
295        Ok(())
296    }
297
298    /// Add the capacity of the given GeometryCollection
299    #[inline]
300    fn add_geometry_collection<T: WktNum>(
301        &mut self,
302        gc: Option<&impl GeometryCollectionTrait<T = T>>,
303    ) -> GeoArrowResult<()> {
304        if let Some(gc) = gc {
305            self.gcs[Dimension::try_from(gc.dim())?.order()].add_geometry_collection(Some(gc))?;
306        } else {
307            self.nulls += 1;
308        };
309        Ok(())
310    }
311
312    /// Construct a new counter pre-filled with the given geometries
313    pub fn from_geometries<'a, T: WktNum>(
314        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = T> + 'a)>>,
315    ) -> GeoArrowResult<Self> {
316        let mut counter = Self::new_empty();
317        for maybe_geom in geoms.into_iter() {
318            counter.add_geometry(maybe_geom)?;
319        }
320        Ok(counter)
321    }
322
323    /// The number of bytes an array with this capacity would occupy.
324    pub fn num_bytes(&self) -> usize {
325        let mut count = 0;
326
327        self.points.iter().for_each(|c| count += c * 2 * 8);
328        self.line_strings
329            .iter()
330            .for_each(|c| count += c.num_bytes());
331        self.polygons.iter().for_each(|c| count += c.num_bytes());
332        self.mpoints.iter().for_each(|c| count += c.num_bytes());
333        self.mline_strings
334            .iter()
335            .for_each(|c| count += c.num_bytes());
336        self.mpolygons.iter().for_each(|c| count += c.num_bytes());
337        self.gcs.iter().for_each(|c| count += c.num_bytes());
338
339        count
340    }
341}
342
343impl AddAssign for GeometryCapacity {
344    fn add_assign(&mut self, rhs: Self) {
345        self.nulls += rhs.nulls;
346
347        self.points = core::array::from_fn(|i| self.points[i] + rhs.points[i]);
348        self.line_strings = core::array::from_fn(|i| self.line_strings[i] + rhs.line_strings[i]);
349        self.polygons = core::array::from_fn(|i| self.polygons[i] + rhs.polygons[i]);
350        self.mpoints = core::array::from_fn(|i| self.mpoints[i] + rhs.mpoints[i]);
351        self.mline_strings = core::array::from_fn(|i| self.mline_strings[i] + rhs.mline_strings[i]);
352        self.mpolygons = core::array::from_fn(|i| self.mpolygons[i] + rhs.mpolygons[i]);
353        self.gcs = core::array::from_fn(|i| self.gcs[i] + rhs.gcs[i]);
354    }
355}
geoarrow_array/capacity/geometry.rs

geoarrow_array/capacity/
geometry.rs
