geoarrow_array/array/coord/
interleaved.rs

1use std::sync::Arc;
2
3use arrow_array::{Array, FixedSizeListArray, Float64Array};
4use arrow_buffer::ScalarBuffer;
5use arrow_schema::{DataType, Field};
6use geo_traits::CoordTrait;
7use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8use geoarrow_schema::{CoordType, Dimension, PointType};
9
10use crate::builder::InterleavedCoordBufferBuilder;
11use crate::scalar::InterleavedCoord;
12
13/// An array of coordinates stored interleaved in a single buffer.
14///
15/// This stores all coordinates in interleaved fashion in a single underlying buffer: e.g. `xyxyxy`
16/// for 2D coordinates.
17#[derive(Debug, Clone, PartialEq)]
18pub struct InterleavedCoordBuffer {
19    pub(crate) coords: ScalarBuffer<f64>,
20    pub(crate) dim: Dimension,
21}
22
23fn check(coords: &ScalarBuffer<f64>, dim: Dimension) -> GeoArrowResult<()> {
24    if coords.len() % dim.size() != 0 {
25        return Err(GeoArrowError::InvalidGeoArrow(
26            "Length of interleaved coordinate buffer must be a multiple of the dimension size"
27                .to_string(),
28        ));
29    }
30
31    Ok(())
32}
33
34impl InterleavedCoordBuffer {
35    /// The underlying coordinate type
36    pub const COORD_TYPE: CoordType = CoordType::Interleaved;
37
38    /// Construct a new InterleavedCoordBuffer
39    ///
40    /// # Panics
41    ///
42    /// - if coords.len() % dim.size() != 0
43    pub fn new(coords: ScalarBuffer<f64>, dim: Dimension) -> Self {
44        Self::try_new(coords, dim).unwrap()
45    }
46
47    /// Construct a new InterleavedCoordBuffer
48    ///
49    /// # Errors
50    ///
51    /// - if the coordinate buffer have different lengths
52    pub fn try_new(coords: ScalarBuffer<f64>, dim: Dimension) -> GeoArrowResult<Self> {
53        check(&coords, dim)?;
54        Ok(Self { coords, dim })
55    }
56
57    /// Construct from an iterator of coordinates.
58    pub fn from_coords<'a>(
59        coords: impl ExactSizeIterator<Item = &'a (impl CoordTrait<T = f64> + 'a)>,
60        dim: Dimension,
61    ) -> GeoArrowResult<Self> {
62        Ok(InterleavedCoordBufferBuilder::from_coords(coords, dim)?.finish())
63    }
64
65    /// Access the underlying coordinate buffer.
66    pub fn coords(&self) -> &ScalarBuffer<f64> {
67        &self.coords
68    }
69
70    pub(crate) fn values_array(&self) -> Float64Array {
71        Float64Array::new(self.coords.clone(), None)
72    }
73
74    /// The dimension of this coordinate buffer
75    pub fn dim(&self) -> Dimension {
76        self.dim
77    }
78
79    pub(crate) fn values_field(&self) -> Field {
80        match self.dim {
81            Dimension::XY => Field::new("xy", DataType::Float64, false),
82            Dimension::XYZ => Field::new("xyz", DataType::Float64, false),
83            Dimension::XYM => Field::new("xym", DataType::Float64, false),
84            Dimension::XYZM => Field::new("xyzm", DataType::Float64, false),
85        }
86    }
87
88    pub(crate) fn slice(&self, offset: usize, length: usize) -> Self {
89        assert!(
90            offset + length <= self.len(),
91            "offset + length may not exceed length of array"
92        );
93        Self {
94            coords: self
95                .coords
96                .slice(offset * self.dim.size(), length * self.dim.size()),
97            dim: self.dim,
98        }
99    }
100
101    pub(crate) fn storage_type(&self) -> DataType {
102        PointType::new(self.dim, Default::default())
103            .with_coord_type(Self::COORD_TYPE)
104            .data_type()
105    }
106
107    /// The number of coordinates
108    pub fn len(&self) -> usize {
109        self.coords.len() / self.dim.size()
110    }
111
112    /// Whether this buffer is empty
113    pub fn is_empty(&self) -> bool {
114        self.len() == 0
115    }
116
117    pub(crate) fn value(&self, index: usize) -> InterleavedCoord<'_> {
118        assert!(index <= self.len());
119        self.value_unchecked(index)
120    }
121
122    pub(crate) fn value_unchecked(&self, index: usize) -> InterleavedCoord<'_> {
123        InterleavedCoord {
124            coords: &self.coords,
125            i: index,
126            dim: self.dim,
127        }
128    }
129
130    pub(crate) fn from_arrow(array: &FixedSizeListArray, dim: Dimension) -> GeoArrowResult<Self> {
131        if array.value_length() != dim.size() as i32 {
132            return Err(GeoArrowError::InvalidGeoArrow(format!(
133                "Expected the FixedSizeListArray to match the dimension. Array length is {}, dimension is: {:?} have size 2",
134                array.value_length(),
135                dim
136            )));
137        }
138
139        let coord_array_values = array
140            .values()
141            .as_any()
142            .downcast_ref::<Float64Array>()
143            .unwrap();
144
145        Ok(InterleavedCoordBuffer::new(
146            coord_array_values.values().clone(),
147            dim,
148        ))
149    }
150}
151
152impl From<InterleavedCoordBuffer> for FixedSizeListArray {
153    fn from(value: InterleavedCoordBuffer) -> Self {
154        FixedSizeListArray::new(
155            Arc::new(value.values_field()),
156            value.dim.size() as i32,
157            Arc::new(value.values_array()),
158            None,
159        )
160    }
161}
162
163#[cfg(test)]
164mod test {
165    use super::*;
166
167    #[test]
168    fn test_eq_slicing() {
169        let coords1 = vec![0., 3., 1., 4., 2., 5.];
170        let buf1 = InterleavedCoordBuffer::new(coords1.into(), Dimension::XY).slice(1, 1);
171
172        let coords2 = vec![1., 4.];
173        let buf2 = InterleavedCoordBuffer::new(coords2.into(), Dimension::XY);
174
175        assert_eq!(buf1, buf2);
176    }
177}
geoarrow_array/array/coord/interleaved.rs

geoarrow_array/array/coord/
interleaved.rs
