geoarrow_array/array/
multipoint.rs

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

geoarrow_array/array/
multipoint.rs
