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    /// Consume the builder and convert to an immutable [`GeometryCollectionArray`]
103    pub fn finish(mut self) -> GeometryCollectionArray {
104        let validity = self.validity.finish();
105        GeometryCollectionArray::new(
106            self.geoms.finish(),
107            self.geom_offsets.finish(),
108            validity,
109            self.data_type.metadata().clone(),
110        )
111    }
112
113    /// Push a Point onto the end of this builder
114    #[inline]
115    fn push_point(&mut self, value: Option<&impl PointTrait<T = f64>>) -> GeoArrowResult<()> {
116        if let Some(geom) = value {
117            self.geoms.push_point(geom)?;
118            self.geom_offsets.try_push_usize(1)?;
119            self.validity.append(value.is_some());
120        } else {
121            self.push_null();
122        }
123        Ok(())
124    }
125
126    /// Push a LineString onto the end of this builder
127    #[inline]
128    fn push_line_string(
129        &mut self,
130        value: Option<&impl LineStringTrait<T = f64>>,
131    ) -> GeoArrowResult<()> {
132        if let Some(geom) = value {
133            self.geoms.push_line_string(geom)?;
134            self.geom_offsets.try_push_usize(1)?;
135            self.validity.append(value.is_some());
136        } else {
137            self.push_null();
138        }
139        Ok(())
140    }
141
142    /// Push a Polygon onto the end of this builder
143    #[inline]
144    fn push_polygon(&mut self, value: Option<&impl PolygonTrait<T = f64>>) -> GeoArrowResult<()> {
145        if let Some(geom) = value {
146            self.geoms.push_polygon(geom)?;
147            self.geom_offsets.try_push_usize(1)?;
148            self.validity.append(value.is_some());
149        } else {
150            self.push_null();
151        }
152        Ok(())
153    }
154
155    /// Push a MultiPoint onto the end of this builder
156    #[inline]
157    fn push_multi_point(
158        &mut self,
159        value: Option<&impl MultiPointTrait<T = f64>>,
160    ) -> GeoArrowResult<()> {
161        if let Some(geom) = value {
162            self.geoms.push_multi_point(geom)?;
163            self.geom_offsets.try_push_usize(1)?;
164            self.validity.append(value.is_some());
165        } else {
166            self.push_null();
167        }
168        Ok(())
169    }
170
171    /// Push a MultiLineString onto the end of this builder
172    #[inline]
173    fn push_multi_line_string(
174        &mut self,
175        value: Option<&impl MultiLineStringTrait<T = f64>>,
176    ) -> GeoArrowResult<()> {
177        if let Some(geom) = value {
178            self.geoms.push_multi_line_string(geom)?;
179            self.geom_offsets.try_push_usize(1)?;
180            self.validity.append(value.is_some());
181        } else {
182            self.push_null();
183        }
184        Ok(())
185    }
186
187    /// Push a MultiPolygon onto the end of this builder
188    #[inline]
189    fn push_multi_polygon(
190        &mut self,
191        value: Option<&impl MultiPolygonTrait<T = f64>>,
192    ) -> GeoArrowResult<()> {
193        if let Some(geom) = value {
194            self.geoms.push_multi_polygon(geom)?;
195            self.geom_offsets.try_push_usize(1)?;
196            self.validity.append(value.is_some());
197        } else {
198            self.push_null();
199        }
200        Ok(())
201    }
202
203    /// Push a Geometry onto the end of this builder
204    #[inline]
205    pub fn push_geometry(
206        &mut self,
207        value: Option<&impl GeometryTrait<T = f64>>,
208    ) -> GeoArrowResult<()> {
209        use geo_traits::GeometryType::*;
210
211        if let Some(g) = value {
212            match g.as_type() {
213                Point(p) => self.push_point(Some(p))?,
214                LineString(p) => {
215                    self.push_line_string(Some(p))?;
216                }
217                Polygon(p) => self.push_polygon(Some(p))?,
218                MultiPoint(p) => self.push_multi_point(Some(p))?,
219                MultiLineString(p) => self.push_multi_line_string(Some(p))?,
220                MultiPolygon(p) => self.push_multi_polygon(Some(p))?,
221                GeometryCollection(p) => self.push_geometry_collection(Some(p))?,
222                Rect(r) => self.push_polygon(Some(&RectWrapper::try_new(r)?))?,
223                Triangle(tri) => self.push_polygon(Some(&TriangleWrapper(tri)))?,
224                Line(l) => self.push_line_string(Some(&LineWrapper(l)))?,
225            }
226        } else {
227            self.push_null();
228        };
229        Ok(())
230    }
231
232    /// Push a GeometryCollection onto the end of this builder
233    #[inline]
234    pub fn push_geometry_collection(
235        &mut self,
236        value: Option<&impl GeometryCollectionTrait<T = f64>>,
237    ) -> GeoArrowResult<()> {
238        if let Some(gc) = value {
239            let num_geoms = gc.num_geometries();
240            for g in gc.geometries() {
241                self.geoms.push_geometry(&g)?;
242            }
243            self.try_push_length(num_geoms)?;
244        } else {
245            self.push_null();
246        }
247        Ok(())
248    }
249
250    /// Extend this builder with the given geometries
251    pub fn extend_from_iter(
252        &mut self,
253        geoms: impl Iterator<Item = Option<&'a (impl GeometryCollectionTrait<T = f64> + 'a)>>,
254    ) {
255        geoms
256            .into_iter()
257            .try_for_each(|maybe_gc| self.push_geometry_collection(maybe_gc))
258            .unwrap();
259    }
260
261    #[inline]
262    pub(crate) fn try_push_length(&mut self, geom_offsets_length: usize) -> GeoArrowResult<()> {
263        self.geom_offsets.try_push_usize(geom_offsets_length)?;
264        self.validity.append(true);
265        Ok(())
266    }
267
268    #[inline]
269    pub(crate) fn push_null(&mut self) {
270        self.geom_offsets.extend_constant(1);
271        self.validity.append(false);
272    }
273
274    /// Construct a new builder, pre-filling it with the provided geometries
275    pub fn from_geometry_collections(
276        geoms: &[impl GeometryCollectionTrait<T = f64>],
277        typ: GeometryCollectionType,
278    ) -> GeoArrowResult<Self> {
279        let capacity =
280            GeometryCollectionCapacity::from_geometry_collections(geoms.iter().map(Some))?;
281        let mut array = Self::with_capacity(typ, capacity);
282        array.extend_from_iter(geoms.iter().map(Some));
283        Ok(array)
284    }
285
286    /// Construct a new builder, pre-filling it with the provided geometries
287    pub fn from_nullable_geometry_collections(
288        geoms: &[Option<impl GeometryCollectionTrait<T = f64>>],
289        typ: GeometryCollectionType,
290    ) -> GeoArrowResult<Self> {
291        let capacity = GeometryCollectionCapacity::from_geometry_collections(
292            geoms.iter().map(|x| x.as_ref()),
293        )?;
294        let mut array = Self::with_capacity(typ, capacity);
295        array.extend_from_iter(geoms.iter().map(|x| x.as_ref()));
296        Ok(array)
297    }
298
299    /// Construct a new builder, pre-filling it with the provided geometries
300    pub fn from_nullable_geometries(
301        geoms: &[Option<impl GeometryTrait<T = f64>>],
302        typ: GeometryCollectionType,
303    ) -> GeoArrowResult<Self> {
304        let capacity =
305            GeometryCollectionCapacity::from_geometries(geoms.iter().map(|x| x.as_ref()))?;
306        let mut array = Self::with_capacity(typ, capacity);
307        for geom in geoms {
308            array.push_geometry(geom.as_ref())?;
309        }
310        Ok(array)
311    }
312}
313
314impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, GeometryCollectionType)>
315    for GeometryCollectionBuilder
316{
317    type Error = GeoArrowError;
318
319    fn try_from(
320        (value, typ): (GenericWkbArray<O>, GeometryCollectionType),
321    ) -> GeoArrowResult<Self> {
322        let wkb_objects = value
323            .iter()
324            .map(|x| x.transpose())
325            .collect::<GeoArrowResult<Vec<_>>>()?;
326        Self::from_nullable_geometries(&wkb_objects, typ)
327    }
328}
329
330impl GeoArrowArrayBuilder for GeometryCollectionBuilder {
331    fn len(&self) -> usize {
332        self.geom_offsets.len_proxy()
333    }
334
335    fn push_null(&mut self) {
336        self.push_null();
337    }
338
339    fn push_geometry(
340        &mut self,
341        geometry: Option<&impl GeometryTrait<T = f64>>,
342    ) -> GeoArrowResult<()> {
343        self.push_geometry(geometry)
344    }
345
346    fn finish(self) -> Arc<dyn GeoArrowArray> {
347        Arc::new(self.finish())
348    }
349}
geoarrow_array/builder/geometrycollection.rs

geoarrow_array/builder/
geometrycollection.rs
