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