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    /// Shrinks the capacity of self to fit.
64    pub fn shrink_to_fit(&mut self) {
65        self.lower.shrink_to_fit();
66        self.upper.shrink_to_fit();
67        // self.validity.shrink_to_fit();
68    }
69
70    /// The canonical method to create a [`RectBuilder`] out of its internal components.
71    ///
72    /// # Implementation
73    ///
74    /// This function is `O(1)`.
75    ///
76    /// # Errors
77    ///
78    /// This function errors iff:
79    ///
80    /// - The validity is not `None` and its length is different from the number of geometries
81    pub fn try_new(
82        lower: SeparatedCoordBufferBuilder,
83        upper: SeparatedCoordBufferBuilder,
84        validity: NullBufferBuilder,
85        data_type: BoxType,
86    ) -> GeoArrowResult<Self> {
87        if lower.len() != upper.len() {
88            return Err(GeoArrowError::InvalidGeoArrow(
89                "Lower and upper lengths must match".to_string(),
90            ));
91        }
92        Ok(Self {
93            lower,
94            upper,
95            validity,
96            data_type,
97        })
98    }
99
100    /// Consume the builder and convert to an immutable [`RectArray`]
101    pub fn finish(mut self) -> RectArray {
102        RectArray::new(
103            self.lower.finish(),
104            self.upper.finish(),
105            self.validity.finish(),
106            self.data_type.metadata().clone(),
107        )
108    }
109
110    /// Add a new Rect to the end of this builder.
111    #[inline]
112    pub fn push_rect(&mut self, value: Option<&impl RectTrait<T = f64>>) {
113        if let Some(value) = value {
114            let min_coord = value.min();
115            let max_coord = value.max();
116
117            self.lower.push_coord(&min_coord);
118            self.upper.push_coord(&max_coord);
119            self.validity.append_non_null()
120        } else {
121            // Since it's a struct, we still need to push coords when null
122            self.lower.push_constant(f64::NAN);
123            self.upper.push_constant(f64::NAN);
124            self.validity.append_null();
125        }
126    }
127
128    /// Add a new null value to the end of this builder.
129    #[inline]
130    pub fn push_null(&mut self) {
131        self.push_rect(None::<&Rect>);
132    }
133
134    /// Push min and max coordinates of a rect to the builder.
135    #[inline]
136    pub fn push_min_max(&mut self, min: &impl CoordTrait<T = f64>, max: &impl CoordTrait<T = f64>) {
137        self.lower.push_coord(min);
138        self.upper.push_coord(max);
139        self.validity.append_non_null()
140    }
141
142    /// Create this builder from a iterator of Rects.
143    pub fn from_rects<'a>(
144        geoms: impl ExactSizeIterator<Item = &'a (impl RectTrait<T = f64> + 'a)>,
145        typ: BoxType,
146    ) -> Self {
147        let mut mutable_array = Self::with_capacity(typ, geoms.len());
148        geoms
149            .into_iter()
150            .for_each(|rect| mutable_array.push_rect(Some(rect)));
151        mutable_array
152    }
153
154    /// Create this builder from a iterator of nullable Rects.
155    pub fn from_nullable_rects<'a>(
156        geoms: impl ExactSizeIterator<Item = Option<&'a (impl RectTrait<T = f64> + 'a)>>,
157        typ: BoxType,
158    ) -> Self {
159        let mut mutable_array = Self::with_capacity(typ, geoms.len());
160        geoms
161            .into_iter()
162            .for_each(|maybe_rect| mutable_array.push_rect(maybe_rect));
163        mutable_array
164    }
165}