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