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::reader::Wkb;
8use wkb::writer::{WriteOptions, write_geometry};
9
10use crate::array::GenericWkbArray;
11use crate::capacity::WkbCapacity;
12
13/// The GeoArrow equivalent to `Vec<Option<Wkb>>`: a mutable collection of Wkb buffers.
14///
15/// Converting a [`WkbBuilder`] into a [`GenericWkbArray`] is `O(1)`.
16#[derive(Debug)]
17pub struct WkbBuilder<O: OffsetSizeTrait>(GenericBinaryBuilder<O>, WkbType);
18
19impl<O: OffsetSizeTrait> WkbBuilder<O> {
20    /// Creates a new empty [`WkbBuilder`].
21    pub fn new(typ: WkbType) -> Self {
22        Self::with_capacity(typ, Default::default())
23    }
24
25    /// Initializes a new [`WkbBuilder`] with a pre-allocated capacity of slots and values.
26    pub fn with_capacity(typ: WkbType, capacity: WkbCapacity) -> Self {
27        Self(
28            GenericBinaryBuilder::with_capacity(
29                capacity.offsets_capacity,
30                capacity.buffer_capacity,
31            ),
32            typ,
33        )
34    }
35
36    // Upstream APIs don't exist for this yet. To implement this without upstream changes, we could
37    // change to using manual `Vec`'s ourselves
38    // pub fn reserve(&mut self, capacity: WkbCapacity) {
39    // }
40
41    /// Push a Geometry onto the end of this builder
42    #[inline]
43    pub fn push_geometry(
44        &mut self,
45        geom: Option<&impl GeometryTrait<T = f64>>,
46    ) -> GeoArrowResult<()> {
47        if let Some(geom) = geom {
48            let wkb_options = WriteOptions {
49                endianness: Endianness::LittleEndian,
50            };
51            write_geometry(&mut self.0, geom, &wkb_options)
52                .map_err(|err| GeoArrowError::Wkb(err.to_string()))?;
53            self.0.append_value("")
54        } else {
55            self.0.append_null()
56        };
57        Ok(())
58    }
59
60    /// Extend this builder from an iterator of Geometries.
61    pub fn extend_from_iter<'a>(
62        &mut self,
63        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = f64> + 'a)>>,
64    ) -> GeoArrowResult<()> {
65        geoms
66            .into_iter()
67            .try_for_each(|maybe_geom| self.push_geometry(maybe_geom))?;
68        Ok(())
69    }
70
71    /// Create this builder from a slice of nullable Geometries.
72    pub fn from_nullable_geometries(
73        geoms: &[Option<impl GeometryTrait<T = f64>>],
74        typ: WkbType,
75    ) -> GeoArrowResult<Self> {
76        let capacity = WkbCapacity::from_geometries(geoms.iter().map(|x| x.as_ref()));
77        let mut array = Self::with_capacity(typ, capacity);
78        array.extend_from_iter(geoms.iter().map(|x| x.as_ref()))?;
79        Ok(array)
80    }
81
82    /// Push raw WKB bytes onto the end of this builder.
83    ///
84    /// This method validates that the input bytes represent valid WKB before appending.
85    /// If the bytes are `None`, a null value is appended.
86    ///
87    /// # Errors
88    ///
89    /// Returns an error if the input bytes are not valid WKB format.
90    ///
91    /// # Example
92    ///
93    /// ```
94    /// use geoarrow_array::builder::WkbBuilder;
95    /// use geoarrow_array::GeoArrowArray;
96    /// use geoarrow_schema::WkbType;
97    ///
98    /// let mut builder = WkbBuilder::<i32>::new(WkbType::default());
99    ///
100    /// // Valid WKB for a Point(1.0, 2.0) in little-endian
101    /// let wkb_bytes = vec![
102    ///     0x01, // Little-endian
103    ///     0x01, 0x00, 0x00, 0x00, // Point type
104    ///     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0, 0x3f, // x = 1.0
105    ///     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, // y = 2.0
106    /// ];
107    ///
108    /// builder.push_wkb(Some(&wkb_bytes)).unwrap();
109    /// builder.push_wkb(None).unwrap(); // Append null
110    ///
111    /// let array = builder.finish();
112    /// assert_eq!(array.len(), 2);
113    /// ```
114    #[inline]
115    pub fn push_wkb(&mut self, wkb: Option<&[u8]>) -> GeoArrowResult<()> {
116        if let Some(bytes) = wkb {
117            // Validate that the bytes are valid WKB
118            Wkb::try_new(bytes).map_err(|err| GeoArrowError::Wkb(err.to_string()))?;
119            self.0.append_value(bytes);
120        } else {
121            self.0.append_null();
122        }
123        Ok(())
124    }
125
126    /// Push raw WKB bytes onto the end of this builder without validation.
127    ///
128    /// This method directly appends the input bytes to the underlying buffer without
129    /// validating that they represent valid WKB. If the bytes are `None`, a null value
130    /// is appended.
131    ///
132    /// # Safety
133    ///
134    /// This function is unsafe because it does not validate that the input bytes are
135    /// valid WKB format. Calling this with invalid WKB data may result in undefined
136    /// behavior when the resulting array is used with operations that assume valid WKB.
137    ///
138    /// The caller must ensure that:
139    /// - The bytes represent valid WKB according to the OGC WKB specification
140    /// - The byte order (endianness) is correctly specified in the WKB header
141    /// - The geometry type and coordinates are properly encoded
142    ///
143    /// # Example
144    ///
145    /// ```
146    /// use geoarrow_array::builder::WkbBuilder;
147    /// use geoarrow_array::GeoArrowArray;
148    /// use geoarrow_schema::WkbType;
149    ///
150    /// let mut builder = WkbBuilder::<i32>::new(WkbType::default());
151    ///
152    /// // Valid WKB for a Point(1.0, 2.0) in little-endian
153    /// let wkb_bytes = vec![
154    ///     0x01, // Little-endian
155    ///     0x01, 0x00, 0x00, 0x00, // Point type
156    ///     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0, 0x3f, // x = 1.0
157    ///     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, // y = 2.0
158    /// ];
159    ///
160    /// unsafe {
161    ///     builder.push_wkb_unchecked(Some(&wkb_bytes));
162    ///     builder.push_wkb_unchecked(None); // Append null
163    /// }
164    ///
165    /// let array = builder.finish();
166    /// assert_eq!(array.len(), 2);
167    /// ```
168    #[inline]
169    pub unsafe fn push_wkb_unchecked(&mut self, wkb: Option<&[u8]>) {
170        if let Some(bytes) = wkb {
171            self.0.append_value(bytes);
172        } else {
173            self.0.append_null();
174        }
175    }
176
177    /// Consume this builder and convert to a [GenericWkbArray].
178    ///
179    /// This is `O(1)`.
180    pub fn finish(mut self) -> GenericWkbArray<O> {
181        GenericWkbArray::new(self.0.finish(), self.1.metadata().clone())
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use crate::trait_::GeoArrowArray;
189
190    /// Valid WKB for Point(1.0, 2.0) in little-endian format
191    fn point_wkb() -> Vec<u8> {
192        let point = geo::Point::new(1.0, 2.0);
193        let mut buf = Vec::new();
194        wkb::writer::write_point(&mut buf, &point, &Default::default()).unwrap();
195        buf
196    }
197
198    /// Valid WKB for Point(3.0, 4.0) in little-endian format
199    fn point_wkb_2() -> Vec<u8> {
200        let point = geo::Point::new(3.0, 4.0);
201        let mut buf = Vec::new();
202        wkb::writer::write_point(&mut buf, &point, &Default::default()).unwrap();
203        buf
204    }
205
206    /// Invalid WKB (too short)
207    fn invalid_wkb() -> Vec<u8> {
208        vec![0x01, 0x01]
209    }
210
211    #[test]
212    fn test_push_raw_valid() {
213        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
214        let wkb = point_wkb();
215
216        // Should succeed with valid WKB
217        builder.push_wkb(Some(&wkb)).unwrap();
218
219        let array = builder.finish();
220        assert_eq!(array.len(), 1);
221        assert!(!array.is_null(0));
222    }
223
224    #[test]
225    fn test_push_raw_multiple() {
226        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
227        let wkb1 = point_wkb();
228        let wkb2 = point_wkb_2();
229
230        builder.push_wkb(Some(&wkb1)).unwrap();
231        builder.push_wkb(Some(&wkb2)).unwrap();
232
233        let array = builder.finish();
234        assert_eq!(array.len(), 2);
235        assert!(!array.is_null(0));
236        assert!(!array.is_null(1));
237    }
238
239    #[test]
240    fn test_push_raw_null() {
241        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
242
243        // Push null value
244        builder.push_wkb(None).unwrap();
245
246        let array = builder.finish();
247        assert_eq!(array.len(), 1);
248        assert!(array.is_null(0));
249    }
250
251    #[test]
252    fn test_push_raw_mixed_with_nulls() {
253        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
254        let wkb = point_wkb();
255
256        builder.push_wkb(Some(&wkb)).unwrap();
257        builder.push_wkb(None).unwrap();
258        builder.push_wkb(Some(&wkb)).unwrap();
259
260        let array = builder.finish();
261        assert_eq!(array.len(), 3);
262        assert!(!array.is_null(0));
263        assert!(array.is_null(1));
264        assert!(!array.is_null(2));
265    }
266
267    #[test]
268    fn test_push_raw_invalid() {
269        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
270        let invalid = invalid_wkb();
271
272        // Should fail with invalid WKB
273        let result = builder.push_wkb(Some(&invalid));
274        assert!(result.is_err());
275    }
276
277    #[test]
278    fn test_push_raw_unchecked_valid() {
279        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
280        let wkb = point_wkb();
281
282        unsafe {
283            builder.push_wkb_unchecked(Some(&wkb));
284        }
285
286        let array = builder.finish();
287        assert_eq!(array.len(), 1);
288        assert!(!array.is_null(0));
289    }
290
291    #[test]
292    fn test_push_raw_unchecked_null() {
293        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
294
295        unsafe {
296            builder.push_wkb_unchecked(None);
297        }
298
299        let array = builder.finish();
300        assert_eq!(array.len(), 1);
301        assert!(array.is_null(0));
302    }
303
304    #[test]
305    fn test_push_raw_unchecked_multiple() {
306        let mut builder = WkbBuilder::<i32>::new(WkbType::default());
307        let wkb1 = point_wkb();
308        let wkb2 = point_wkb_2();
309
310        unsafe {
311            builder.push_wkb_unchecked(Some(&wkb1));
312            builder.push_wkb_unchecked(None);
313            builder.push_wkb_unchecked(Some(&wkb2));
314        }
315
316        let array = builder.finish();
317        assert_eq!(array.len(), 3);
318        assert!(!array.is_null(0));
319        assert!(array.is_null(1));
320        assert!(!array.is_null(2));
321    }
322
323    #[test]
324    fn test_push_raw_with_i64_offset() {
325        let mut builder = WkbBuilder::<i64>::new(WkbType::default());
326        let wkb = point_wkb();
327
328        builder.push_wkb(Some(&wkb)).unwrap();
329        builder.push_wkb(None).unwrap();
330
331        let array = builder.finish();
332        assert_eq!(array.len(), 2);
333        assert!(!array.is_null(0));
334        assert!(array.is_null(1));
335    }
336}
geoarrow_array/builder/wkb.rs

geoarrow_array/builder/
wkb.rs
