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    /// Consume the builder and convert to an immutable [`PointArray`]
66    pub fn finish(mut self) -> PointArray {
67        let validity = self.validity.finish();
68        PointArray::new(
69            self.coords.finish(),
70            validity,
71            self.data_type.metadata().clone(),
72        )
73    }
74
75    /// Add a new coord to the end of this array, interpreting the coord as a non-empty point.
76    ///
77    /// ## Panics
78    ///
79    /// - If the added coordinate does not have the same dimension as the point array.
80    #[inline]
81    pub fn push_coord(&mut self, value: Option<&impl CoordTrait<T = f64>>) {
82        self.try_push_coord(value).unwrap()
83    }
84
85    /// Add a new point to the end of this array.
86    ///
87    /// ## Panics
88    ///
89    /// - If the added point does not have the same dimension as the point array.
90    #[inline]
91    pub fn push_point(&mut self, value: Option<&impl PointTrait<T = f64>>) {
92        self.try_push_point(value).unwrap()
93    }
94
95    /// Add a new coord to the end of this array, where the coord is a non-empty point
96    ///
97    /// ## Errors
98    ///
99    /// - If the added coordinate does not have the same dimension as the point array.
100    #[inline]
101    pub fn try_push_coord(
102        &mut self,
103        value: Option<&impl CoordTrait<T = f64>>,
104    ) -> GeoArrowResult<()> {
105        if let Some(value) = value {
106            self.coords.try_push_coord(value)?;
107            self.validity.append(true);
108        } else {
109            self.push_null()
110        };
111        Ok(())
112    }
113
114    /// Add a new point to the end of this array.
115    ///
116    /// ## Errors
117    ///
118    /// - If the added point does not have the same dimension as the point array.
119    #[inline]
120    pub fn try_push_point(
121        &mut self,
122        value: Option<&impl PointTrait<T = f64>>,
123    ) -> GeoArrowResult<()> {
124        if let Some(value) = value {
125            self.coords.try_push_point(value)?;
126            self.validity.append(true);
127        } else {
128            self.push_null()
129        };
130        Ok(())
131    }
132
133    /// Add a valid but empty point to the end of this array.
134    #[inline]
135    pub fn push_empty(&mut self) {
136        self.coords.push_constant(f64::NAN);
137        self.validity.append_non_null();
138    }
139
140    /// Add a new null value to the end of this array.
141    #[inline]
142    pub fn push_null(&mut self) {
143        self.coords.push_constant(f64::NAN);
144        self.validity.append_null();
145    }
146
147    /// Add a new geometry to this builder
148    ///
149    /// This will error if the geometry type is not Point or a MultiPoint with length 1.
150    #[inline]
151    pub fn push_geometry(
152        &mut self,
153        value: Option<&impl GeometryTrait<T = f64>>,
154    ) -> GeoArrowResult<()> {
155        if let Some(value) = value {
156            match value.as_type() {
157                GeometryType::Point(p) => self.push_point(Some(p)),
158                GeometryType::MultiPoint(mp) => {
159                    let num_points = mp.num_points();
160                    if num_points == 0 {
161                        self.push_empty();
162                    } else if num_points == 1 {
163                        self.push_point(Some(&mp.point(0).unwrap()))
164                    } else {
165                        return Err(GeoArrowError::IncorrectGeometryType(format!(
166                            "Expected MultiPoint with only one point in PointBuilder, got {} points",
167                            num_points
168                        )));
169                    }
170                }
171                gt => {
172                    return Err(GeoArrowError::IncorrectGeometryType(format!(
173                        "Expected point, got {}",
174                        gt.name()
175                    )));
176                }
177            }
178        } else {
179            self.push_null()
180        };
181        Ok(())
182    }
183
184    /// Extend this builder with the given geometries
185    pub fn extend_from_iter<'a>(
186        &mut self,
187        geoms: impl Iterator<Item = Option<&'a (impl PointTrait<T = f64> + 'a)>>,
188    ) {
189        geoms
190            .into_iter()
191            .for_each(|maybe_polygon| self.push_point(maybe_polygon));
192    }
193
194    /// Extend this builder with the given geometries
195    pub fn extend_from_geometry_iter<'a>(
196        &mut self,
197        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = f64> + 'a)>>,
198    ) -> GeoArrowResult<()> {
199        geoms.into_iter().try_for_each(|g| self.push_geometry(g))?;
200        Ok(())
201    }
202
203    /// Construct a new builder, pre-filling it with the provided geometries
204    pub fn from_points<'a>(
205        geoms: impl ExactSizeIterator<Item = &'a (impl PointTrait<T = f64> + 'a)>,
206        typ: PointType,
207    ) -> Self {
208        let mut mutable_array = Self::with_capacity(typ, geoms.len());
209        geoms
210            .into_iter()
211            .for_each(|maybe_point| mutable_array.push_point(Some(maybe_point)));
212        mutable_array
213    }
214
215    /// Construct a new builder, pre-filling it with the provided geometries
216    pub fn from_nullable_points<'a>(
217        geoms: impl ExactSizeIterator<Item = Option<&'a (impl PointTrait<T = f64> + 'a)>>,
218        typ: PointType,
219    ) -> Self {
220        let mut mutable_array = Self::with_capacity(typ, geoms.len());
221        geoms
222            .into_iter()
223            .for_each(|maybe_point| mutable_array.push_point(maybe_point));
224        mutable_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: PointType,
231    ) -> GeoArrowResult<Self> {
232        let capacity = geoms.len();
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>, PointType)> for PointBuilder {
240    type Error = GeoArrowError;
241
242    fn try_from((value, typ): (GenericWkbArray<O>, PointType)) -> 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 PointBuilder {
252    fn len(&self) -> usize {
253        self.coords.len()
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}