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

geoarrow_array/capacity/
geometrycollection.rs
