geoarrow_array/builder/
wkb.rs

1use arrow_array::OffsetSizeTrait;
2use arrow_array::builder::GenericBinaryBuilder;
3use geo_traits::GeometryTrait;
4use geoarrow_schema::WkbType;
5use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
6use wkb::Endianness;
7use wkb::writer::{WriteOptions, write_geometry};
8
9use crate::array::GenericWkbArray;
10use crate::capacity::WkbCapacity;
11
12/// The GeoArrow equivalent to `Vec<Option<Wkb>>`: a mutable collection of Wkb buffers.
13///
14/// Converting a [`WkbBuilder`] into a [`GenericWkbArray`] is `O(1)`.
15#[derive(Debug)]
16pub struct WkbBuilder<O: OffsetSizeTrait>(GenericBinaryBuilder<O>, WkbType);
17
18impl<O: OffsetSizeTrait> WkbBuilder<O> {
19    /// Creates a new empty [`WkbBuilder`].
20    pub fn new(typ: WkbType) -> Self {
21        Self::with_capacity(typ, Default::default())
22    }
23
24    /// Initializes a new [`WkbBuilder`] with a pre-allocated capacity of slots and values.
25    pub fn with_capacity(typ: WkbType, capacity: WkbCapacity) -> Self {
26        Self(
27            GenericBinaryBuilder::with_capacity(
28                capacity.offsets_capacity,
29                capacity.buffer_capacity,
30            ),
31            typ,
32        )
33    }
34
35    // Upstream APIs don't exist for this yet. To implement this without upstream changes, we could
36    // change to using manual `Vec`'s ourselves
37    // pub fn reserve(&mut self, capacity: WkbCapacity) {
38    // }
39
40    /// Push a Geometry onto the end of this builder
41    #[inline]
42    pub fn push_geometry(
43        &mut self,
44        geom: Option<&impl GeometryTrait<T = f64>>,
45    ) -> GeoArrowResult<()> {
46        if let Some(geom) = geom {
47            let wkb_options = WriteOptions {
48                endianness: Endianness::LittleEndian,
49            };
50            write_geometry(&mut self.0, geom, &wkb_options)
51                .map_err(|err| GeoArrowError::Wkb(err.to_string()))?;
52            self.0.append_value("")
53        } else {
54            self.0.append_null()
55        };
56        Ok(())
57    }
58
59    /// Extend this builder from an iterator of Geometries.
60    pub fn extend_from_iter<'a>(
61        &mut self,
62        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = f64> + 'a)>>,
63    ) -> GeoArrowResult<()> {
64        geoms
65            .into_iter()
66            .try_for_each(|maybe_geom| self.push_geometry(maybe_geom))?;
67        Ok(())
68    }
69
70    /// Create this builder from a slice of nullable Geometries.
71    pub fn from_nullable_geometries(
72        geoms: &[Option<impl GeometryTrait<T = f64>>],
73        typ: WkbType,
74    ) -> GeoArrowResult<Self> {
75        let capacity = WkbCapacity::from_geometries(geoms.iter().map(|x| x.as_ref()));
76        let mut array = Self::with_capacity(typ, capacity);
77        array.extend_from_iter(geoms.iter().map(|x| x.as_ref()))?;
78        Ok(array)
79    }
80
81    /// Consume this builder and convert to a [GenericWkbArray].
82    ///
83    /// This is `O(1)`.
84    pub fn finish(mut self) -> GenericWkbArray<O> {
85        GenericWkbArray::new(self.0.finish(), self.1.metadata().clone())
86    }
87}
geoarrow_array/builder/wkb.rs

geoarrow_array/builder/
wkb.rs
