geoarrow_array/array/
linestring.rs

1use std::sync::Arc;
2
3use arrow_array::cast::AsArray;
4use arrow_array::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
5use arrow_buffer::{NullBuffer, OffsetBuffer};
6use arrow_schema::{DataType, Field};
7use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8use geoarrow_schema::type_id::GeometryTypeId;
9use geoarrow_schema::{CoordType, Dimension, GeoArrowType, LineStringType, Metadata};
10
11use crate::array::{CoordBuffer, GenericWkbArray};
12use crate::builder::LineStringBuilder;
13use crate::capacity::LineStringCapacity;
14use crate::eq::offset_buffer_eq;
15use crate::scalar::LineString;
16use crate::trait_::{GeoArrowArray, GeoArrowArrayAccessor, IntoArrow};
17use crate::util::{OffsetBufferUtils, offsets_buffer_i64_to_i32};
18
19/// An immutable array of LineString geometries.
20///
21/// This is semantically equivalent to `Vec<Option<LineString>>` due to the internal validity
22/// bitmap.
23#[derive(Debug, Clone)]
24pub struct LineStringArray {
25    pub(crate) data_type: LineStringType,
26
27    pub(crate) coords: CoordBuffer,
28
29    /// Offsets into the coordinate array where each geometry starts
30    pub(crate) geom_offsets: OffsetBuffer<i32>,
31
32    /// Validity bitmap
33    pub(crate) nulls: Option<NullBuffer>,
34}
35
36pub(super) fn check(
37    coords: &CoordBuffer,
38    validity_len: Option<usize>,
39    geom_offsets: &OffsetBuffer<i32>,
40) -> GeoArrowResult<()> {
41    if validity_len.is_some_and(|len| len != geom_offsets.len_proxy()) {
42        return Err(GeoArrowError::InvalidGeoArrow(
43            "nulls mask length must match the number of values".to_string(),
44        ));
45    }
46
47    if *geom_offsets.last() as usize != coords.len() {
48        return Err(GeoArrowError::InvalidGeoArrow(
49            "largest geometry offset must match coords length".to_string(),
50        ));
51    }
52
53    Ok(())
54}
55
56impl LineStringArray {
57    /// Create a new LineStringArray from parts
58    ///
59    /// # Implementation
60    ///
61    /// This function is `O(1)`.
62    ///
63    /// # Panics
64    ///
65    /// - if the nulls is not `None` and its length is different from the number of geometries
66    /// - if the largest geometry offset does not match the number of coordinates
67    pub fn new(
68        coords: CoordBuffer,
69        geom_offsets: OffsetBuffer<i32>,
70        nulls: Option<NullBuffer>,
71        metadata: Arc<Metadata>,
72    ) -> Self {
73        Self::try_new(coords, geom_offsets, nulls, metadata).unwrap()
74    }
75
76    /// Create a new LineStringArray from parts
77    ///
78    /// # Implementation
79    ///
80    /// This function is `O(1)`.
81    ///
82    /// # Errors
83    ///
84    /// - if the nulls buffer does not have the same length as the number of geometries
85    /// - if the geometry offsets do not match the number of coordinates
86    pub fn try_new(
87        coords: CoordBuffer,
88        geom_offsets: OffsetBuffer<i32>,
89        nulls: Option<NullBuffer>,
90        metadata: Arc<Metadata>,
91    ) -> GeoArrowResult<Self> {
92        check(&coords, nulls.as_ref().map(|v| v.len()), &geom_offsets)?;
93        Ok(Self {
94            data_type: LineStringType::new(coords.dim(), metadata)
95                .with_coord_type(coords.coord_type()),
96            coords,
97            geom_offsets,
98            nulls,
99        })
100    }
101
102    /// Access the underlying coordinate buffer
103    pub fn coords(&self) -> &CoordBuffer {
104        &self.coords
105    }
106
107    /// Access the underlying geometry offsets buffer
108    pub fn geom_offsets(&self) -> &OffsetBuffer<i32> {
109        &self.geom_offsets
110    }
111
112    /// The lengths of each buffer contained in this array.
113    pub fn buffer_lengths(&self) -> LineStringCapacity {
114        LineStringCapacity::new(*self.geom_offsets.last() as usize, self.len())
115    }
116
117    /// The number of bytes occupied by this array.
118    pub fn num_bytes(&self) -> usize {
119        let validity_len = self.nulls.as_ref().map(|v| v.buffer().len()).unwrap_or(0);
120        validity_len + self.buffer_lengths().num_bytes(self.data_type.dimension())
121    }
122
123    /// Slice this [`LineStringArray`].
124    ///
125    /// # Implementation
126    ///
127    /// This operation is `O(1)` as it amounts to increasing a few ref counts.
128    ///
129    /// # Panic
130    ///
131    /// This function panics iff `offset + length > self.len()`.
132    #[inline]
133    pub fn slice(&self, offset: usize, length: usize) -> Self {
134        assert!(
135            offset + length <= self.len(),
136            "offset + length may not exceed length of array"
137        );
138        // Note: we **only** slice the geom_offsets and not any actual data. Otherwise the offsets
139        // would be in the wrong location.
140        Self {
141            data_type: self.data_type.clone(),
142            coords: self.coords.clone(),
143            geom_offsets: self.geom_offsets.slice(offset, length),
144            nulls: self.nulls.as_ref().map(|v| v.slice(offset, length)),
145        }
146    }
147
148    /// Change the [`CoordType`] of this array.
149    pub fn into_coord_type(self, coord_type: CoordType) -> Self {
150        Self {
151            data_type: self.data_type.with_coord_type(coord_type),
152            coords: self.coords.into_coord_type(coord_type),
153            ..self
154        }
155    }
156
157    /// Change the [`Metadata`] of this array.
158    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
159        Self {
160            data_type: self.data_type.with_metadata(metadata),
161            ..self
162        }
163    }
164}
165
166impl GeoArrowArray for LineStringArray {
167    fn as_any(&self) -> &dyn std::any::Any {
168        self
169    }
170
171    fn into_array_ref(self) -> ArrayRef {
172        Arc::new(self.into_arrow())
173    }
174
175    fn to_array_ref(&self) -> ArrayRef {
176        self.clone().into_array_ref()
177    }
178
179    #[inline]
180    fn len(&self) -> usize {
181        self.geom_offsets.len_proxy()
182    }
183
184    #[inline]
185    fn logical_nulls(&self) -> Option<NullBuffer> {
186        self.nulls.clone()
187    }
188
189    #[inline]
190    fn logical_null_count(&self) -> usize {
191        self.nulls.as_ref().map(|v| v.null_count()).unwrap_or(0)
192    }
193
194    fn is_null(&self, i: usize) -> bool {
195        self.nulls
196            .as_ref()
197            .map(|n| n.is_null(i))
198            .unwrap_or_default()
199    }
200
201    fn data_type(&self) -> GeoArrowType {
202        GeoArrowType::LineString(self.data_type.clone())
203    }
204
205    fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
206        Arc::new(self.slice(offset, length))
207    }
208
209    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
210        Arc::new(self.with_metadata(metadata))
211    }
212}
213
214impl<'a> GeoArrowArrayAccessor<'a> for LineStringArray {
215    type Item = LineString<'a>;
216
217    unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
218        Ok(LineString::new(&self.coords, &self.geom_offsets, index))
219    }
220}
221
222impl IntoArrow for LineStringArray {
223    type ArrowArray = GenericListArray<i32>;
224    type ExtensionType = LineStringType;
225
226    fn into_arrow(self) -> Self::ArrowArray {
227        let vertices_field = match self.data_type.data_type() {
228            DataType::List(inner_field) => inner_field,
229            _ => unreachable!(),
230        };
231        let nulls = self.nulls;
232        let coord_array = self.coords.into_array_ref();
233        GenericListArray::new(vertices_field, self.geom_offsets, coord_array, nulls)
234    }
235
236    fn extension_type(&self) -> &Self::ExtensionType {
237        &self.data_type
238    }
239}
240
241impl TryFrom<(&GenericListArray<i32>, LineStringType)> for LineStringArray {
242    type Error = GeoArrowError;
243
244    fn try_from((value, typ): (&GenericListArray<i32>, LineStringType)) -> GeoArrowResult<Self> {
245        let coords = CoordBuffer::from_arrow(value.values().as_ref(), typ.dimension())?;
246        let geom_offsets = value.offsets();
247        let nulls = value.nulls();
248
249        Ok(Self::new(
250            coords,
251            geom_offsets.clone(),
252            nulls.cloned(),
253            typ.metadata().clone(),
254        ))
255    }
256}
257
258impl TryFrom<(&GenericListArray<i64>, LineStringType)> for LineStringArray {
259    type Error = GeoArrowError;
260
261    fn try_from((value, typ): (&GenericListArray<i64>, LineStringType)) -> GeoArrowResult<Self> {
262        let coords = CoordBuffer::from_arrow(value.values().as_ref(), typ.dimension())?;
263        let geom_offsets = offsets_buffer_i64_to_i32(value.offsets())?;
264        let nulls = value.nulls();
265
266        Ok(Self::new(
267            coords,
268            geom_offsets,
269            nulls.cloned(),
270            typ.metadata().clone(),
271        ))
272    }
273}
274impl TryFrom<(&dyn Array, LineStringType)> for LineStringArray {
275    type Error = GeoArrowError;
276
277    fn try_from((value, typ): (&dyn Array, LineStringType)) -> GeoArrowResult<Self> {
278        match value.data_type() {
279            DataType::List(_) => (value.as_list::<i32>(), typ).try_into(),
280            DataType::LargeList(_) => (value.as_list::<i64>(), typ).try_into(),
281            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
282                "Unexpected LineString DataType: {dt:?}",
283            ))),
284        }
285    }
286}
287
288impl TryFrom<(&dyn Array, &Field)> for LineStringArray {
289    type Error = GeoArrowError;
290
291    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
292        let typ = field.try_extension_type::<LineStringType>()?;
293        (arr, typ).try_into()
294    }
295}
296
297impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, LineStringType)> for LineStringArray {
298    type Error = GeoArrowError;
299
300    fn try_from(value: (GenericWkbArray<O>, LineStringType)) -> GeoArrowResult<Self> {
301        let mut_arr: LineStringBuilder = value.try_into()?;
302        Ok(mut_arr.finish())
303    }
304}
305
306impl PartialEq for LineStringArray {
307    fn eq(&self, other: &Self) -> bool {
308        self.nulls == other.nulls
309            && offset_buffer_eq(&self.geom_offsets, &other.geom_offsets)
310            && self.coords == other.coords
311    }
312}
313
314impl GeometryTypeId for LineStringArray {
315    const GEOMETRY_TYPE_OFFSET: i8 = 2;
316
317    fn dimension(&self) -> Dimension {
318        self.data_type.dimension()
319    }
320}
321
322#[cfg(test)]
323mod test {
324    use geo_traits::to_geo::ToGeoLineString;
325    use geoarrow_schema::{CoordType, Dimension};
326
327    use super::*;
328    use crate::test::linestring;
329
330    #[test]
331    fn geo_round_trip() {
332        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
333            let geoms = [Some(linestring::ls0()), None, Some(linestring::ls1()), None];
334            let typ =
335                LineStringType::new(Dimension::XY, Default::default()).with_coord_type(coord_type);
336            let geo_arr = LineStringBuilder::from_nullable_line_strings(&geoms, typ).finish();
337
338            for (i, g) in geo_arr.iter().enumerate() {
339                assert_eq!(geoms[i], g.transpose().unwrap().map(|g| g.to_line_string()));
340            }
341
342            // Test sliced
343            for (i, g) in geo_arr.slice(2, 2).iter().enumerate() {
344                assert_eq!(
345                    geoms[i + 2],
346                    g.transpose().unwrap().map(|g| g.to_line_string())
347                );
348            }
349        }
350    }
351
352    #[test]
353    fn geo_round_trip2() {
354        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
355            let geo_arr = linestring::array(coord_type, Dimension::XY);
356            let geo_geoms = geo_arr
357                .iter()
358                .map(|x| x.transpose().unwrap().map(|g| g.to_line_string()))
359                .collect::<Vec<_>>();
360
361            let typ =
362                LineStringType::new(Dimension::XY, Default::default()).with_coord_type(coord_type);
363            let geo_arr2 = LineStringBuilder::from_nullable_line_strings(&geo_geoms, typ).finish();
364            assert_eq!(geo_arr, geo_arr2);
365        }
366    }
367
368    #[test]
369    fn try_from_arrow() {
370        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
371            for dim in [
372                Dimension::XY,
373                Dimension::XYZ,
374                Dimension::XYM,
375                Dimension::XYZM,
376            ] {
377                let geo_arr = linestring::array(coord_type, dim);
378
379                let extension_type = geo_arr.extension_type().clone();
380                let field = extension_type.to_field("geometry", true);
381
382                let arrow_arr = geo_arr.to_array_ref();
383
384                let geo_arr2: LineStringArray =
385                    (arrow_arr.as_ref(), extension_type).try_into().unwrap();
386                let geo_arr3: LineStringArray = (arrow_arr.as_ref(), &field).try_into().unwrap();
387
388                assert_eq!(geo_arr, geo_arr2);
389                assert_eq!(geo_arr, geo_arr3);
390            }
391        }
392    }
393
394    #[test]
395    fn partial_eq() {
396        let arr1 = linestring::ls_array(CoordType::Interleaved);
397        let arr2 = linestring::ls_array(CoordType::Separated);
398        assert_eq!(arr1, arr1);
399        assert_eq!(arr2, arr2);
400        assert_eq!(arr1, arr2);
401
402        assert_ne!(arr1, arr2.slice(0, 2));
403    }
404}
geoarrow_array/array/linestring.rs

geoarrow_array/array/
linestring.rs
