geoarrow_array/builder/
rect.rs

1use arrow_buffer::NullBufferBuilder;
2use geo_traits::{CoordTrait, RectTrait};
3use geoarrow_schema::BoxType;
4use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
5
6use crate::array::RectArray;
7use crate::builder::SeparatedCoordBufferBuilder;
8use crate::scalar::Rect;
9
10/// The GeoArrow equivalent to `Vec<Option<Rect>>`: a mutable collection of Rects.
11///
12/// Converting an [`RectBuilder`] into a [`RectArray`] is `O(1)`.
13#[derive(Debug)]
14pub struct RectBuilder {
15    pub(crate) data_type: BoxType,
16    pub(crate) lower: SeparatedCoordBufferBuilder,
17    pub(crate) upper: SeparatedCoordBufferBuilder,
18    pub(crate) validity: NullBufferBuilder,
19}
20
21impl RectBuilder {
22    /// Creates a new empty [`RectBuilder`].
23    pub fn new(typ: BoxType) -> Self {
24        Self::with_capacity(typ, Default::default())
25    }
26
27    /// Creates a new [`RectBuilder`] with a capacity.
28    pub fn with_capacity(typ: BoxType, capacity: usize) -> Self {
29        Self {
30            lower: SeparatedCoordBufferBuilder::with_capacity(capacity, typ.dimension()),
31            upper: SeparatedCoordBufferBuilder::with_capacity(capacity, typ.dimension()),
32            validity: NullBufferBuilder::new(capacity),
33            data_type: typ,
34        }
35    }
36
37    /// Reserves capacity for at least `additional` more Rects.
38    ///
39    /// The collection may reserve more space to speculatively avoid frequent reallocations. After
40    /// calling `reserve`, capacity will be greater than or equal to `self.len() + additional`.
41    /// Does nothing if capacity is already sufficient.
42    pub fn reserve(&mut self, additional: usize) {
43        self.lower.reserve(additional);
44        self.upper.reserve(additional);
45    }
46
47    /// Reserves the minimum capacity for at least `additional` more Rects.
48    ///
49    /// Unlike [`reserve`], this will not deliberately over-allocate to speculatively avoid
50    /// frequent allocations. After calling `reserve_exact`, capacity will be greater than or equal
51    /// to `self.len() + additional`. Does nothing if the capacity is already sufficient.
52    ///
53    /// Note that the allocator may give the collection more space than it
54    /// requests. Therefore, capacity can not be relied upon to be precisely
55    /// minimal. Prefer [`reserve`] if future insertions are expected.
56    ///
57    /// [`reserve`]: Self::reserve
58    pub fn reserve_exact(&mut self, additional: usize) {
59        self.lower.reserve_exact(additional);
60        self.upper.reserve_exact(additional);
61    }
62
63    /// The canonical method to create a [`RectBuilder`] out of its internal components.
64    ///
65    /// # Implementation
66    ///
67    /// This function is `O(1)`.
68    ///
69    /// # Errors
70    ///
71    /// This function errors iff:
72    ///
73    /// - The validity is not `None` and its length is different from the number of geometries
74    pub fn try_new(
75        lower: SeparatedCoordBufferBuilder,
76        upper: SeparatedCoordBufferBuilder,
77        validity: NullBufferBuilder,
78        data_type: BoxType,
79    ) -> GeoArrowResult<Self> {
80        if lower.len() != upper.len() {
81            return Err(GeoArrowError::InvalidGeoArrow(
82                "Lower and upper lengths must match".to_string(),
83            ));
84        }
85        Ok(Self {
86            lower,
87            upper,
88            validity,
89            data_type,
90        })
91    }
92
93    /// Consume the builder and convert to an immutable [`RectArray`]
94    pub fn finish(mut self) -> RectArray {
95        RectArray::new(
96            self.lower.finish(),
97            self.upper.finish(),
98            self.validity.finish(),
99            self.data_type.metadata().clone(),
100        )
101    }
102
103    /// Add a new Rect to the end of this builder.
104    #[inline]
105    pub fn push_rect(&mut self, value: Option<&impl RectTrait<T = f64>>) {
106        if let Some(value) = value {
107            let min_coord = value.min();
108            let max_coord = value.max();
109
110            self.lower.push_coord(&min_coord);
111            self.upper.push_coord(&max_coord);
112            self.validity.append_non_null()
113        } else {
114            // Since it's a struct, we still need to push coords when null
115            self.lower.push_constant(f64::NAN);
116            self.upper.push_constant(f64::NAN);
117            self.validity.append_null();
118        }
119    }
120
121    /// Add a new null value to the end of this builder.
122    #[inline]
123    pub fn push_null(&mut self) {
124        self.push_rect(None::<&Rect>);
125    }
126
127    /// Push min and max coordinates of a rect to the builder.
128    #[inline]
129    pub fn push_min_max(&mut self, min: &impl CoordTrait<T = f64>, max: &impl CoordTrait<T = f64>) {
130        self.lower.push_coord(min);
131        self.upper.push_coord(max);
132        self.validity.append_non_null()
133    }
134
135    /// Create this builder from a iterator of Rects.
136    pub fn from_rects<'a>(
137        geoms: impl ExactSizeIterator<Item = &'a (impl RectTrait<T = f64> + 'a)>,
138        typ: BoxType,
139    ) -> Self {
140        let mut mutable_array = Self::with_capacity(typ, geoms.len());
141        geoms
142            .into_iter()
143            .for_each(|rect| mutable_array.push_rect(Some(rect)));
144        mutable_array
145    }
146
147    /// Create this builder from a iterator of nullable Rects.
148    pub fn from_nullable_rects<'a>(
149        geoms: impl ExactSizeIterator<Item = Option<&'a (impl RectTrait<T = f64> + 'a)>>,
150        typ: BoxType,
151    ) -> Self {
152        let mut mutable_array = Self::with_capacity(typ, geoms.len());
153        geoms
154            .into_iter()
155            .for_each(|maybe_rect| mutable_array.push_rect(maybe_rect));
156        mutable_array
157    }
158}