geoarrow_array/array/
wkt_view.rs

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

geoarrow_array/array/
wkt_view.rs
