geoarrow_array/builder/
linestring.rs

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

geoarrow_array/builder/
linestring.rs
