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

geoarrow_array/array/
linestring.rs
