geoarrow_array/builder/
multipoint.rs

1use std::sync::Arc;
2
3use arrow_array::OffsetSizeTrait;
4use arrow_buffer::NullBufferBuilder;
5use geo_traits::{CoordTrait, GeometryTrait, GeometryType, MultiPointTrait, PointTrait};
6use geoarrow_schema::MultiPointType;
7use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8
9use crate::GeoArrowArray;
10use crate::array::{GenericWkbArray, MultiPointArray};
11use crate::builder::{CoordBufferBuilder, OffsetsBuilder};
12use crate::capacity::MultiPointCapacity;
13use crate::trait_::{GeoArrowArrayAccessor, GeoArrowArrayBuilder};
14use crate::util::GeometryTypeName;
15
16/// The GeoArrow equivalent to `Vec<Option<MultiPoint>>`: a mutable collection of MultiPoints.
17///
18/// Converting an [`MultiPointBuilder`] into a [`MultiPointArray`] is `O(1)`.
19#[derive(Debug)]
20pub struct MultiPointBuilder {
21    data_type: MultiPointType,
22
23    coords: CoordBufferBuilder,
24
25    geom_offsets: OffsetsBuilder<i32>,
26
27    /// Validity is only defined at the geometry level
28    validity: NullBufferBuilder,
29}
30
31impl MultiPointBuilder {
32    /// Creates a new empty [`MultiPointBuilder`].
33    pub fn new(typ: MultiPointType) -> Self {
34        Self::with_capacity(typ, Default::default())
35    }
36
37    /// Creates a new [`MultiPointBuilder`] with a capacity.
38    pub fn with_capacity(typ: MultiPointType, capacity: MultiPointCapacity) -> Self {
39        let coords = CoordBufferBuilder::with_capacity(
40            capacity.coord_capacity,
41            typ.coord_type(),
42            typ.dimension(),
43        );
44        Self {
45            coords,
46            geom_offsets: OffsetsBuilder::with_capacity(capacity.geom_capacity),
47            validity: NullBufferBuilder::new(capacity.geom_capacity),
48            data_type: typ,
49        }
50    }
51
52    /// Reserves capacity for at least `additional` more MultiPoints.
53    ///
54    /// The collection may reserve more space to speculatively avoid frequent reallocations. After
55    /// calling `reserve`, capacity will be greater than or equal to `self.len() + additional`.
56    /// Does nothing if capacity is already sufficient.
57    pub fn reserve(&mut self, capacity: MultiPointCapacity) {
58        self.coords.reserve(capacity.coord_capacity);
59        self.geom_offsets.reserve(capacity.geom_capacity);
60    }
61
62    /// Reserves the minimum capacity for at least `additional` more MultiPoints.
63    ///
64    /// Unlike [`reserve`], this will not deliberately over-allocate to speculatively avoid
65    /// frequent allocations. After calling `reserve_exact`, capacity will be greater than or equal
66    /// to `self.len() + additional`. Does nothing if the capacity is already sufficient.
67    ///
68    /// Note that the allocator may give the collection more space than it
69    /// requests. Therefore, capacity can not be relied upon to be precisely
70    /// minimal. Prefer [`reserve`] if future insertions are expected.
71    ///
72    /// [`reserve`]: Self::reserve
73    pub fn reserve_exact(&mut self, capacity: MultiPointCapacity) {
74        self.coords.reserve_exact(capacity.coord_capacity);
75        self.geom_offsets.reserve_exact(capacity.geom_capacity);
76    }
77
78    /// Shrinks the capacity of self to fit.
79    pub fn shrink_to_fit(&mut self) {
80        self.coords.shrink_to_fit();
81        self.geom_offsets.shrink_to_fit();
82        // self.validity.shrink_to_fit();
83    }
84
85    /// Consume the builder and convert to an immutable [`MultiPointArray`]
86    pub fn finish(mut self) -> MultiPointArray {
87        let validity = self.validity.finish();
88        MultiPointArray::new(
89            self.coords.finish(),
90            self.geom_offsets.finish(),
91            validity,
92            self.data_type.metadata().clone(),
93        )
94    }
95
96    /// Extend this builder with the given geometries
97    pub fn extend_from_iter<'a>(
98        &mut self,
99        geoms: impl Iterator<Item = Option<&'a (impl MultiPointTrait<T = f64> + 'a)>>,
100    ) {
101        geoms
102            .into_iter()
103            .try_for_each(|maybe_multi_point| self.push_multi_point(maybe_multi_point))
104            .unwrap();
105    }
106
107    /// Extend this builder with the given geometries
108    pub fn extend_from_geometry_iter<'a>(
109        &mut self,
110        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = f64> + 'a)>>,
111    ) -> GeoArrowResult<()> {
112        geoms.into_iter().try_for_each(|g| self.push_geometry(g))?;
113        Ok(())
114    }
115
116    /// Add a new Point to the end of this array.
117    ///
118    /// # Errors
119    ///
120    /// This function errors iff the new last item is larger than what O supports.
121    #[inline]
122    pub fn push_point(&mut self, value: Option<&impl PointTrait<T = f64>>) -> GeoArrowResult<()> {
123        if let Some(point) = value {
124            self.coords.push_point(point);
125            self.try_push_length(1)?;
126        } else {
127            self.push_null();
128        }
129
130        Ok(())
131    }
132
133    /// Add a new MultiPoint to the end of this array.
134    ///
135    /// # Errors
136    ///
137    /// This function errors iff the new last item is larger than what O supports.
138    #[inline]
139    pub fn push_multi_point(
140        &mut self,
141        value: Option<&impl MultiPointTrait<T = f64>>,
142    ) -> GeoArrowResult<()> {
143        if let Some(multi_point) = value {
144            let num_points = multi_point.num_points();
145            for point in multi_point.points() {
146                self.coords.push_point(&point);
147            }
148            self.try_push_length(num_points)?;
149        } else {
150            self.push_null();
151        }
152        Ok(())
153    }
154
155    /// Add a new geometry to this builder
156    ///
157    /// This will error if the geometry type is not Point or MultiPoint.
158    #[inline]
159    pub fn push_geometry(
160        &mut self,
161        value: Option<&impl GeometryTrait<T = f64>>,
162    ) -> GeoArrowResult<()> {
163        if let Some(value) = value {
164            match value.as_type() {
165                GeometryType::Point(g) => self.push_point(Some(g))?,
166                GeometryType::MultiPoint(g) => self.push_multi_point(Some(g))?,
167                gt => {
168                    return Err(GeoArrowError::IncorrectGeometryType(format!(
169                        "Expected MultiPoint compatible geometry, got {}",
170                        gt.name()
171                    )));
172                }
173            }
174        } else {
175            self.push_null();
176        };
177        Ok(())
178    }
179
180    /// Push a raw coordinate to the underlying coordinate array.
181    ///
182    /// # Invariant
183    ///
184    /// Care must be taken to ensure that pushing raw coordinates to the array upholds the
185    /// necessary invariants of the array.
186    #[inline]
187    pub(crate) fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> GeoArrowResult<()> {
188        self.coords.try_push_coord(coord)
189    }
190
191    /// Needs to be called when a valid value was extended to this array.
192    /// This is a relatively low level function, prefer `try_push` when you can.
193    #[inline]
194    pub(crate) fn try_push_length(&mut self, geom_offsets_length: usize) -> GeoArrowResult<()> {
195        self.geom_offsets.try_push_usize(geom_offsets_length)?;
196        self.validity.append(true);
197        Ok(())
198    }
199
200    #[inline]
201    pub(crate) fn push_null(&mut self) {
202        self.geom_offsets.extend_constant(1);
203        self.validity.append(false);
204    }
205
206    /// Construct a new builder, pre-filling it with the provided geometries
207    pub fn from_multi_points(geoms: &[impl MultiPointTrait<T = f64>], typ: MultiPointType) -> Self {
208        let capacity = MultiPointCapacity::from_multi_points(geoms.iter().map(Some));
209        let mut array = Self::with_capacity(typ, capacity);
210        array.extend_from_iter(geoms.iter().map(Some));
211        array
212    }
213
214    /// Construct a new builder, pre-filling it with the provided geometries
215    pub fn from_nullable_multi_points(
216        geoms: &[Option<impl MultiPointTrait<T = f64>>],
217        typ: MultiPointType,
218    ) -> Self {
219        let capacity = MultiPointCapacity::from_multi_points(geoms.iter().map(|x| x.as_ref()));
220        let mut array = Self::with_capacity(typ, capacity);
221        array.extend_from_iter(geoms.iter().map(|x| x.as_ref()));
222        array
223    }
224
225    /// Construct a new builder, pre-filling it with the provided geometries
226    pub fn from_nullable_geometries(
227        geoms: &[Option<impl GeometryTrait<T = f64>>],
228        typ: MultiPointType,
229    ) -> GeoArrowResult<Self> {
230        let capacity = MultiPointCapacity::from_geometries(geoms.iter().map(|x| x.as_ref()))?;
231        let mut array = Self::with_capacity(typ, capacity);
232        array.extend_from_geometry_iter(geoms.iter().map(|x| x.as_ref()))?;
233        Ok(array)
234    }
235}
236
237impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, MultiPointType)> for MultiPointBuilder {
238    type Error = GeoArrowError;
239
240    fn try_from((value, typ): (GenericWkbArray<O>, MultiPointType)) -> GeoArrowResult<Self> {
241        let wkb_objects = value
242            .iter()
243            .map(|x| x.transpose())
244            .collect::<GeoArrowResult<Vec<_>>>()?;
245        Self::from_nullable_geometries(&wkb_objects, typ)
246    }
247}
248
249impl GeoArrowArrayBuilder for MultiPointBuilder {
250    fn len(&self) -> usize {
251        self.geom_offsets.len_proxy()
252    }
253
254    fn push_null(&mut self) {
255        self.push_null();
256    }
257
258    fn push_geometry(
259        &mut self,
260        geometry: Option<&impl GeometryTrait<T = f64>>,
261    ) -> GeoArrowResult<()> {
262        self.push_geometry(geometry)
263    }
264
265    fn finish(self) -> Arc<dyn GeoArrowArray> {
266        Arc::new(self.finish())
267    }
268}