geoarrow_array/array/
rect.rs

1use std::sync::Arc;
2
3use arrow_array::cast::AsArray;
4use arrow_array::types::Float64Type;
5use arrow_array::{Array, ArrayRef, StructArray};
6use arrow_buffer::NullBuffer;
7use arrow_schema::{DataType, Field};
8use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
9use geoarrow_schema::{BoxType, GeoArrowType, Metadata};
10
11use crate::array::SeparatedCoordBuffer;
12use crate::scalar::Rect;
13use crate::trait_::{GeoArrowArray, GeoArrowArrayAccessor, IntoArrow};
14
15/// An immutable array of Rect or Box geometries.
16///
17/// A rect is an axis-aligned bounded rectangle whose area is defined by minimum and maximum
18/// coordinates.
19///
20/// All rects must have the same dimension.
21///
22/// This is **not** an array type defined by the GeoArrow specification (as of spec version 0.1)
23/// but is included here for parity with georust/geo, and to save memory for the output of
24/// `bounds()`.
25///
26/// Internally this is implemented as a FixedSizeList, laid out as minx, miny, maxx, maxy.
27#[derive(Debug, Clone)]
28pub struct RectArray {
29    pub(crate) data_type: BoxType,
30
31    /// Separated arrays for each of the "lower" dimensions
32    lower: SeparatedCoordBuffer,
33
34    /// Separated arrays for each of the "upper" dimensions
35    upper: SeparatedCoordBuffer,
36
37    nulls: Option<NullBuffer>,
38}
39
40impl RectArray {
41    /// Construct a new [`RectArray`] from parts
42    pub fn new(
43        lower: SeparatedCoordBuffer,
44        upper: SeparatedCoordBuffer,
45        nulls: Option<NullBuffer>,
46        metadata: Arc<Metadata>,
47    ) -> Self {
48        assert_eq!(lower.dim(), upper.dim());
49        Self {
50            data_type: BoxType::new(lower.dim(), metadata),
51            lower,
52            upper,
53            nulls,
54        }
55    }
56
57    /// Access the coordinate buffer of the "lower" corner of the RectArray
58    ///
59    /// Note that this needs to be interpreted in conjunction with the [null
60    /// buffer][Self::logical_nulls].
61    pub fn lower(&self) -> &SeparatedCoordBuffer {
62        &self.lower
63    }
64
65    /// Access the coordinate buffer of the "upper" corner of the RectArray
66    ///
67    /// Note that this needs to be interpreted in conjunction with the [null
68    /// buffer][Self::logical_nulls].
69    pub fn upper(&self) -> &SeparatedCoordBuffer {
70        &self.upper
71    }
72
73    /// Slice this [`RectArray`].
74    ///
75    /// # Panic
76    /// This function panics iff `offset + length > self.len()`.
77    #[inline]
78    pub fn slice(&self, offset: usize, length: usize) -> Self {
79        assert!(
80            offset + length <= self.len(),
81            "offset + length may not exceed length of array"
82        );
83
84        Self {
85            data_type: self.data_type.clone(),
86            lower: self.lower().slice(offset, length),
87            upper: self.upper().slice(offset, length),
88            nulls: self.nulls.as_ref().map(|v| v.slice(offset, length)),
89        }
90    }
91
92    /// Change the [`Metadata`] of this array.
93    pub fn with_metadata(self, metadata: Arc<Metadata>) -> Self {
94        Self {
95            data_type: self.data_type.with_metadata(metadata),
96            ..self
97        }
98    }
99}
100
101impl GeoArrowArray for RectArray {
102    fn as_any(&self) -> &dyn std::any::Any {
103        self
104    }
105
106    fn into_array_ref(self) -> ArrayRef {
107        Arc::new(self.into_arrow())
108    }
109
110    fn to_array_ref(&self) -> ArrayRef {
111        self.clone().into_array_ref()
112    }
113
114    #[inline]
115    fn len(&self) -> usize {
116        self.lower.len()
117    }
118
119    #[inline]
120    fn logical_nulls(&self) -> Option<NullBuffer> {
121        self.nulls.clone()
122    }
123
124    #[inline]
125    fn logical_null_count(&self) -> usize {
126        self.nulls.as_ref().map(|v| v.null_count()).unwrap_or(0)
127    }
128
129    #[inline]
130    fn is_null(&self, i: usize) -> bool {
131        self.nulls
132            .as_ref()
133            .map(|n| n.is_null(i))
134            .unwrap_or_default()
135    }
136
137    fn data_type(&self) -> GeoArrowType {
138        GeoArrowType::Rect(self.data_type.clone())
139    }
140
141    fn slice(&self, offset: usize, length: usize) -> Arc<dyn GeoArrowArray> {
142        Arc::new(self.slice(offset, length))
143    }
144
145    fn with_metadata(self, metadata: Arc<Metadata>) -> Arc<dyn GeoArrowArray> {
146        Arc::new(self.with_metadata(metadata))
147    }
148}
149
150impl<'a> GeoArrowArrayAccessor<'a> for RectArray {
151    type Item = Rect<'a>;
152
153    unsafe fn value_unchecked(&'a self, index: usize) -> GeoArrowResult<Self::Item> {
154        Ok(Rect::new(&self.lower, &self.upper, index))
155    }
156}
157
158impl IntoArrow for RectArray {
159    type ArrowArray = StructArray;
160    type ExtensionType = BoxType;
161
162    fn into_arrow(self) -> Self::ArrowArray {
163        let fields = match self.data_type.data_type() {
164            DataType::Struct(fields) => fields,
165            _ => unreachable!(),
166        };
167
168        let mut arrays: Vec<ArrayRef> = vec![];
169
170        // values_array takes care of the correct number of dimensions
171        arrays.extend_from_slice(self.lower.values_array().as_slice());
172        arrays.extend_from_slice(self.upper.values_array().as_slice());
173
174        let nulls = self.nulls;
175        StructArray::new(fields, arrays, nulls)
176    }
177
178    fn extension_type(&self) -> &Self::ExtensionType {
179        &self.data_type
180    }
181}
182
183impl TryFrom<(&StructArray, BoxType)> for RectArray {
184    type Error = GeoArrowError;
185
186    fn try_from((value, typ): (&StructArray, BoxType)) -> GeoArrowResult<Self> {
187        let dim = typ.dimension();
188        let nulls = value.nulls();
189        let columns = value.columns();
190        if columns.len() != dim.size() * 2 {
191            return Err(GeoArrowError::InvalidGeoArrow(format!(
192                "Invalid number of columns for RectArray: expected {} but got {}",
193                dim.size() * 2,
194                columns.len()
195            )));
196        }
197
198        let lower = columns[0..dim.size()]
199            .iter()
200            .map(|c| c.as_primitive::<Float64Type>().values().clone())
201            .collect::<Vec<_>>();
202        let lower = SeparatedCoordBuffer::from_vec(lower, dim)?;
203
204        let upper = columns[dim.size()..]
205            .iter()
206            .map(|c| c.as_primitive::<Float64Type>().values().clone())
207            .collect::<Vec<_>>();
208        let upper = SeparatedCoordBuffer::from_vec(upper, dim)?;
209
210        Ok(Self::new(
211            lower,
212            upper,
213            nulls.cloned(),
214            typ.metadata().clone(),
215        ))
216    }
217}
218
219impl TryFrom<(&dyn Array, BoxType)> for RectArray {
220    type Error = GeoArrowError;
221
222    fn try_from((value, dim): (&dyn Array, BoxType)) -> GeoArrowResult<Self> {
223        match value.data_type() {
224            DataType::Struct(_) => (value.as_struct(), dim).try_into(),
225            dt => Err(GeoArrowError::InvalidGeoArrow(format!(
226                "Unexpected Rect DataType: {:?}",
227                dt
228            ))),
229        }
230    }
231}
232
233impl TryFrom<(&dyn Array, &Field)> for RectArray {
234    type Error = GeoArrowError;
235
236    fn try_from((arr, field): (&dyn Array, &Field)) -> GeoArrowResult<Self> {
237        let typ = field.try_extension_type::<BoxType>()?;
238        (arr, typ).try_into()
239    }
240}
241
242impl PartialEq for RectArray {
243    fn eq(&self, other: &Self) -> bool {
244        // A naive implementation of PartialEq would check for buffer equality. This won't always
245        // work for null elements where the actual value can be undefined and doesn't have to be
246        // equal. As such, it's simplest to reuse the upstream PartialEq impl, especially since
247        // RectArray only has one coordinate type.
248        self.clone().into_arrow() == other.clone().into_arrow()
249    }
250}
251
252#[cfg(test)]
253mod test {
254    use geo_traits::to_geo::ToGeoRect;
255    use geoarrow_schema::Dimension;
256
257    use super::*;
258    use crate::builder::RectBuilder;
259    use crate::test::rect;
260
261    #[test]
262    fn geo_round_trip() {
263        let geoms = [Some(rect::r0()), None, Some(rect::r1()), None];
264        let typ = BoxType::new(Dimension::XY, Default::default());
265        let geo_arr =
266            RectBuilder::from_nullable_rects(geoms.iter().map(|x| x.as_ref()), typ).finish();
267
268        for (i, g) in geo_arr.iter().enumerate() {
269            assert_eq!(geoms[i], g.transpose().unwrap().map(|g| g.to_rect()));
270        }
271
272        // Test sliced
273        for (i, g) in geo_arr.slice(2, 2).iter().enumerate() {
274            assert_eq!(geoms[i + 2], g.transpose().unwrap().map(|g| g.to_rect()));
275        }
276    }
277
278    #[test]
279    fn try_from_arrow() {
280        let geo_arr = rect::r_array();
281
282        let extension_type = geo_arr.extension_type().clone();
283        let field = extension_type.to_field("geometry", true);
284
285        let arrow_arr = geo_arr.to_array_ref();
286
287        let geo_arr2: RectArray = (arrow_arr.as_ref(), extension_type).try_into().unwrap();
288        let geo_arr3: RectArray = (arrow_arr.as_ref(), &field).try_into().unwrap();
289
290        assert_eq!(geo_arr, geo_arr2);
291        assert_eq!(geo_arr, geo_arr3);
292    }
293
294    #[test]
295    fn partial_eq() {
296        let arr1 = rect::r_array();
297        assert_eq!(arr1, arr1);
298    }
299}
geoarrow_array/array/rect.rs

geoarrow_array/array/
rect.rs
