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    /// Returns the element at index `i`, not considering validity.
118    ///
119    /// # Examples
120    ///
121    /// ```
122    /// use geo_traits::CoordTrait;
123    /// use geoarrow_array::array::InterleavedCoordBuffer;
124    /// use geoarrow_schema::Dimension;
125    ///
126    /// let coords = [
127    ///     geo_types::coord! { x: 1.0, y: 2.0 },
128    ///     geo_types::coord! { x: 3.0, y: 4.0 },
129    /// ];
130    /// let coord_buffer = InterleavedCoordBuffer::from_coords(coords.iter(), Dimension::XY).unwrap();
131    /// let coord = coord_buffer.value(0);
132    /// assert_eq!(coord.x(), 1.0);
133    /// assert_eq!(coord.y(), 2.0);
134    /// ```
135    ///
136    /// # Panics
137    ///
138    /// Panics if the value is outside the bounds of the buffer.
139    pub fn value(&self, index: usize) -> InterleavedCoord<'_> {
140        assert!(index <= self.len());
141        unsafe { self.value_unchecked(index) }
142    }
143
144    /// Returns the element at index `i`, not considering validity.
145    ///
146    /// # Examples
147    ///
148    /// ```
149    /// use geo_traits::CoordTrait;
150    /// use geoarrow_array::array::InterleavedCoordBuffer;
151    /// use geoarrow_schema::Dimension;
152    ///
153    /// let coords = [
154    ///     geo_types::coord! { x: 1.0, y: 2.0 },
155    ///     geo_types::coord! { x: 3.0, y: 4.0 },
156    /// ];
157    /// let coord_buffer = InterleavedCoordBuffer::from_coords(coords.iter(), Dimension::XY).unwrap();
158    /// let coord = unsafe { coord_buffer.value_unchecked(0) };
159    /// assert_eq!(coord.x(), 1.0);
160    /// assert_eq!(coord.y(), 2.0);
161    /// ```
162    ///
163    /// # Safety
164    ///
165    /// Caller is responsible for ensuring that the index is within the bounds of the buffer.
166    pub unsafe fn value_unchecked(&self, index: usize) -> InterleavedCoord<'_> {
167        InterleavedCoord {
168            coords: &self.coords,
169            i: index,
170            dim: self.dim,
171        }
172    }
173
174    pub(crate) fn from_arrow(array: &FixedSizeListArray, dim: Dimension) -> GeoArrowResult<Self> {
175        if array.value_length() != dim.size() as i32 {
176            return Err(GeoArrowError::InvalidGeoArrow(format!(
177                "Expected the FixedSizeListArray to match the dimension. Array length is {}, dimension is: {:?} have size 2",
178                array.value_length(),
179                dim
180            )));
181        }
182
183        let coord_array_values = array
184            .values()
185            .as_any()
186            .downcast_ref::<Float64Array>()
187            .unwrap();
188
189        Ok(InterleavedCoordBuffer::new(
190            coord_array_values.values().clone(),
191            dim,
192        ))
193    }
194}
195
196impl From<InterleavedCoordBuffer> for FixedSizeListArray {
197    fn from(value: InterleavedCoordBuffer) -> Self {
198        FixedSizeListArray::new(
199            Arc::new(value.values_field()),
200            value.dim.size() as i32,
201            Arc::new(value.values_array()),
202            None,
203        )
204    }
205}
206
207#[cfg(test)]
208mod test {
209    use super::*;
210
211    #[test]
212    fn test_eq_slicing() {
213        let coords1 = vec![0., 3., 1., 4., 2., 5.];
214        let buf1 = InterleavedCoordBuffer::new(coords1.into(), Dimension::XY).slice(1, 1);
215
216        let coords2 = vec![1., 4.];
217        let buf2 = InterleavedCoordBuffer::new(coords2.into(), Dimension::XY);
218
219        assert_eq!(buf1, buf2);
220    }
221}
geoarrow_array/array/coord/interleaved.rs

geoarrow_array/array/coord/
interleaved.rs
