geoarrow_array/array/
point.rs

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

geoarrow_array/array/
point.rs
