geoarrow_array/array/
wkb.rs

1use std::sync::Arc;
2
3use arrow_array::builder::GenericByteBuilder;
4use arrow_array::cast::AsArray;
5use arrow_array::{
6    Array, ArrayRef, BinaryArray, GenericBinaryArray, LargeBinaryArray, OffsetSizeTrait,
7};
8use arrow_buffer::NullBuffer;
9use arrow_schema::{DataType, Field};
10use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
11use geoarrow_schema::{GeoArrowType, Metadata, WkbType};
12use wkb::reader::Wkb;
13
14use crate::array::WkbViewArray;
15use crate::capacity::WkbCapacity;
16use crate::trait_::{GeoArrowArray, GeoArrowArrayAccessor, IntoArrow};
17use crate::util::{offsets_buffer_i32_to_i64, offsets_buffer_i64_to_i32};
18
19/// An immutable array of WKB geometries.
20///
21/// This is stored either as an Arrow [`BinaryArray`] or [`LargeBinaryArray`] and is semantically
22/// equivalent to `Vec<Option<Wkb>>` due to the internal validity bitmap.
23///
24/// Refer to [`crate::cast`] for converting this array to other GeoArrow array types.
25#[derive(Debug, Clone, PartialEq)]
26pub struct GenericWkbArray<O: OffsetSizeTrait> {
27    pub(crate) data_type: WkbType,
28    pub(crate) array: GenericBinaryArray<O>,
29}
30
31// Implement geometry accessors
32impl<O: OffsetSizeTrait> GenericWkbArray<O> {
33    /// Create a new GenericWkbArray from a BinaryArray
34    pub fn new(array: GenericBinaryArray<O>, metadata: Arc<Metadata>) -> Self {
35        Self {
36            data_type: WkbType::new(metadata),
37            array,
38        }
39    }
40
41    /// Returns true if the array is empty
42    pub fn is_empty(&self) -> bool {
43        self.len() == 0
44    }
45
46    /// The lengths of each buffer contained in this array.
47    pub fn buffer_lengths(&self) -> WkbCapacity {
48        WkbCapacity::new(
49            self.array.offsets().last().unwrap().to_usize().unwrap(),
50            self.len(),
51        )
52    }
53
54    /// The number of bytes occupied by this array.
55    pub fn num_bytes(&self) -> usize {
56        let validity_len = self
57            .array
58            .nulls()
59            .as_ref()
60            .map(|v| v.buffer().len())
61            .unwrap_or(0);
62        validity_len + self.buffer_lengths().num_bytes::<O>()
63    }
64
65    /// Slice this [`GenericWkbArray`].
66    ///
67    ///
68    /// # Panic
69    /// This function panics iff `offset + length > self.len()`.
70    #[inline]
71    pub fn slice(&self, offset: usize, length: usize) -> Self {
72        assert!(
73            offset + length <= self.len(),
74            "offset + length may not exceed length of array"
75        );
76        Self {
77            array: self.array.slice(offset, length),
78            data_type: self.data_type.clone(),
79        }
80    }
81
82    /// Replace the [Metadata] in the array with the given metadata
83    pub fn with_metadata(&self, metadata: Arc<Metadata>) -> Self {
84        let mut arr = self.clone();
85        arr.data_type = self.data_type.clone().with_metadata(metadata);
86        arr
87    }
88}
89
90impl<O: OffsetSizeTrait> GeoArrowArray for GenericWkbArray<O> {
91    fn as_any(&self) -> &dyn std::any::Any {
92        self
93    }
94
95    fn into_array_ref(self) -> ArrayRef {
96        Arc::new(self.into_arrow())
97    }
98
99    fn to_array_ref(&self) -> ArrayRef {
100        self.clone().into_array_ref()
101    }
102
103    #[inline]
104    fn len(&self) -> usize {
105        self.array.len()
106    }
107
108    #[inline]
109    fn logical_nulls(&self) -> Option<NullBuffer> {
110        self.array.logical_nulls()
111    }
112
113    #[inline]
114    fn logical_null_count(&self) -> usize {
115        self.array.logical_null_count()
116    }
117
118    #[inline]
119    fn is_null(&self, i: usize) -> bool {
120        self.array.is_null(i)
121    }
122
123    fn data_type(&self) -> GeoArrowType {
124        if O::IS_LARGE {
125            GeoArrowType::LargeWkb(self.data_type.clone())
126        } else {
127            GeoArrowType::Wkb(self.data_type.clone())
128        }
129    }
130
131    fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
132        Arc::new(self.slice(offset, length))
133    }
134
135    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
136        Arc::new(Self::with_metadata(&self, metadata))
137    }
138}
139
140impl<'a, O: OffsetSizeTrait> GeoArrowArrayAccessor<'a> for GenericWkbArray<O> {
141    type Item = Wkb<'a>;
142
143    unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
144        let buf = self.array.value(index);
145        Wkb::try_new(buf).map_err(|err| GeoArrowError::External(Box::new(err)))
146    }
147}
148
149impl<O: OffsetSizeTrait> IntoArrow for GenericWkbArray<O> {
150    type ArrowArray = GenericBinaryArray<O>;
151    type ExtensionType = WkbType;
152
153    fn into_arrow(self) -> Self::ArrowArray {
154        self.array
155    }
156
157    fn extension_type(&self) -> &Self::ExtensionType {
158        &self.data_type
159    }
160}
161
162impl<O: OffsetSizeTrait> From<(GenericBinaryArray<O>, WkbType)> for GenericWkbArray<O> {
163    fn from((value, typ): (GenericBinaryArray<O>, WkbType)) -> Self {
164        Self {
165            data_type: typ,
166            array: value,
167        }
168    }
169}
170
171impl TryFrom<(&dyn Array, WkbType)> for GenericWkbArray<i32> {
172    type Error = GeoArrowError;
173    fn try_from((value, typ): (&dyn Array, WkbType)) -> GeoArrowResult<Self> {
174        match value.data_type() {
175            DataType::Binary => Ok((value.as_binary::<i32>().clone(), typ).into()),
176            DataType::LargeBinary => {
177                let geom_array: GenericWkbArray<i64> =
178                    (value.as_binary::<i64>().clone(), typ).into();
179                geom_array.try_into()
180            }
181            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
182                "Unexpected GenericWkbArray DataType: {:?}",
183                dt
184            ))),
185        }
186    }
187}
188
189impl TryFrom<(&dyn Array, WkbType)> for GenericWkbArray<i64> {
190    type Error = GeoArrowError;
191    fn try_from((value, typ): (&dyn Array, WkbType)) -> GeoArrowResult<Self> {
192        match value.data_type() {
193            DataType::Binary => {
194                let geom_array: GenericWkbArray<i32> =
195                    (value.as_binary::<i32>().clone(), typ).into();
196                Ok(geom_array.into())
197            }
198            DataType::LargeBinary => Ok((value.as_binary::<i64>().clone(), typ).into()),
199            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
200                "Unexpected GenericWkbArray DataType: {:?}",
201                dt
202            ))),
203        }
204    }
205}
206
207impl TryFrom<(&dyn Array, &Field)> for GenericWkbArray<i32> {
208    type Error = GeoArrowError;
209
210    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
211        let typ = field
212            .try_extension_type::<WkbType>()
213            .ok()
214            .unwrap_or_default();
215        (arr, typ).try_into()
216    }
217}
218
219impl TryFrom<(&dyn Array, &Field)> for GenericWkbArray<i64> {
220    type Error = GeoArrowError;
221
222    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
223        let typ = field
224            .try_extension_type::<WkbType>()
225            .ok()
226            .unwrap_or_default();
227        (arr, typ).try_into()
228    }
229}
230
231impl From<GenericWkbArray<i32>> for GenericWkbArray<i64> {
232    fn from(value: GenericWkbArray<i32>) -> Self {
233        let binary_array = value.array;
234        let (offsets, values, nulls) = binary_array.into_parts();
235        let array = LargeBinaryArray::new(offsets_buffer_i32_to_i64(&offsets), values, nulls);
236        Self {
237            data_type: value.data_type,
238            array,
239        }
240    }
241}
242
243impl TryFrom<GenericWkbArray<i64>> for GenericWkbArray<i32> {
244    type Error = GeoArrowError;
245
246    fn try_from(value: GenericWkbArray<i64>) -> GeoArrowResult<Self> {
247        let binary_array = value.array;
248        let (offsets, values, nulls) = binary_array.into_parts();
249        let array = BinaryArray::new(offsets_buffer_i64_to_i32(&offsets)?, values, nulls);
250        Ok(Self {
251            data_type: value.data_type,
252            array,
253        })
254    }
255}
256
257impl<O: OffsetSizeTrait> From<WkbViewArray> for GenericWkbArray<O> {
258    fn from(value: WkbViewArray) -> Self {
259        let wkb_type = value.data_type;
260        let binary_view_array = value.array;
261
262        // Copy the bytes from the binary view array into a new byte array
263        let mut builder = GenericByteBuilder::new();
264        binary_view_array
265            .iter()
266            .for_each(|value| builder.append_option(value));
267
268        Self {
269            data_type: wkb_type,
270            array: builder.finish(),
271        }
272    }
273}
274
275/// A [`GenericWkbArray`] using `i32` offsets
276///
277/// The byte length of each element is represented by an i32.
278///
279/// See [`GenericWkbArray`] for more information and examples
280pub type WkbArray = GenericWkbArray<i32>;
281
282/// A [`GenericWkbArray`] using `i64` offsets
283///
284/// The byte length of each element is represented by an i64.
285///
286/// See [`GenericWkbArray`] for more information and examples
287pub type LargeWkbArray = GenericWkbArray<i64>;
288
289#[cfg(test)]
290mod test {
291    use arrow_array::builder::{BinaryBuilder, LargeBinaryBuilder};
292
293    use super::*;
294    use crate::GeoArrowArray;
295    use crate::builder::WkbBuilder;
296    use crate::test::point;
297
298    fn wkb_data<O: OffsetSizeTrait>() -> GenericWkbArray<O> {
299        let mut builder = WkbBuilder::new(WkbType::new(Default::default()));
300        builder.push_geometry(Some(&point::p0()));
301        builder.push_geometry(Some(&point::p1()));
302        builder.push_geometry(Some(&point::p2()));
303        builder.finish()
304    }
305
306    #[test]
307    fn parse_dyn_array_i32() {
308        let wkb_array = wkb_data::<i32>();
309        let array = wkb_array.to_array_ref();
310        let field = Field::new("geometry", array.data_type().clone(), true)
311            .with_extension_type(wkb_array.data_type.clone());
312        let wkb_array_retour: GenericWkbArray<i32> = (array.as_ref(), &field).try_into().unwrap();
313
314        assert_eq!(wkb_array, wkb_array_retour);
315    }
316
317    #[test]
318    fn parse_dyn_array_i64() {
319        let wkb_array = wkb_data::<i64>();
320        let array = wkb_array.to_array_ref();
321        let field = Field::new("geometry", array.data_type().clone(), true)
322            .with_extension_type(wkb_array.data_type.clone());
323        let wkb_array_retour: GenericWkbArray<i64> = (array.as_ref(), &field).try_into().unwrap();
324
325        assert_eq!(wkb_array, wkb_array_retour);
326    }
327
328    #[test]
329    fn convert_i32_to_i64() {
330        let wkb_array = wkb_data::<i32>();
331        let wkb_array_i64: GenericWkbArray<i64> = wkb_array.clone().into();
332        let wkb_array_i32: GenericWkbArray<i32> = wkb_array_i64.clone().try_into().unwrap();
333
334        assert_eq!(wkb_array, wkb_array_i32);
335    }
336
337    #[test]
338    fn convert_i64_to_i32_to_i64() {
339        let wkb_array = wkb_data::<i64>();
340        let wkb_array_i32: GenericWkbArray<i32> = wkb_array.clone().try_into().unwrap();
341        let wkb_array_i64: GenericWkbArray<i64> = wkb_array_i32.clone().into();
342
343        assert_eq!(wkb_array, wkb_array_i64);
344    }
345
346    /// Passing a field without an extension name should not panic
347    #[test]
348    fn allow_field_without_extension_name() {
349        // String array
350        let mut builder = BinaryBuilder::new();
351        builder.append_value(b"a");
352        let array = Arc::new(builder.finish()) as ArrayRef;
353        let field = Field::new("geometry", array.data_type().clone(), true);
354        let _wkt_arr = GenericWkbArray::<i32>::try_from((array.as_ref(), &field)).unwrap();
355
356        // Large string
357        let mut builder = LargeBinaryBuilder::new();
358        builder.append_value(b"a");
359        let array = Arc::new(builder.finish()) as ArrayRef;
360        let field = Field::new("geometry", array.data_type().clone(), true);
361        let _wkt_arr = GenericWkbArray::<i64>::try_from((array.as_ref(), &field)).unwrap();
362    }
363}
geoarrow_array/array/wkb.rs

geoarrow_array/array/
wkb.rs
