geoarrow_array/builder/
wkb.rs

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