geoarrow_array/array/
wkt.rs

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

geoarrow_array/array/
wkt.rs
