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(self.data_type.dimension())
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: {dt:?}",
285            ))),
286        }
287    }
288}
289
290impl TryFrom<(&dyn Array, &Field)> for MultiPointArray {
291    type Error = GeoArrowError;
292
293    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
294        let typ = field.try_extension_type::<MultiPointType>()?;
295        (arr, typ).try_into()
296    }
297}
298
299impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, MultiPointType)> for MultiPointArray {
300    type Error = GeoArrowError;
301
302    fn try_from(value: (GenericWkbArray<O>, MultiPointType)) -> GeoArrowResult<Self> {
303        let mut_arr: MultiPointBuilder = value.try_into()?;
304        Ok(mut_arr.finish())
305    }
306}
307
308impl From<PointArray> for MultiPointArray {
309    fn from(value: PointArray) -> Self {
310        let (coord_type, dimension, metadata) = value.data_type.into_inner();
311        let new_type = MultiPointType::new(dimension, metadata).with_coord_type(coord_type);
312
313        let coords = value.coords;
314        let geom_offsets = OffsetBuffer::from_lengths(vec![1; coords.len()]);
315        let nulls = value.nulls;
316        Self {
317            data_type: new_type,
318            coords,
319            geom_offsets,
320            nulls,
321        }
322    }
323}
324
325impl PartialEq for MultiPointArray {
326    fn eq(&self, other: &Self) -> bool {
327        self.nulls == other.nulls
328            && offset_buffer_eq(&self.geom_offsets, &other.geom_offsets)
329            && self.coords == other.coords
330    }
331}
332
333#[cfg(test)]
334mod test {
335    use geo_traits::to_geo::ToGeoMultiPoint;
336    use geoarrow_schema::{CoordType, Dimension};
337
338    use super::*;
339    use crate::test::multipoint;
340
341    #[test]
342    fn geo_round_trip() {
343        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
344            let geoms = [Some(multipoint::mp0()), None, Some(multipoint::mp1()), None];
345            let typ =
346                MultiPointType::new(Dimension::XY, Default::default()).with_coord_type(coord_type);
347            let geo_arr = MultiPointBuilder::from_nullable_multi_points(&geoms, typ).finish();
348
349            for (i, g) in geo_arr.iter().enumerate() {
350                assert_eq!(geoms[i], g.transpose().unwrap().map(|g| g.to_multi_point()));
351            }
352
353            // Test sliced
354            for (i, g) in geo_arr.slice(2, 2).iter().enumerate() {
355                assert_eq!(
356                    geoms[i + 2],
357                    g.transpose().unwrap().map(|g| g.to_multi_point())
358                );
359            }
360        }
361    }
362
363    #[test]
364    fn geo_round_trip2() {
365        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
366            let geo_arr = multipoint::array(coord_type, Dimension::XY);
367            let geo_geoms = geo_arr
368                .iter()
369                .map(|x| x.transpose().unwrap().map(|g| g.to_multi_point()))
370                .collect::<Vec<_>>();
371
372            let typ =
373                MultiPointType::new(Dimension::XY, Default::default()).with_coord_type(coord_type);
374            let geo_arr2 = MultiPointBuilder::from_nullable_multi_points(&geo_geoms, typ).finish();
375            assert_eq!(geo_arr, geo_arr2);
376        }
377    }
378
379    #[test]
380    fn try_from_arrow() {
381        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
382            for dim in [
383                Dimension::XY,
384                Dimension::XYZ,
385                Dimension::XYM,
386                Dimension::XYZM,
387            ] {
388                let geo_arr = multipoint::array(coord_type, dim);
389
390                let extension_type = geo_arr.extension_type().clone();
391                let field = extension_type.to_field("geometry", true);
392
393                let arrow_arr = geo_arr.to_array_ref();
394
395                let geo_arr2: MultiPointArray =
396                    (arrow_arr.as_ref(), extension_type).try_into().unwrap();
397                let geo_arr3: MultiPointArray = (arrow_arr.as_ref(), &field).try_into().unwrap();
398
399                assert_eq!(geo_arr, geo_arr2);
400                assert_eq!(geo_arr, geo_arr3);
401            }
402        }
403    }
404
405    #[test]
406    fn partial_eq() {
407        for dim in [
408            Dimension::XY,
409            Dimension::XYZ,
410            Dimension::XYM,
411            Dimension::XYZM,
412        ] {
413            let arr1 = multipoint::array(CoordType::Interleaved, dim);
414            let arr2 = multipoint::array(CoordType::Separated, dim);
415            assert_eq!(arr1, arr1);
416            assert_eq!(arr2, arr2);
417            assert_eq!(arr1, arr2);
418
419            assert_ne!(arr1, arr2.slice(0, 2));
420        }
421    }
422}
geoarrow_array/array/multipoint.rs

geoarrow_array/array/
multipoint.rs
