geoarrow_array/array/
multipolygon.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, MultiPolygonType};
9
10use crate::array::{CoordBuffer, GenericWkbArray, PolygonArray};
11use crate::builder::MultiPolygonBuilder;
12use crate::capacity::MultiPolygonCapacity;
13use crate::eq::offset_buffer_eq;
14use crate::scalar::MultiPolygon;
15use crate::trait_::{GeoArrowArray, GeoArrowArrayAccessor, IntoArrow};
16use crate::util::{OffsetBufferUtils, offsets_buffer_i64_to_i32};
17
18/// An immutable array of MultiPolygon geometries.
19///
20/// This is semantically equivalent to `Vec<Option<MultiPolygon>>` due to the internal validity
21/// bitmap.
22#[derive(Debug, Clone)]
23pub struct MultiPolygonArray {
24    pub(crate) data_type: MultiPolygonType,
25
26    pub(crate) coords: CoordBuffer,
27
28    /// Offsets into the polygon array where each geometry starts
29    pub(crate) geom_offsets: OffsetBuffer<i32>,
30
31    /// Offsets into the ring array where each polygon starts
32    pub(crate) polygon_offsets: OffsetBuffer<i32>,
33
34    /// Offsets into the coordinate array where each ring starts
35    pub(crate) ring_offsets: OffsetBuffer<i32>,
36
37    /// Validity bitmap
38    pub(crate) nulls: Option<NullBuffer>,
39}
40
41pub(super) fn check(
42    coords: &CoordBuffer,
43    geom_offsets: &OffsetBuffer<i32>,
44    polygon_offsets: &OffsetBuffer<i32>,
45    ring_offsets: &OffsetBuffer<i32>,
46    validity_len: Option<usize>,
47) -> GeoArrowResult<()> {
48    if validity_len.is_some_and(|len| len != geom_offsets.len_proxy()) {
49        return Err(GeoArrowError::InvalidGeoArrow(
50            "nulls mask length must match the number of values".to_string(),
51        ));
52    }
53    if *ring_offsets.last() as usize != coords.len() {
54        return Err(GeoArrowError::InvalidGeoArrow(
55            "largest ring offset must match coords length".to_string(),
56        ));
57    }
58
59    if *polygon_offsets.last() as usize != ring_offsets.len_proxy() {
60        return Err(GeoArrowError::InvalidGeoArrow(
61            "largest polygon offset must match ring offsets length".to_string(),
62        ));
63    }
64
65    if *geom_offsets.last() as usize != polygon_offsets.len_proxy() {
66        return Err(GeoArrowError::InvalidGeoArrow(
67            "largest geometry offset must match polygon offsets length".to_string(),
68        ));
69    }
70
71    Ok(())
72}
73
74impl MultiPolygonArray {
75    /// Create a new MultiPolygonArray from parts
76    ///
77    /// # Implementation
78    ///
79    /// This function is `O(1)`.
80    ///
81    /// # Panics
82    ///
83    /// - if the nulls is not `None` and its length is different from the number of geometries
84    /// - if the largest ring offset does not match the number of coordinates
85    /// - if the largest polygon offset does not match the size of ring offsets
86    /// - if the largest geometry offset does not match the size of polygon offsets
87    pub fn new(
88        coords: CoordBuffer,
89        geom_offsets: OffsetBuffer<i32>,
90        polygon_offsets: OffsetBuffer<i32>,
91        ring_offsets: OffsetBuffer<i32>,
92        nulls: Option<NullBuffer>,
93        metadata: Arc<Metadata>,
94    ) -> Self {
95        Self::try_new(
96            coords,
97            geom_offsets,
98            polygon_offsets,
99            ring_offsets,
100            nulls,
101            metadata,
102        )
103        .unwrap()
104    }
105
106    /// Create a new MultiPolygonArray from parts
107    ///
108    /// # Implementation
109    ///
110    /// This function is `O(1)`.
111    ///
112    /// # Errors
113    ///
114    /// - if the nulls is not `None` and its length is different from the number of geometries
115    /// - if the largest ring offset does not match the number of coordinates
116    /// - if the largest polygon offset does not match the size of ring offsets
117    /// - if the largest geometry offset does not match the size of polygon offsets
118    pub fn try_new(
119        coords: CoordBuffer,
120        geom_offsets: OffsetBuffer<i32>,
121        polygon_offsets: OffsetBuffer<i32>,
122        ring_offsets: OffsetBuffer<i32>,
123        nulls: Option<NullBuffer>,
124        metadata: Arc<Metadata>,
125    ) -> GeoArrowResult<Self> {
126        check(
127            &coords,
128            &geom_offsets,
129            &polygon_offsets,
130            &ring_offsets,
131            nulls.as_ref().map(|v| v.len()),
132        )?;
133        Ok(Self {
134            data_type: MultiPolygonType::new(coords.dim(), metadata)
135                .with_coord_type(coords.coord_type()),
136            coords,
137            geom_offsets,
138            polygon_offsets,
139            ring_offsets,
140            nulls,
141        })
142    }
143
144    fn vertices_field(&self) -> Arc<Field> {
145        Field::new("vertices", self.coords.storage_type(), false).into()
146    }
147
148    fn rings_field(&self) -> Arc<Field> {
149        let name = "rings";
150        Field::new_list(name, self.vertices_field(), false).into()
151    }
152
153    fn polygons_field(&self) -> Arc<Field> {
154        let name = "polygons";
155        Field::new_list(name, self.rings_field(), false).into()
156    }
157
158    /// Access the underlying coordinate buffer
159    pub fn coords(&self) -> &CoordBuffer {
160        &self.coords
161    }
162
163    /// Access the underlying geometry offsets buffer
164    pub fn geom_offsets(&self) -> &OffsetBuffer<i32> {
165        &self.geom_offsets
166    }
167
168    /// Access the underlying polygon offsets buffer
169    pub fn polygon_offsets(&self) -> &OffsetBuffer<i32> {
170        &self.polygon_offsets
171    }
172
173    /// Access the underlying ring offsets buffer
174    pub fn ring_offsets(&self) -> &OffsetBuffer<i32> {
175        &self.ring_offsets
176    }
177
178    /// The lengths of each buffer contained in this array.
179    pub fn buffer_lengths(&self) -> MultiPolygonCapacity {
180        MultiPolygonCapacity::new(
181            *self.ring_offsets.last() as usize,
182            *self.polygon_offsets.last() as usize,
183            *self.geom_offsets.last() as usize,
184            self.len(),
185        )
186    }
187
188    /// The number of bytes occupied by this array.
189    pub fn num_bytes(&self) -> usize {
190        let validity_len = self.nulls.as_ref().map(|v| v.buffer().len()).unwrap_or(0);
191        validity_len + self.buffer_lengths().num_bytes()
192    }
193
194    /// Slice this [`MultiPolygonArray`].
195    ///
196    /// # Panic
197    /// This function panics iff `offset + length > self.len()`.
198    #[inline]
199    pub fn slice(&self, offset: usize, length: usize) -> Self {
200        assert!(
201            offset + length <= self.len(),
202            "offset + length may not exceed length of array"
203        );
204        // Note: we **only** slice the geom_offsets and not any actual data. Otherwise the offsets
205        // would be in the wrong location.
206        Self {
207            data_type: self.data_type.clone(),
208            coords: self.coords.clone(),
209            geom_offsets: self.geom_offsets.slice(offset, length),
210            polygon_offsets: self.polygon_offsets.clone(),
211            ring_offsets: self.ring_offsets.clone(),
212            nulls: self.nulls.as_ref().map(|v| v.slice(offset, length)),
213        }
214    }
215
216    /// Change the [`CoordType`] of this array.
217    pub fn into_coord_type(self, coord_type: CoordType) -> Self {
218        Self {
219            data_type: self.data_type.with_coord_type(coord_type),
220            coords: self.coords.into_coord_type(coord_type),
221            ..self
222        }
223    }
224
225    /// Change the [`Metadata`] of this array.
226    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
227        Self {
228            data_type: self.data_type.with_metadata(metadata),
229            ..self
230        }
231    }
232}
233
234impl GeoArrowArray for MultiPolygonArray {
235    fn as_any(&self) -> &dyn std::any::Any {
236        self
237    }
238
239    fn into_array_ref(self) -> ArrayRef {
240        Arc::new(self.into_arrow())
241    }
242
243    fn to_array_ref(&self) -> ArrayRef {
244        self.clone().into_array_ref()
245    }
246
247    #[inline]
248    fn len(&self) -> usize {
249        self.geom_offsets.len_proxy()
250    }
251
252    #[inline]
253    fn logical_nulls(&self) -> Option<NullBuffer> {
254        self.nulls.clone()
255    }
256
257    #[inline]
258    fn logical_null_count(&self) -> usize {
259        self.nulls.as_ref().map(|v| v.null_count()).unwrap_or(0)
260    }
261
262    #[inline]
263    fn is_null(&self, i: usize) -> bool {
264        self.nulls
265            .as_ref()
266            .map(|n| n.is_null(i))
267            .unwrap_or_default()
268    }
269
270    fn data_type(&self) -> GeoArrowType {
271        GeoArrowType::MultiPolygon(self.data_type.clone())
272    }
273
274    fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
275        Arc::new(self.slice(offset, length))
276    }
277
278    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
279        Arc::new(self.with_metadata(metadata))
280    }
281}
282
283impl<'a> GeoArrowArrayAccessor<'a> for MultiPolygonArray {
284    type Item = MultiPolygon<'a>;
285
286    unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
287        Ok(MultiPolygon::new(
288            &self.coords,
289            &self.geom_offsets,
290            &self.polygon_offsets,
291            &self.ring_offsets,
292            index,
293        ))
294    }
295}
296
297impl IntoArrow for MultiPolygonArray {
298    type ArrowArray = GenericListArray<i32>;
299    type ExtensionType = MultiPolygonType;
300
301    fn into_arrow(self) -> Self::ArrowArray {
302        let vertices_field = self.vertices_field();
303        let rings_field = self.rings_field();
304        let polygons_field = self.polygons_field();
305
306        let nulls = self.nulls;
307        let coord_array = ArrayRef::from(self.coords);
308        let ring_array = Arc::new(GenericListArray::new(
309            vertices_field,
310            self.ring_offsets,
311            coord_array,
312            None,
313        ));
314        let polygons_array = Arc::new(GenericListArray::new(
315            rings_field,
316            self.polygon_offsets,
317            ring_array,
318            None,
319        ));
320        GenericListArray::new(polygons_field, self.geom_offsets, polygons_array, nulls)
321    }
322
323    fn extension_type(&self) -> &Self::ExtensionType {
324        &self.data_type
325    }
326}
327
328impl TryFrom<(&GenericListArray<i32>, MultiPolygonType)> for MultiPolygonArray {
329    type Error = GeoArrowError;
330
331    fn try_from(
332        (geom_array, typ): (&GenericListArray<i32>, MultiPolygonType),
333    ) -> GeoArrowResult<Self> {
334        let geom_offsets = geom_array.offsets();
335        let nulls = geom_array.nulls();
336
337        let polygons_dyn_array = geom_array.values();
338        let polygons_array = polygons_dyn_array.as_list::<i32>();
339
340        let polygon_offsets = polygons_array.offsets();
341        let rings_dyn_array = polygons_array.values();
342        let rings_array = rings_dyn_array.as_list::<i32>();
343
344        let ring_offsets = rings_array.offsets();
345        let coords = CoordBuffer::from_arrow(rings_array.values().as_ref(), typ.dimension())?;
346
347        Ok(Self::new(
348            coords,
349            geom_offsets.clone(),
350            polygon_offsets.clone(),
351            ring_offsets.clone(),
352            nulls.cloned(),
353            typ.metadata().clone(),
354        ))
355    }
356}
357
358impl TryFrom<(&GenericListArray<i64>, MultiPolygonType)> for MultiPolygonArray {
359    type Error = GeoArrowError;
360
361    fn try_from(
362        (geom_array, typ): (&GenericListArray<i64>, MultiPolygonType),
363    ) -> GeoArrowResult<Self> {
364        let geom_offsets = offsets_buffer_i64_to_i32(geom_array.offsets())?;
365        let nulls = geom_array.nulls();
366
367        let polygons_dyn_array = geom_array.values();
368        let polygons_array = polygons_dyn_array.as_list::<i64>();
369
370        let polygon_offsets = offsets_buffer_i64_to_i32(polygons_array.offsets())?;
371        let rings_dyn_array = polygons_array.values();
372        let rings_array = rings_dyn_array.as_list::<i64>();
373
374        let ring_offsets = offsets_buffer_i64_to_i32(rings_array.offsets())?;
375        let coords = CoordBuffer::from_arrow(rings_array.values().as_ref(), typ.dimension())?;
376
377        Ok(Self::new(
378            coords,
379            geom_offsets,
380            polygon_offsets,
381            ring_offsets,
382            nulls.cloned(),
383            typ.metadata().clone(),
384        ))
385    }
386}
387
388impl TryFrom<(&dyn Array, MultiPolygonType)> for MultiPolygonArray {
389    type Error = GeoArrowError;
390
391    fn try_from((value, typ): (&dyn Array, MultiPolygonType)) -> GeoArrowResult<Self> {
392        match value.data_type() {
393            DataType::List(_) => (value.as_list::<i32>(), typ).try_into(),
394            DataType::LargeList(_) => (value.as_list::<i64>(), typ).try_into(),
395            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
396                "Unexpected MultiPolygon DataType: {:?}",
397                dt
398            ))),
399        }
400    }
401}
402
403impl TryFrom<(&dyn Array, &Field)> for MultiPolygonArray {
404    type Error = GeoArrowError;
405
406    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
407        let typ = field.try_extension_type::<MultiPolygonType>()?;
408        (arr, typ).try_into()
409    }
410}
411
412impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, MultiPolygonType)> for MultiPolygonArray {
413    type Error = GeoArrowError;
414
415    fn try_from(value: (GenericWkbArray<O>, MultiPolygonType)) -> GeoArrowResult<Self> {
416        let mut_arr: MultiPolygonBuilder = value.try_into()?;
417        Ok(mut_arr.finish())
418    }
419}
420
421impl From<PolygonArray> for MultiPolygonArray {
422    fn from(value: PolygonArray) -> Self {
423        let (coord_type, dimension, metadata) = value.data_type.into_inner();
424        let new_type = MultiPolygonType::new(dimension, metadata).with_coord_type(coord_type);
425
426        let coords = value.coords;
427        let geom_offsets = OffsetBuffer::from_lengths(vec![1; coords.len()]);
428        let ring_offsets = value.ring_offsets;
429        let polygon_offsets = value.geom_offsets;
430        let nulls = value.nulls;
431        Self {
432            data_type: new_type,
433            coords,
434            geom_offsets,
435            polygon_offsets,
436            ring_offsets,
437            nulls,
438        }
439    }
440}
441
442impl PartialEq for MultiPolygonArray {
443    fn eq(&self, other: &Self) -> bool {
444        self.nulls == other.nulls
445            && offset_buffer_eq(&self.geom_offsets, &other.geom_offsets)
446            && offset_buffer_eq(&self.polygon_offsets, &other.polygon_offsets)
447            && offset_buffer_eq(&self.ring_offsets, &other.ring_offsets)
448            && self.coords == other.coords
449    }
450}
451
452#[cfg(test)]
453mod test {
454    use geo_traits::to_geo::ToGeoMultiPolygon;
455    use geoarrow_schema::{CoordType, Dimension};
456
457    use super::*;
458    use crate::test::multipolygon;
459
460    #[test]
461    fn geo_round_trip() {
462        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
463            let geoms = [
464                Some(multipolygon::mp0()),
465                None,
466                Some(multipolygon::mp1()),
467                None,
468            ];
469            let typ = MultiPolygonType::new(Dimension::XY, Default::default())
470                .with_coord_type(coord_type);
471            let geo_arr = MultiPolygonBuilder::from_nullable_multi_polygons(&geoms, typ).finish();
472
473            for (i, g) in geo_arr.iter().enumerate() {
474                assert_eq!(
475                    geoms[i],
476                    g.transpose().unwrap().map(|g| g.to_multi_polygon())
477                );
478            }
479
480            // Test sliced
481            for (i, g) in geo_arr.slice(2, 2).iter().enumerate() {
482                assert_eq!(
483                    geoms[i + 2],
484                    g.transpose().unwrap().map(|g| g.to_multi_polygon())
485                );
486            }
487        }
488    }
489
490    #[test]
491    fn geo_round_trip2() {
492        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
493            let geo_arr = multipolygon::array(coord_type, Dimension::XY);
494            let geo_geoms = geo_arr
495                .iter()
496                .map(|x| x.transpose().unwrap().map(|g| g.to_multi_polygon()))
497                .collect::<Vec<_>>();
498
499            let typ = MultiPolygonType::new(Dimension::XY, Default::default())
500                .with_coord_type(coord_type);
501            let geo_arr2 =
502                MultiPolygonBuilder::from_nullable_multi_polygons(&geo_geoms, typ).finish();
503            assert_eq!(geo_arr, geo_arr2);
504        }
505    }
506
507    #[test]
508    fn try_from_arrow() {
509        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
510            for dim in [
511                Dimension::XY,
512                Dimension::XYZ,
513                Dimension::XYM,
514                Dimension::XYZM,
515            ] {
516                let geo_arr = multipolygon::array(coord_type, dim);
517
518                let extension_type = geo_arr.extension_type().clone();
519                let field = extension_type.to_field("geometry", true);
520
521                let arrow_arr = geo_arr.to_array_ref();
522
523                let geo_arr2: MultiPolygonArray =
524                    (arrow_arr.as_ref(), extension_type).try_into().unwrap();
525                let geo_arr3: MultiPolygonArray = (arrow_arr.as_ref(), &field).try_into().unwrap();
526
527                assert_eq!(geo_arr, geo_arr2);
528                assert_eq!(geo_arr, geo_arr3);
529            }
530        }
531    }
532
533    #[test]
534    fn partial_eq() {
535        for dim in [
536            Dimension::XY,
537            Dimension::XYZ,
538            Dimension::XYM,
539            Dimension::XYZM,
540        ] {
541            let arr1 = multipolygon::array(CoordType::Interleaved, dim);
542            let arr2 = multipolygon::array(CoordType::Separated, dim);
543            assert_eq!(arr1, arr1);
544            assert_eq!(arr2, arr2);
545            assert_eq!(arr1, arr2);
546
547            assert_ne!(arr1, arr2.slice(0, 2));
548        }
549    }
550}
geoarrow_array/array/multipolygon.rs

geoarrow_array/array/
multipolygon.rs
