geoarrow_array/array/
geometrycollection.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, GeometryCollectionType, Metadata};
9
10use crate::array::{GenericWkbArray, MixedGeometryArray};
11use crate::builder::GeometryCollectionBuilder;
12use crate::capacity::GeometryCollectionCapacity;
13use crate::eq::offset_buffer_eq;
14use crate::scalar::GeometryCollection;
15use crate::trait_::{GeoArrowArray, GeoArrowArrayAccessor, IntoArrow};
16use crate::util::{OffsetBufferUtils, offsets_buffer_i64_to_i32};
17
18/// An immutable array of GeometryCollection geometries.
19///
20/// This is semantically equivalent to `Vec<Option<GeometryCollection>>` due to the internal
21/// validity bitmap.
22#[derive(Debug, Clone)]
23pub struct GeometryCollectionArray {
24    pub(crate) data_type: GeometryCollectionType,
25
26    pub(crate) array: MixedGeometryArray,
27
28    /// Offsets into the mixed geometry array where each geometry starts
29    pub(crate) geom_offsets: OffsetBuffer<i32>,
30
31    /// Validity bitmap
32    pub(crate) nulls: Option<NullBuffer>,
33}
34
35impl GeometryCollectionArray {
36    /// Create a new GeometryCollectionArray from parts
37    ///
38    /// # Implementation
39    ///
40    /// This function is `O(1)`.
41    pub fn new(
42        array: MixedGeometryArray,
43        geom_offsets: OffsetBuffer<i32>,
44        nulls: Option<NullBuffer>,
45        metadata: Arc<Metadata>,
46    ) -> Self {
47        Self {
48            data_type: GeometryCollectionType::new(array.dim, metadata)
49                .with_coord_type(array.coord_type),
50            array,
51            geom_offsets,
52            nulls,
53        }
54    }
55
56    fn geometries_field(&self) -> Arc<Field> {
57        Field::new("geometries", self.array.storage_type(), false).into()
58    }
59
60    /// The lengths of each buffer contained in this array.
61    pub fn buffer_lengths(&self) -> GeometryCollectionCapacity {
62        GeometryCollectionCapacity::new(
63            self.array.buffer_lengths(),
64            *self.geom_offsets.last() as usize,
65        )
66    }
67
68    /// The number of bytes occupied by this array.
69    pub fn num_bytes(&self) -> usize {
70        let validity_len = self.nulls.as_ref().map(|v| v.buffer().len()).unwrap_or(0);
71        validity_len + self.buffer_lengths().num_bytes(self.data_type.dimension())
72    }
73
74    /// Slice this [`GeometryCollectionArray`].
75    ///
76    /// # Implementation
77    ///
78    /// This operation is `O(1)` as it amounts to increasing a few ref counts.
79    ///
80    /// # Panic
81    /// This function panics iff `offset + length > self.len()`.
82    #[inline]
83    pub fn slice(&self, offset: usize, length: usize) -> Self {
84        assert!(
85            offset + length <= self.len(),
86            "offset + length may not exceed length of array"
87        );
88        // Note: we **only** slice the geom_offsets and not any actual data
89        Self {
90            data_type: self.data_type.clone(),
91            array: self.array.clone(),
92            geom_offsets: self.geom_offsets.slice(offset, length),
93            nulls: self.nulls.as_ref().map(|v| v.slice(offset, length)),
94        }
95    }
96
97    /// Change the [`CoordType`] of this array.
98    pub fn into_coord_type(self, coord_type: CoordType) -> Self {
99        Self {
100            data_type: self.data_type.with_coord_type(coord_type),
101            array: self.array.into_coord_type(coord_type),
102            ..self
103        }
104    }
105
106    /// Change the [`Metadata`] of this array.
107    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
108        Self {
109            data_type: self.data_type.with_metadata(metadata),
110            ..self
111        }
112    }
113}
114
115impl GeoArrowArray for GeometryCollectionArray {
116    fn as_any(&self) -> &dyn std::any::Any {
117        self
118    }
119
120    fn into_array_ref(self) -> ArrayRef {
121        Arc::new(self.into_arrow())
122    }
123
124    fn to_array_ref(&self) -> ArrayRef {
125        self.clone().into_array_ref()
126    }
127
128    #[inline]
129    fn len(&self) -> usize {
130        self.geom_offsets.len_proxy()
131    }
132
133    #[inline]
134    fn logical_nulls(&self) -> Option<NullBuffer> {
135        self.nulls.clone()
136    }
137
138    #[inline]
139    fn logical_null_count(&self) -> usize {
140        self.nulls.as_ref().map(|v| v.null_count()).unwrap_or(0)
141    }
142
143    #[inline]
144    fn is_null(&self, i: usize) -> bool {
145        self.nulls
146            .as_ref()
147            .map(|n| n.is_null(i))
148            .unwrap_or_default()
149    }
150
151    fn data_type(&self) -> GeoArrowType {
152        GeoArrowType::GeometryCollection(self.data_type.clone())
153    }
154
155    fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
156        Arc::new(self.slice(offset, length))
157    }
158
159    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
160        Arc::new(self.with_metadata(metadata))
161    }
162}
163
164impl<'a> GeoArrowArrayAccessor<'a> for GeometryCollectionArray {
165    type Item = GeometryCollection<'a>;
166
167    unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
168        Ok(GeometryCollection::new(
169            &self.array,
170            &self.geom_offsets,
171            index,
172        ))
173    }
174}
175
176impl IntoArrow for GeometryCollectionArray {
177    type ArrowArray = GenericListArray<i32>;
178    type ExtensionType = GeometryCollectionType;
179
180    fn into_arrow(self) -> Self::ArrowArray {
181        let geometries_field = self.geometries_field();
182        let nulls = self.nulls;
183        let values = self.array.into_array_ref();
184        GenericListArray::new(geometries_field, self.geom_offsets, values, nulls)
185    }
186
187    fn extension_type(&self) -> &Self::ExtensionType {
188        &self.data_type
189    }
190}
191
192impl TryFrom<(&GenericListArray<i32>, GeometryCollectionType)> for GeometryCollectionArray {
193    type Error = GeoArrowError;
194
195    fn try_from(
196        (value, typ): (&GenericListArray<i32>, GeometryCollectionType),
197    ) -> GeoArrowResult<Self> {
198        let geoms: MixedGeometryArray =
199            (value.values().as_ref(), typ.dimension(), typ.coord_type()).try_into()?;
200        let geom_offsets = value.offsets();
201        let nulls = value.nulls();
202
203        Ok(Self::new(
204            geoms,
205            geom_offsets.clone(),
206            nulls.cloned(),
207            typ.metadata().clone(),
208        ))
209    }
210}
211
212impl TryFrom<(&GenericListArray<i64>, GeometryCollectionType)> for GeometryCollectionArray {
213    type Error = GeoArrowError;
214
215    fn try_from(
216        (value, typ): (&GenericListArray<i64>, GeometryCollectionType),
217    ) -> GeoArrowResult<Self> {
218        let geoms: MixedGeometryArray =
219            (value.values().as_ref(), typ.dimension(), typ.coord_type()).try_into()?;
220        let geom_offsets = offsets_buffer_i64_to_i32(value.offsets())?;
221        let nulls = value.nulls();
222
223        Ok(Self::new(
224            geoms,
225            geom_offsets,
226            nulls.cloned(),
227            typ.metadata().clone(),
228        ))
229    }
230}
231
232impl TryFrom<(&dyn Array, GeometryCollectionType)> for GeometryCollectionArray {
233    type Error = GeoArrowError;
234
235    fn try_from((value, typ): (&dyn Array, GeometryCollectionType)) -> GeoArrowResult<Self> {
236        match value.data_type() {
237            DataType::List(_) => (value.as_list::<i32>(), typ).try_into(),
238            DataType::LargeList(_) => (value.as_list::<i64>(), typ).try_into(),
239            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
240                "Unexpected GeometryCollection Arrow DataType: {dt:?}"
241            ))),
242        }
243    }
244}
245
246impl TryFrom<(&dyn Array, &Field)> for GeometryCollectionArray {
247    type Error = GeoArrowError;
248
249    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
250        let typ = field.try_extension_type::<GeometryCollectionType>()?;
251        (arr, typ).try_into()
252    }
253}
254
255impl<O: OffsetSizeTrait> TryFrom<(GenericWkbArray<O>, GeometryCollectionType)>
256    for GeometryCollectionArray
257{
258    type Error = GeoArrowError;
259
260    fn try_from(value: (GenericWkbArray<O>, GeometryCollectionType)) -> GeoArrowResult<Self> {
261        let mut_arr: GeometryCollectionBuilder = value.try_into()?;
262        Ok(mut_arr.finish())
263    }
264}
265
266impl PartialEq for GeometryCollectionArray {
267    fn eq(&self, other: &Self) -> bool {
268        self.nulls == other.nulls
269            && offset_buffer_eq(&self.geom_offsets, &other.geom_offsets)
270            && self.array == other.array
271    }
272}
273
274#[cfg(test)]
275mod test {
276    use geoarrow_schema::{CoordType, Dimension};
277    use geoarrow_test::raw;
278
279    use super::*;
280    use crate::test::geometrycollection;
281
282    #[test]
283    fn try_from_arrow() {
284        for coord_type in [CoordType::Interleaved, CoordType::Separated] {
285            for dim in [
286                Dimension::XY,
287                Dimension::XYZ,
288                Dimension::XYM,
289                Dimension::XYZM,
290            ] {
291                for prefer_multi in [true, false] {
292                    let geo_arr = geometrycollection::array(coord_type, dim, prefer_multi);
293
294                    let point_type = geo_arr.extension_type().clone();
295                    let field = point_type.to_field("geometry", true);
296
297                    let arrow_arr = geo_arr.to_array_ref();
298
299                    let geo_arr2: GeometryCollectionArray =
300                        (arrow_arr.as_ref(), point_type).try_into().unwrap();
301                    let geo_arr3: GeometryCollectionArray =
302                        (arrow_arr.as_ref(), &field).try_into().unwrap();
303
304                    assert_eq!(geo_arr, geo_arr2);
305                    assert_eq!(geo_arr, geo_arr3);
306                }
307            }
308        }
309    }
310
311    #[test]
312    fn test_nullability() {
313        let geoms = raw::geometrycollection::xy::geoms();
314        let null_idxs = geoms
315            .iter()
316            .enumerate()
317            .filter_map(|(i, geom)| if geom.is_none() { Some(i) } else { None })
318            .collect::<Vec<_>>();
319
320        let typ = GeometryCollectionType::new(Dimension::XY, Default::default());
321        let geo_arr = GeometryCollectionBuilder::from_nullable_geometry_collections(&geoms, typ)
322            .unwrap()
323            .finish();
324
325        for null_idx in &null_idxs {
326            assert!(geo_arr.is_null(*null_idx));
327        }
328    }
329
330    #[test]
331    fn test_logical_nulls() {
332        let geoms = raw::geometrycollection::xy::geoms();
333        let expected_nulls = NullBuffer::from_iter(geoms.iter().map(|g| g.is_some()));
334
335        let typ = GeometryCollectionType::new(Dimension::XY, Default::default());
336        let geo_arr = GeometryCollectionBuilder::from_nullable_geometry_collections(&geoms, typ)
337            .unwrap()
338            .finish();
339
340        assert_eq!(geo_arr.logical_nulls().unwrap(), expected_nulls);
341    }
342}
geoarrow_array/array/geometrycollection.rs

geoarrow_array/array/
geometrycollection.rs
