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

geoarrow_array/builder/
geometrycollection.rs
