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

geoarrow_array/builder/
linestring.rs
