geoarrow_array/array/coord/
combined.rs

1use std::sync::Arc;
2
3use arrow_array::{Array, ArrayRef, FixedSizeListArray, StructArray};
4use arrow_schema::DataType;
5use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
6use geoarrow_schema::{CoordType, Dimension};
7
8use crate::array::{InterleavedCoordBuffer, SeparatedCoordBuffer};
9use crate::builder::{InterleavedCoordBufferBuilder, SeparatedCoordBufferBuilder};
10use crate::scalar::Coord;
11
12/// An Arrow representation of an array of coordinates.
13///
14/// As defined in the GeoArrow spec, coordinates can either be interleaved (i.e. a single array of
15/// XYXYXY) or separated (i.e. two arrays, one XXX and another YYY).
16///
17/// This CoordBuffer abstracts over an `InterleavedCoordBuffer` and a `SeparatedCoordBuffer`.
18///
19/// For now all coordinate buffers support only two dimensions.
20///
21/// This is named `CoordBuffer` instead of `CoordArray` because the buffer does not store its own
22/// validity bitmask. Rather the geometry arrays that build on top of this maintain their own
23/// validity masks.
24#[derive(Debug, Clone)]
25pub enum CoordBuffer {
26    /// Interleaved coordinates
27    Interleaved(InterleavedCoordBuffer),
28    /// Separated coordinates
29    Separated(SeparatedCoordBuffer),
30}
31
32impl CoordBuffer {
33    /// Slice this buffer
34    pub(crate) fn slice(&self, offset: usize, length: usize) -> Self {
35        match self {
36            CoordBuffer::Interleaved(c) => CoordBuffer::Interleaved(c.slice(offset, length)),
37            CoordBuffer::Separated(c) => CoordBuffer::Separated(c.slice(offset, length)),
38        }
39    }
40
41    /// The underlying coordinate type
42    pub fn coord_type(&self) -> CoordType {
43        match self {
44            CoordBuffer::Interleaved(_) => CoordType::Interleaved,
45            CoordBuffer::Separated(_) => CoordType::Separated,
46        }
47    }
48
49    /// The arrow [DataType] for this coordinate buffer.
50    pub(crate) fn storage_type(&self) -> DataType {
51        match self {
52            CoordBuffer::Interleaved(c) => c.storage_type(),
53            CoordBuffer::Separated(c) => c.storage_type(),
54        }
55    }
56
57    /// The length of this coordinate buffer
58    pub fn len(&self) -> usize {
59        match self {
60            CoordBuffer::Interleaved(c) => c.len(),
61            CoordBuffer::Separated(c) => c.len(),
62        }
63    }
64
65    /// Whether this coordinate buffer is empty
66    pub fn is_empty(&self) -> bool {
67        self.len() == 0
68    }
69
70    pub(crate) fn value(&self, index: usize) -> Coord<'_> {
71        match self {
72            CoordBuffer::Interleaved(c) => Coord::Interleaved(c.value(index)),
73            CoordBuffer::Separated(c) => Coord::Separated(c.value(index)),
74        }
75    }
76
77    pub(crate) fn into_array_ref(self) -> ArrayRef {
78        self.into()
79    }
80
81    /// The dimension of this coordinate buffer
82    pub fn dim(&self) -> Dimension {
83        match self {
84            CoordBuffer::Interleaved(c) => c.dim(),
85            CoordBuffer::Separated(c) => c.dim(),
86        }
87    }
88
89    /// Convert this coordinate array into the given [CoordType]
90    ///
91    /// This is a no-op if the coord_type matches the existing coord type. Otherwise a full clone
92    /// of the underlying coordinate buffers will be performed.
93    pub fn into_coord_type(self, coord_type: CoordType) -> Self {
94        let dim = self.dim();
95        match (self, coord_type) {
96            (CoordBuffer::Interleaved(cb), CoordType::Interleaved) => CoordBuffer::Interleaved(cb),
97            (CoordBuffer::Interleaved(cb), CoordType::Separated) => {
98                let mut new_buffer = SeparatedCoordBufferBuilder::with_capacity(cb.len(), dim);
99                for i in 0..cb.len() {
100                    let coord = cb.value(i);
101                    new_buffer.push_coord(&coord);
102                }
103                CoordBuffer::Separated(new_buffer.finish())
104            }
105            (CoordBuffer::Separated(cb), CoordType::Separated) => CoordBuffer::Separated(cb),
106            (CoordBuffer::Separated(cb), CoordType::Interleaved) => {
107                let mut new_buffer = InterleavedCoordBufferBuilder::with_capacity(cb.len(), dim);
108                for i in 0..cb.len() {
109                    let coord = cb.value(i);
110                    new_buffer.push_coord(&coord);
111                }
112                CoordBuffer::Interleaved(new_buffer.finish())
113            }
114        }
115    }
116
117    pub(crate) fn from_arrow(value: &dyn Array, dim: Dimension) -> GeoArrowResult<Self> {
118        match value.data_type() {
119            DataType::Struct(_) => {
120                let downcasted = value.as_any().downcast_ref::<StructArray>().unwrap();
121                Ok(CoordBuffer::Separated(SeparatedCoordBuffer::from_arrow(
122                    downcasted, dim,
123                )?))
124            }
125            DataType::FixedSizeList(_, _) => {
126                let downcasted = value.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
127                Ok(CoordBuffer::Interleaved(
128                    InterleavedCoordBuffer::from_arrow(downcasted, dim)?,
129                ))
130            }
131            _ => Err(GeoArrowError::InvalidGeoArrow(format!(
132                "Unexpected coord buffer type: {:?}",
133                value.data_type()
134            ))),
135        }
136    }
137}
138
139impl From<CoordBuffer> for ArrayRef {
140    fn from(value: CoordBuffer) -> Self {
141        match value {
142            CoordBuffer::Interleaved(c) => Arc::new(FixedSizeListArray::from(c)),
143            CoordBuffer::Separated(c) => Arc::new(StructArray::from(c)),
144        }
145    }
146}
147
148impl PartialEq for CoordBuffer {
149    fn eq(&self, other: &Self) -> bool {
150        match (self, other) {
151            (CoordBuffer::Interleaved(a), CoordBuffer::Interleaved(b)) => PartialEq::eq(a, b),
152            (CoordBuffer::Interleaved(left), CoordBuffer::Separated(right)) => {
153                if left.len() != right.len() {
154                    return false;
155                }
156
157                for i in 0..left.len() {
158                    let left_coord = left.value(i);
159                    let right_coord = right.value(i);
160
161                    if left_coord != right_coord {
162                        return false;
163                    }
164                }
165
166                true
167            }
168            (CoordBuffer::Separated(a), CoordBuffer::Separated(b)) => PartialEq::eq(a, b),
169            (CoordBuffer::Separated(left), CoordBuffer::Interleaved(right)) => {
170                if left.len() != right.len() {
171                    return false;
172                }
173
174                for i in 0..left.len() {
175                    let left_coord = left.value(i);
176                    let right_coord = right.value(i);
177
178                    if left_coord != right_coord {
179                        return false;
180                    }
181                }
182
183                true
184            }
185        }
186    }
187}
188
189impl From<InterleavedCoordBuffer> for CoordBuffer {
190    fn from(value: InterleavedCoordBuffer) -> Self {
191        Self::Interleaved(value)
192    }
193}
194
195impl From<SeparatedCoordBuffer> for CoordBuffer {
196    fn from(value: SeparatedCoordBuffer) -> Self {
197        Self::Separated(value)
198    }
199}
200
201// #[cfg(test)]
202// mod test {
203//     use crate::error::Result;
204
205//     use super::*;
206
207//     #[test]
208//     fn test_eq_both_interleaved() -> Result<()> {
209//         let coords1 = vec![0., 3., 1., 4., 2., 5.];
210//         let buf1 =
211//             CoordBuffer::Interleaved(InterleavedCoordBuffer::from_vec(coords1, Dimension::XY)?);
212
213//         let coords2 = vec![0., 3., 1., 4., 2., 5.];
214//         let buf2 =
215//             CoordBuffer::Interleaved(InterleavedCoordBuffer::from_vec(coords2, Dimension::XY)?);
216
217//         assert_eq!(buf1, buf2);
218//         Ok(())
219//     }
220
221//     #[test]
222//     fn test_eq_across_types() -> Result<()> {
223//         let x1 = vec![0., 1., 2.];
224//         let y1 = vec![3., 4., 5.];
225
226//         let buf1 = CoordBuffer::Separated(SeparatedCoordBuffer::new(
227//             [x1.into(), y1.into(), vec![].into(), vec![].into()],
228//             Dimension::XY,
229//         ));
230
231//         let coords2 = vec![0., 3., 1., 4., 2., 5.];
232//         let buf2 =
233//             CoordBuffer::Interleaved(InterleavedCoordBuffer::new(coords2.into(), Dimension::XY));
234
235//         assert_eq!(buf1, buf2);
236//         Ok(())
237//     }
238
239//     #[test]
240//     fn test_eq_across_types_slicing() -> Result<()> {
241//         let x1 = vec![0., 1., 2.];
242//         let y1 = vec![3., 4., 5.];
243
244//         let buf1 = CoordBuffer::Separated((x1, y1).try_into()?).slice(1, 1);
245
246//         let coords2 = vec![0., 3., 1., 4., 2., 5.];
247//         let buf2 =
248//             CoordBuffer::Interleaved(InterleavedCoordBuffer::new(coords2.into(), Dimension::XY))
249//                 .slice(1, 1);
250
251//         assert_eq!(buf1, buf2);
252//         Ok(())
253//     }
254// }
geoarrow_array/array/coord/combined.rs

geoarrow_array/array/coord/
combined.rs
