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