geoarrow_array/array/
wkb_view.rs

Help
1use std::sync::Arc;
2
3use arrow_array::builder::BinaryViewBuilder;
4use arrow_array::cast::AsArray;
5use arrow_array::{Array, ArrayRef, BinaryViewArray, OffsetSizeTrait};
6use arrow_buffer::NullBuffer;
7use arrow_schema::{DataType, Field};
8use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
9use geoarrow_schema::{GeoArrowType, Metadata, WkbType};
10use wkb::reader::Wkb;
11
12use crate::array::GenericWkbArray;
13use crate::{GeoArrowArray, GeoArrowArrayAccessor, IntoArrow};
14
15/// An immutable array of WKB geometries.
16///
17/// This is stored as an Arrow [`BinaryViewArray`] and is semantically equivalent to
18/// `Vec<Option<Wkb>>` due to the internal validity bitmap.
19///
20/// Refer to [`crate::cast`] for converting this array to other GeoArrow array types.
21#[derive(Debug, Clone, PartialEq)]
22pub struct WkbViewArray {
23    pub(crate) data_type: WkbType,
24    pub(crate) array: BinaryViewArray,
25}
26
27impl WkbViewArray {
28    /// Create a new GenericWkbArray from a BinaryArray
29    pub fn new(array: BinaryViewArray, metadata: Arc<Metadata>) -> Self {
30        Self {
31            data_type: WkbType::new(metadata),
32            array,
33        }
34    }
35
36    /// Returns true if the array is empty
37    pub fn is_empty(&self) -> bool {
38        self.len() == 0
39    }
40
41    /// Slice this [`GenericWkbArray`].
42    ///
43    ///
44    /// # Panic
45    /// This function panics iff `offset + length > self.len()`.
46    #[inline]
47    pub fn slice(&self, offset: usize, length: usize) -> Self {
48        assert!(
49            offset + length <= self.len(),
50            "offset + length may not exceed length of array"
51        );
52        Self {
53            array: self.array.slice(offset, length),
54            data_type: self.data_type.clone(),
55        }
56    }
57
58    /// Replace the [Metadata] in the array with the given metadata
59    pub fn with_metadata(&self, metadata: Arc<Metadata>) -> Self {
60        let mut arr = self.clone();
61        arr.data_type = self.data_type.clone().with_metadata(metadata);
62        arr
63    }
64}
65
66impl GeoArrowArray for WkbViewArray {
67    fn as_any(&self) -> &dyn std::any::Any {
68        self
69    }
70
71    fn into_array_ref(self) -> ArrayRef {
72        Arc::new(self.into_arrow())
73    }
74
75    fn to_array_ref(&self) -> ArrayRef {
76        self.clone().into_array_ref()
77    }
78
79    #[inline]
80    fn len(&self) -> usize {
81        self.array.len()
82    }
83
84    #[inline]
85    fn logical_nulls(&self) -> Option<NullBuffer> {
86        self.array.logical_nulls()
87    }
88
89    #[inline]
90    fn logical_null_count(&self) -> usize {
91        self.array.logical_null_count()
92    }
93
94    #[inline]
95    fn is_null(&self, i: usize) -> bool {
96        self.array.is_null(i)
97    }
98
99    fn data_type(&self) -> GeoArrowType {
100        GeoArrowType::WkbView(self.data_type.clone())
101    }
102
103    fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
104        Arc::new(self.slice(offset, length))
105    }
106
107    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
108        Arc::new(Self::with_metadata(&self, metadata))
109    }
110}
111
112impl<'a> GeoArrowArrayAccessor<'a> for WkbViewArray {
113    type Item = Wkb<'a>;
114
115    unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
116        let buf = self.array.value(index);
117        Wkb::try_new(buf).map_err(|err| GeoArrowError::External(Box::new(err)))
118    }
119}
120
121impl IntoArrow for WkbViewArray {
122    type ArrowArray = BinaryViewArray;
123    type ExtensionType = WkbType;
124
125    fn into_arrow(self) -> Self::ArrowArray {
126        self.array
127    }
128
129    fn extension_type(&self) -> &Self::ExtensionType {
130        &self.data_type
131    }
132}
133
134impl From<(BinaryViewArray, WkbType)> for WkbViewArray {
135    fn from((value, typ): (BinaryViewArray, WkbType)) -> Self {
136        Self {
137            data_type: typ,
138            array: value,
139        }
140    }
141}
142
143impl TryFrom<(&dyn Array, WkbType)> for WkbViewArray {
144    type Error = GeoArrowError;
145
146    fn try_from((value, typ): (&dyn Array, WkbType)) -> GeoArrowResult<Self> {
147        match value.data_type() {
148            DataType::BinaryView => Ok((value.as_binary_view().clone(), typ).into()),
149            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
150                "Unexpected WkbView DataType: {:?}",
151                dt
152            ))),
153        }
154    }
155}
156
157impl TryFrom<(&dyn Array, &Field)> for WkbViewArray {
158    type Error = GeoArrowError;
159
160    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
161        let typ = field
162            .try_extension_type::<WkbType>()
163            .ok()
164            .unwrap_or_default();
165        (arr, typ).try_into()
166    }
167}
168
169impl<O: OffsetSizeTrait> From<GenericWkbArray<O>> for WkbViewArray {
170    fn from(value: GenericWkbArray<O>) -> Self {
171        let wkb_type = value.data_type;
172        let binary_view_array = value.array;
173
174        // Copy the bytes from the binary view array into a new byte array
175        let mut builder = BinaryViewBuilder::new();
176        binary_view_array
177            .iter()
178            .for_each(|value| builder.append_option(value));
179
180        Self {
181            data_type: wkb_type,
182            array: builder.finish(),
183        }
184    }
185}
geoarrow_array/array/wkb_view.rs

geoarrow_array/array/
wkb_view.rs
