geoarrow_array/builder/
geometrycollection.rs

1use std::sync::Arc;
2
3use arrow_array::OffsetSizeTrait;
4use arrow_buffer::NullBufferBuilder;
5use geo_traits::{
6    GeometryCollectionTrait, GeometryTrait, LineStringTrait, MultiLineStringTrait, MultiPointTrait,
7    MultiPolygonTrait, PointTrait, PolygonTrait,
8};
9use geoarrow_schema::GeometryCollectionType;
10use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
11
12use crate::GeoArrowArray;
13use crate::array::{GenericWkbArray, GeometryCollectionArray};
14use crate::builder::geo_trait_wrappers::{LineWrapper, RectWrapper, TriangleWrapper};
15use crate::builder::{MixedGeometryBuilder, OffsetsBuilder};
16use crate::capacity::GeometryCollectionCapacity;
17use crate::trait_::{GeoArrowArrayAccessor, GeoArrowArrayBuilder};
18
19/// The GeoArrow equivalent to `Vec<Option<GeometryCollection>>`: a mutable collection of
20/// GeometryCollections.
21///
22/// Converting an [`GeometryCollectionBuilder`] into a [`GeometryCollectionArray`] is `O(1)`.
23#[derive(Debug)]
24pub struct GeometryCollectionBuilder {
25    data_type: GeometryCollectionType,
26
27    pub(crate) geoms: MixedGeometryBuilder,
28
29    pub(crate) geom_offsets: OffsetsBuilder<i32>,
30
31    pub(crate) validity: NullBufferBuilder,
32}
33
34impl<'a> GeometryCollectionBuilder {
35    /// Creates a new empty [`GeometryCollectionBuilder`].
36    pub fn new(typ: GeometryCollectionType) -> Self {
37        Self::with_capacity(typ, Default::default())
38    }
39
40    /// Creates a new empty [`GeometryCollectionBuilder`] with the provided
41    /// [capacity][GeometryCollectionCapacity].
42    pub fn with_capacity(
43        typ: GeometryCollectionType,
44        capacity: GeometryCollectionCapacity,
45    ) -> Self {
46        Self {
47            geoms: MixedGeometryBuilder::with_capacity_and_options(
48                typ.dimension(),
49                capacity.mixed_capacity,
50                typ.coord_type(),
51            ),
52            geom_offsets: OffsetsBuilder::with_capacity(capacity.geom_capacity),
53            validity: NullBufferBuilder::new(capacity.geom_capacity),
54            data_type: typ,
55        }
56    }
57
58    /// Change whether to prefer multi or single arrays for new single-part geometries.
59    ///
60    /// If `true`, a new `Point` will be added to the `MultiPointBuilder` child array, a new
61    /// `LineString` will be added to the `MultiLineStringBuilder` child array, and a new `Polygon`
62    /// will be added to the `MultiPolygonBuilder` child array.
63    ///
64    /// This can be desired when the user wants to downcast the array to a single geometry array
65    /// later, as casting to a, say, `MultiPointArray` from a `GeometryCollectionArray` could be
66    /// done zero-copy.
67    ///
68    /// Note that only geometries added _after_ this method is called will be affected.
69    pub fn with_prefer_multi(self, prefer_multi: bool) -> Self {
70        Self {
71            geoms: self.geoms.with_prefer_multi(prefer_multi),
72            ..self
73        }
74    }
75
76    /// Reserves capacity for at least `additional` more GeometryCollections.
77    ///
78    /// The collection may reserve more space to speculatively avoid frequent reallocations. After
79    /// calling `reserve`, capacity will be greater than or equal to `self.len() + additional`.
80    /// Does nothing if capacity is already sufficient.
81    pub fn reserve(&mut self, additional: GeometryCollectionCapacity) {
82        self.geoms.reserve(additional.mixed_capacity);
83        self.geom_offsets.reserve(additional.geom_capacity);
84    }
85
86    /// Reserves the minimum capacity for at least `additional` more GeometryCollections.
87    ///
88    /// Unlike [`reserve`], this will not deliberately over-allocate to speculatively avoid
89    /// frequent allocations. After calling `reserve_exact`, capacity will be greater than or equal
90    /// to `self.len() + additional`. Does nothing if the capacity is already sufficient.
91    ///
92    /// Note that the allocator may give the collection more space than it
93    /// requests. Therefore, capacity can not be relied upon to be precisely
94    /// minimal. Prefer [`reserve`] if future insertions are expected.
95    ///
96    /// [`reserve`]: Self::reserve
97    pub fn reserve_exact(&mut self, additional: GeometryCollectionCapacity) {
98        self.geoms.reserve_exact(additional.mixed_capacity);
99        self.geom_offsets.reserve_exact(additional.geom_capacity);
100    }
101
102    /// Shrinks the capacity of self to fit.
103    pub fn shrink_to_fit(&mut self) {
104        self.geoms.shrink_to_fit();
105        self.geom_offsets.shrink_to_fit();
106        // self.validity.shrink_to_fit();
107    }
108
109    /// Consume the builder and convert to an immutable [`GeometryCollectionArray`]
110    pub fn finish(mut self) -> GeometryCollectionArray {
111        let validity = self.validity.finish();
112        GeometryCollectionArray::new(
113            self.geoms.finish(),
114            self.geom_offsets.finish(),
115            validity,
116            self.data_type.metadata().clone(),
117        )
118    }
119
120    /// Push a Point onto the end of this builder
121    #[inline]
122    fn push_point(&mut self, value: Option<&impl PointTrait<T = f64>>) -> GeoArrowResult<()> {
123        if let Some(geom) = value {
124            self.geoms.push_point(geom)?;
125            self.geom_offsets.try_push_usize(1)?;
126            self.validity.append(value.is_some());
127        } else {
128            self.push_null();
129        }
130        Ok(())
131    }
132
133    /// Push a LineString onto the end of this builder
134    #[inline]
135    fn push_line_string(
136        &mut self,
137        value: Option<&impl LineStringTrait<T = f64>>,
138    ) -> GeoArrowResult<()> {
139        if let Some(geom) = value {
140            self.geoms.push_line_string(geom)?;
141            self.geom_offsets.try_push_usize(1)?;
142            self.validity.append(value.is_some());
143        } else {
144            self.push_null();
145        }
146        Ok(())
147    }
148
149    /// Push a Polygon onto the end of this builder
150    #[inline]
151    fn push_polygon(&mut self, value: Option<&impl PolygonTrait<T = f64>>) -> GeoArrowResult<()> {
152        if let Some(geom) = value {
153            self.geoms.push_polygon(geom)?;
154            self.geom_offsets.try_push_usize(1)?;
155            self.validity.append(value.is_some());
156        } else {
157            self.push_null();
158        }
159        Ok(())
160    }
161
162    /// Push a MultiPoint onto the end of this builder
163    #[inline]
164    fn push_multi_point(
165        &mut self,
166        value: Option<&impl MultiPointTrait<T = f64>>,
167    ) -> GeoArrowResult<()> {
168        if let Some(geom) = value {
169            self.geoms.push_multi_point(geom)?;
170            self.geom_offsets.try_push_usize(1)?;
171            self.validity.append(value.is_some());
172        } else {
173            self.push_null();
174        }
175        Ok(())
176    }
177
178    /// Push a MultiLineString onto the end of this builder
179    #[inline]
180    fn push_multi_line_string(
181        &mut self,
182        value: Option<&impl MultiLineStringTrait<T = f64>>,
183    ) -> GeoArrowResult<()> {
184        if let Some(geom) = value {
185            self.geoms.push_multi_line_string(geom)?;
186            self.geom_offsets.try_push_usize(1)?;
187            self.validity.append(value.is_some());
188        } else {
189            self.push_null();
190        }
191        Ok(())
192    }
193
194    /// Push a MultiPolygon onto the end of this builder
195    #[inline]
196    fn push_multi_polygon(
197        &mut self,
198        value: Option<&impl MultiPolygonTrait<T = f64>>,
199    ) -> GeoArrowResult<()> {
200        if let Some(geom) = value {
201            self.geoms.push_multi_polygon(geom)?;
202            self.geom_offsets.try_push_usize(1)?;
203            self.validity.append(value.is_some());
204        } else {
205            self.push_null();
206        }
207        Ok(())
208    }
209
210    /// Push a Geometry onto the end of this builder
211    #[inline]
212    pub fn push_geometry(
213        &mut self,
214        value: Option<&impl GeometryTrait<T = f64>>,
215    ) -> GeoArrowResult<()> {
216        use geo_traits::GeometryType::*;
217
218        if let Some(g) = value {
219            match g.as_type() {
220                Point(p) => self.push_point(Some(p))?,
221                LineString(p) => {
222                    self.push_line_string(Some(p))?;
223                }
224                Polygon(p) => self.push_polygon(Some(p))?,
225                MultiPoint(p) => self.push_multi_point(Some(p))?,
226                MultiLineString(p) => self.push_multi_line_string(Some(p))?,
227                MultiPolygon(p) => self.push_multi_polygon(Some(p))?,
228                GeometryCollection(p) => self.push_geometry_collection(Some(p))?,
229                Rect(r) => self.push_polygon(Some(&RectWrapper::try_new(r)?))?,
230                Triangle(tri) => self.push_polygon(Some(&TriangleWrapper(tri)))?,
231                Line(l) => self.push_line_string(Some(&LineWrapper(l)))?,
232            }
233        } else {
234            self.push_null();
235        };
236        Ok(())
237    }
238
239    /// Push a GeometryCollection onto the end of this builder
240    #[inline]
241    pub fn push_geometry_collection(
242        &mut self,
243        value: Option<&impl GeometryCollectionTrait<T = f64>>,
244    ) -> GeoArrowResult<()> {
245        if let Some(gc) = value {
246            let num_geoms = gc.num_geometries();
247            for g in gc.geometries() {
248                self.geoms.push_geometry(&g)?;
249            }
250            self.try_push_length(num_geoms)?;
251        } else {
252            self.push_null();
253        }
254        Ok(())
255    }
256
257    /// Extend this builder with the given geometries
258    pub fn extend_from_iter(
259        &mut self,
260        geoms: impl Iterator<Item = Option<&'a (impl GeometryCollectionTrait<T = f64> + 'a)>>,
261    ) {
262        geoms
263            .into_iter()
264            .try_for_each(|maybe_gc| self.push_geometry_collection(maybe_gc))
265            .unwrap();
266    }
267
268    #[inline]
269    pub(crate) fn try_push_length(&mut self, geom_offsets_length: usize) -> GeoArrowResult<()> {
270        self.geom_offsets.try_push_usize(geom_offsets_length)?;
271        self.validity.append(true);
272        Ok(())
273    }
274
275    #[inline]
276    pub(crate) fn push_null(&mut self) {
277        self.geom_offsets.extend_constant(1);
278        self.validity.append(false);
279    }
280
281    /// Construct a new builder, pre-filling it with the provided geometries
282    pub fn from_geometry_collections(
283        geoms: &[impl GeometryCollectionTrait<T = f64>],
284        typ: GeometryCollectionType,
285    ) -> GeoArrowResult<Self> {
286        let capacity =
287            GeometryCollectionCapacity::from_geometry_collections(geoms.iter().map(Some))?;
288        let mut array = Self::with_capacity(typ, capacity);
289        array.extend_from_iter(geoms.iter().map(Some));
290        Ok(array)
291    }
292
293    /// Construct a new builder, pre-filling it with the provided geometries
294    pub fn from_nullable_geometry_collections(
295        geoms: &[Option<impl GeometryCollectionTrait<T = f64>>],
296        typ: GeometryCollectionType,
297    ) -> GeoArrowResult<Self> {
298        let capacity = GeometryCollectionCapacity::from_geometry_collections(
299            geoms.iter().map(|x| x.as_ref()),
300        )?;
301        let mut array = Self::with_capacity(typ, capacity);
302        array.extend_from_iter(geoms.iter().map(|x| x.as_ref()));
303        Ok(array)
304    }
305
306    /// Construct a new builder, pre-filling it with the provided geometries
307    pub fn from_nullable_geometries(
308        geoms: &[Option<impl GeometryTrait<T = f64>>],
309        typ: GeometryCollectionType,
310    ) -> GeoArrowResult<Self> {
311        let capacity =
312            GeometryCollectionCapacity::from_geometries(geoms.iter().map(|x| x.as_ref()))?;
313        let mut array = Self::with_capacity(typ, capacity);
314        for geom in geoms {
315            array.push_geometry(geom.as_ref())?;
316        }
317        Ok(array)
318    }
319}
320
321impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, GeometryCollectionType)>
322    for GeometryCollectionBuilder
323{
324    type Error = GeoArrowError;
325
326    fn try_from(
327        (value, typ): (GenericWkbArray<O>, GeometryCollectionType),
328    ) -> GeoArrowResult<Self> {
329        let wkb_objects = value
330            .iter()
331            .map(|x| x.transpose())
332            .collect::<GeoArrowResult<Vec<_>>>()?;
333        Self::from_nullable_geometries(&wkb_objects, typ)
334    }
335}
336
337impl GeoArrowArrayBuilder for GeometryCollectionBuilder {
338    fn len(&self) -> usize {
339        self.geom_offsets.len_proxy()
340    }
341
342    fn push_null(&mut self) {
343        self.push_null();
344    }
345
346    fn push_geometry(
347        &mut self,
348        geometry: Option<&impl GeometryTrait<T = f64>>,
349    ) -> GeoArrowResult<()> {
350        self.push_geometry(geometry)
351    }
352
353    fn finish(self) -> Arc<dyn GeoArrowArray> {
354        Arc::new(self.finish())
355    }
356}
geoarrow_array/builder/geometrycollection.rs

geoarrow_array/builder/
geometrycollection.rs
