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    /// Consume the builder and convert to an immutable [`MultiPointArray`]
79    pub fn finish(mut self) -> MultiPointArray {
80        let validity = self.validity.finish();
81
82        // TODO: impl shrink_to_fit for all mutable -> * impls
83        // self.coords.shrink_to_fit();
84        self.geom_offsets.shrink_to_fit();
85
86        MultiPointArray::new(
87            self.coords.finish(),
88            self.geom_offsets.finish(),
89            validity,
90            self.data_type.metadata().clone(),
91        )
92    }
93
94    /// Extend this builder with the given geometries
95    pub fn extend_from_iter<'a>(
96        &mut self,
97        geoms: impl Iterator<Item = Option<&'a (impl MultiPointTrait<T = f64> + 'a)>>,
98    ) {
99        geoms
100            .into_iter()
101            .try_for_each(|maybe_multi_point| self.push_multi_point(maybe_multi_point))
102            .unwrap();
103    }
104
105    /// Extend this builder with the given geometries
106    pub fn extend_from_geometry_iter<'a>(
107        &mut self,
108        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = f64> + 'a)>>,
109    ) -> GeoArrowResult<()> {
110        geoms.into_iter().try_for_each(|g| self.push_geometry(g))?;
111        Ok(())
112    }
113
114    /// Add a new Point to the end of this array.
115    ///
116    /// # Errors
117    ///
118    /// This function errors iff the new last item is larger than what O supports.
119    #[inline]
120    pub fn push_point(&mut self, value: Option<&impl PointTrait<T = f64>>) -> GeoArrowResult<()> {
121        if let Some(point) = value {
122            self.coords.push_point(point);
123            self.try_push_length(1)?;
124        } else {
125            self.push_null();
126        }
127
128        Ok(())
129    }
130
131    /// Add a new MultiPoint to the end of this array.
132    ///
133    /// # Errors
134    ///
135    /// This function errors iff the new last item is larger than what O supports.
136    #[inline]
137    pub fn push_multi_point(
138        &mut self,
139        value: Option<&impl MultiPointTrait<T = f64>>,
140    ) -> GeoArrowResult<()> {
141        if let Some(multi_point) = value {
142            let num_points = multi_point.num_points();
143            for point in multi_point.points() {
144                self.coords.push_point(&point);
145            }
146            self.try_push_length(num_points)?;
147        } else {
148            self.push_null();
149        }
150        Ok(())
151    }
152
153    /// Add a new geometry to this builder
154    ///
155    /// This will error if the geometry type is not Point or MultiPoint.
156    #[inline]
157    pub fn push_geometry(
158        &mut self,
159        value: Option<&impl GeometryTrait<T = f64>>,
160    ) -> GeoArrowResult<()> {
161        if let Some(value) = value {
162            match value.as_type() {
163                GeometryType::Point(g) => self.push_point(Some(g))?,
164                GeometryType::MultiPoint(g) => self.push_multi_point(Some(g))?,
165                gt => {
166                    return Err(GeoArrowError::IncorrectGeometryType(format!(
167                        "Expected MultiPoint compatible geometry, got {}",
168                        gt.name()
169                    )));
170                }
171            }
172        } else {
173            self.push_null();
174        };
175        Ok(())
176    }
177
178    /// Push a raw coordinate to the underlying coordinate array.
179    ///
180    /// # Invariant
181    ///
182    /// Care must be taken to ensure that pushing raw coordinates to the array upholds the
183    /// necessary invariants of the array.
184    #[inline]
185    pub(crate) fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> GeoArrowResult<()> {
186        self.coords.try_push_coord(coord)
187    }
188
189    /// Needs to be called when a valid value was extended to this array.
190    /// This is a relatively low level function, prefer `try_push` when you can.
191    #[inline]
192    pub(crate) fn try_push_length(&mut self, geom_offsets_length: usize) -> GeoArrowResult<()> {
193        self.geom_offsets.try_push_usize(geom_offsets_length)?;
194        self.validity.append(true);
195        Ok(())
196    }
197
198    #[inline]
199    pub(crate) fn push_null(&mut self) {
200        self.geom_offsets.extend_constant(1);
201        self.validity.append(false);
202    }
203
204    /// Construct a new builder, pre-filling it with the provided geometries
205    pub fn from_multi_points(geoms: &[impl MultiPointTrait<T = f64>], typ: MultiPointType) -> Self {
206        let capacity = MultiPointCapacity::from_multi_points(geoms.iter().map(Some));
207        let mut array = Self::with_capacity(typ, capacity);
208        array.extend_from_iter(geoms.iter().map(Some));
209        array
210    }
211
212    /// Construct a new builder, pre-filling it with the provided geometries
213    pub fn from_nullable_multi_points(
214        geoms: &[Option<impl MultiPointTrait<T = f64>>],
215        typ: MultiPointType,
216    ) -> Self {
217        let capacity = MultiPointCapacity::from_multi_points(geoms.iter().map(|x| x.as_ref()));
218        let mut array = Self::with_capacity(typ, capacity);
219        array.extend_from_iter(geoms.iter().map(|x| x.as_ref()));
220        array
221    }
222
223    /// Construct a new builder, pre-filling it with the provided geometries
224    pub fn from_nullable_geometries(
225        geoms: &[Option<impl GeometryTrait<T = f64>>],
226        typ: MultiPointType,
227    ) -> GeoArrowResult<Self> {
228        let capacity = MultiPointCapacity::from_geometries(geoms.iter().map(|x| x.as_ref()))?;
229        let mut array = Self::with_capacity(typ, capacity);
230        array.extend_from_geometry_iter(geoms.iter().map(|x| x.as_ref()))?;
231        Ok(array)
232    }
233}
234
235impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, MultiPointType)> for MultiPointBuilder {
236    type Error = GeoArrowError;
237
238    fn try_from((value, typ): (GenericWkbArray<O>, MultiPointType)) -> GeoArrowResult<Self> {
239        let wkb_objects = value
240            .iter()
241            .map(|x| x.transpose())
242            .collect::<GeoArrowResult<Vec<_>>>()?;
243        Self::from_nullable_geometries(&wkb_objects, typ)
244    }
245}
246
247impl GeoArrowArrayBuilder for MultiPointBuilder {
248    fn len(&self) -> usize {
249        self.geom_offsets.len_proxy()
250    }
251
252    fn push_null(&mut self) {
253        self.push_null();
254    }
255
256    fn push_geometry(
257        &mut self,
258        geometry: Option<&impl GeometryTrait<T = f64>>,
259    ) -> GeoArrowResult<()> {
260        self.push_geometry(geometry)
261    }
262
263    fn finish(self) -> Arc<dyn GeoArrowArray> {
264        Arc::new(self.finish())
265    }
266}