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: {:?}",
250                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 the coords are already true, don't check for NaNs
274        // TODO: maybe only iterate once for perf?
275        if self.coords == other.coords {
276            return true;
277        }
278
279        if self.coords.len() != other.coords.len() {
280            return false;
281        }
282
283        // TODO: this should check for point equal.
284        for point_idx in 0..self.len() {
285            let c1 = self.value(point_idx).unwrap();
286            let c2 = other.value(point_idx).unwrap();
287            if !point_eq(&c1, &c2) {
288                return false;
289            }
290        }
291
292        true
293    }
294}
295
296#[cfg(test)]
297mod test {
298    use geo_traits::to_geo::ToGeoPoint;
299    use geoarrow_schema::{CoordType, Dimension};
300
301    use super::*;
302    use crate::builder::PointBuilder;
303    use crate::test::point;
304
305    #[test]
306    fn geo_round_trip() {
307        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
308            let geoms = [
309                Some(point::p0()),
310                Some(point::p1()),
311                None,
312                Some(point::p2()),
313            ];
314            let typ = PointType::new(Dimension::XY, Default::default()).with_coord_type(coord_type);
315            let geo_arr =
316                PointBuilder::from_nullable_points(geoms.iter().map(|x| x.as_ref()), typ).finish();
317
318            for (i, g) in geo_arr.iter().enumerate() {
319                assert_eq!(geoms[i], g.transpose().unwrap().map(|g| g.to_point()));
320            }
321
322            // Test sliced
323            for (i, g) in geo_arr.slice(2, 2).iter().enumerate() {
324                assert_eq!(geoms[i + 2], g.transpose().unwrap().map(|g| g.to_point()));
325            }
326        }
327    }
328
329    #[test]
330    fn try_from_arrow() {
331        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
332            for dim in [
333                Dimension::XY,
334                Dimension::XYZ,
335                Dimension::XYM,
336                Dimension::XYZM,
337            ] {
338                let geo_arr = point::array(coord_type, dim);
339
340                let point_type = geo_arr.extension_type().clone();
341                let field = point_type.to_field("geometry", true);
342
343                let arrow_arr = geo_arr.to_array_ref();
344
345                let geo_arr2: PointArray = (arrow_arr.as_ref(), point_type).try_into().unwrap();
346                let geo_arr3: PointArray = (arrow_arr.as_ref(), &field).try_into().unwrap();
347
348                assert_eq!(geo_arr, geo_arr2);
349                assert_eq!(geo_arr, geo_arr3);
350            }
351        }
352    }
353
354    #[test]
355    fn into_coord_type() {
356        for dim in [
357            Dimension::XY,
358            Dimension::XYZ,
359            Dimension::XYM,
360            Dimension::XYZM,
361        ] {
362            let geo_arr = point::array(CoordType::Interleaved, dim);
363            let geo_arr2 = geo_arr
364                .clone()
365                .into_coord_type(CoordType::Separated)
366                .into_coord_type(CoordType::Interleaved);
367
368            assert_eq!(geo_arr, geo_arr2);
369        }
370    }
371
372    #[test]
373    fn partial_eq() {
374        for dim in [
375            Dimension::XY,
376            Dimension::XYZ,
377            Dimension::XYM,
378            Dimension::XYZM,
379        ] {
380            let arr1 = point::array(CoordType::Interleaved, dim);
381            let arr2 = point::array(CoordType::Separated, dim);
382            assert_eq!(arr1, arr1);
383            assert_eq!(arr2, arr2);
384            assert_eq!(arr1, arr2);
385
386            assert_ne!(arr1, arr2.slice(0, 2));
387        }
388    }
389}
geoarrow_array/array/point.rs

geoarrow_array/array/
point.rs
