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

geoarrow_array/array/
point.rs
