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

geoarrow_array/builder/
linestring.rs
