geoarrow_array/builder/coord/
separated.rs

1use geo_traits::{CoordTrait, PointTrait};
2use geoarrow_schema::Dimension;
3use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
4
5use crate::array::SeparatedCoordBuffer;
6
7/// The GeoArrow equivalent to `Vec<Option<Coord>>`: a mutable collection of coordinates.
8///
9/// This stores all coordinates in separated fashion as multiple arrays: `xxx` and `yyy`.
10///
11/// Converting an [`SeparatedCoordBufferBuilder`] into a [`SeparatedCoordBuffer`] is `O(1)`.
12#[derive(Debug, Clone)]
13pub struct SeparatedCoordBufferBuilder {
14    buffers: [Vec<f64>; 4],
15    dim: Dimension,
16}
17
18impl SeparatedCoordBufferBuilder {
19    /// Create a new empty builder with the given dimension
20    pub fn new(dim: Dimension) -> Self {
21        Self::with_capacity(0, dim)
22    }
23
24    /// Create a new builder with the given capacity and dimension
25    pub fn with_capacity(capacity: usize, dim: Dimension) -> Self {
26        // Only allocate buffers for existant dimensions
27        let buffers = core::array::from_fn(|i| {
28            if i < dim.size() {
29                Vec::with_capacity(capacity)
30            } else {
31                Vec::new()
32            }
33        });
34
35        Self { buffers, dim }
36    }
37
38    /// Initialize a buffer of a given length with all coordinates set to the given value
39    pub fn initialize(len: usize, dim: Dimension, value: f64) -> Self {
40        // Only allocate buffers for existant dimensions
41        let buffers = core::array::from_fn(|i| {
42            if i < dim.size() {
43                vec![value; len]
44            } else {
45                Vec::new()
46            }
47        });
48
49        Self { buffers, dim }
50    }
51
52    /// Reserves capacity for at least `additional` more coordinates.
53    ///
54    /// The collection may reserve more space to speculatively avoid frequent reallocations. After
55    /// calling `reserve`, capacity will be greater than or equal to `self.len() + additional`.
56    /// Does nothing if capacity is already sufficient.
57    pub fn reserve(&mut self, additional: usize) {
58        self.buffers
59            .iter_mut()
60            .for_each(|buffer| buffer.reserve(additional))
61    }
62
63    /// Reserves the minimum capacity for at least `additional` more coordinates.
64    ///
65    /// Unlike [`reserve`], this will not deliberately over-allocate to speculatively avoid
66    /// frequent allocations. After calling `reserve_exact`, capacity will be greater than or equal
67    /// to `self.len() + additional`. Does nothing if the capacity is already sufficient.
68    ///
69    /// Note that the allocator may give the collection more space than it
70    /// requests. Therefore, capacity can not be relied upon to be precisely
71    /// minimal. Prefer [`reserve`] if future insertions are expected.
72    ///
73    /// [`reserve`]: Self::reserve
74    pub fn reserve_exact(&mut self, additional: usize) {
75        self.buffers
76            .iter_mut()
77            .for_each(|buffer| buffer.reserve_exact(additional))
78    }
79
80    /// Shrinks the capacity of self to fit.
81    pub fn shrink_to_fit(&mut self) {
82        self.buffers
83            .iter_mut()
84            .for_each(|buffer| buffer.shrink_to_fit());
85    }
86
87    /// Returns the total number of coordinates the vector can hold without reallocating.
88    pub fn capacity(&self) -> usize {
89        self.buffers[0].capacity()
90    }
91
92    /// The number of coordinates in this builder
93    pub fn len(&self) -> usize {
94        self.buffers[0].len()
95    }
96
97    /// Whether this builder is empty
98    pub fn is_empty(&self) -> bool {
99        self.len() == 0
100    }
101
102    /// Push a new coord onto the end of this coordinate buffer
103    ///
104    /// ## Panics
105    ///
106    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
107    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) {
108        self.try_push_coord(coord).unwrap()
109    }
110
111    /// Push a new coord onto the end of this coordinate buffer
112    ///
113    /// ## Errors
114    ///
115    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
116    pub fn try_push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> GeoArrowResult<()> {
117        // Note duplicated across buffer types; consider refactoring
118        match self.dim {
119            Dimension::XY => match coord.dim() {
120                geo_traits::Dimensions::Xy | geo_traits::Dimensions::Unknown(2) => {}
121                d => {
122                    return Err(GeoArrowError::IncorrectGeometryType(format!(
123                        "coord dimension must be XY for this buffer; got {d:?}."
124                    )));
125                }
126            },
127            Dimension::XYZ => match coord.dim() {
128                geo_traits::Dimensions::Xyz | geo_traits::Dimensions::Unknown(3) => {}
129                d => {
130                    return Err(GeoArrowError::IncorrectGeometryType(format!(
131                        "coord dimension must be XYZ for this buffer; got {d:?}."
132                    )));
133                }
134            },
135            Dimension::XYM => match coord.dim() {
136                geo_traits::Dimensions::Xym | geo_traits::Dimensions::Unknown(3) => {}
137                d => {
138                    return Err(GeoArrowError::IncorrectGeometryType(format!(
139                        "coord dimension must be XYM for this buffer; got {d:?}."
140                    )));
141                }
142            },
143            Dimension::XYZM => match coord.dim() {
144                geo_traits::Dimensions::Xyzm | geo_traits::Dimensions::Unknown(4) => {}
145                d => {
146                    return Err(GeoArrowError::IncorrectGeometryType(format!(
147                        "coord dimension must be XYZM for this buffer; got {d:?}."
148                    )));
149                }
150            },
151        }
152
153        self.buffers[0].push(coord.x());
154        self.buffers[1].push(coord.y());
155        if let Some(z) = coord.nth(2) {
156            self.buffers[2].push(z);
157        };
158        if let Some(m) = coord.nth(3) {
159            self.buffers[3].push(m);
160        };
161        Ok(())
162    }
163
164    /// Push a valid coordinate with the given constant value
165    ///
166    /// Used in the case of point and rect arrays, where a `null` array value still needs to have
167    /// space allocated for it.
168    pub(crate) fn push_constant(&mut self, value: f64) {
169        for i in 0..self.dim.size() {
170            self.buffers[i].push(value);
171        }
172    }
173
174    /// Push a new point onto the end of this coordinate buffer
175    ///
176    /// ## Panics
177    ///
178    /// - If the added point does not have the same dimension as the coordinate buffer.
179    pub(crate) fn push_point(&mut self, point: &impl PointTrait<T = f64>) {
180        self.try_push_point(point).unwrap()
181    }
182
183    /// Push a new point onto the end of this coordinate buffer
184    ///
185    /// ## Errors
186    ///
187    /// - If the added point does not have the same dimension as the coordinate buffer.
188    pub(crate) fn try_push_point(
189        &mut self,
190        point: &impl PointTrait<T = f64>,
191    ) -> GeoArrowResult<()> {
192        if let Some(coord) = point.coord() {
193            self.try_push_coord(&coord)?;
194        } else {
195            self.push_constant(f64::NAN);
196        };
197        Ok(())
198    }
199
200    /// Construct a new builder and pre-fill it with coordinates from the provided iterator
201    pub fn from_coords<'a>(
202        coords: impl ExactSizeIterator<Item = &'a (impl CoordTrait<T = f64> + 'a)>,
203        dim: Dimension,
204    ) -> GeoArrowResult<Self> {
205        let mut buffer = SeparatedCoordBufferBuilder::with_capacity(coords.len(), dim);
206        for coord in coords {
207            buffer.try_push_coord(coord)?;
208        }
209        Ok(buffer)
210    }
211
212    /// Consume the builder and convert to an immutable [`SeparatedCoordBuffer`]
213    pub fn finish(self) -> SeparatedCoordBuffer {
214        // Initialize buffers with empty array, then mutate into it
215        let mut buffers = core::array::from_fn(|_| vec![].into());
216        for (i, buffer) in self.buffers.into_iter().enumerate() {
217            buffers[i] = buffer.into();
218        }
219        SeparatedCoordBuffer::from_array(buffers, self.dim).unwrap()
220    }
221}
222
223#[cfg(test)]
224mod test {
225    use wkt::types::Coord;
226
227    use super::*;
228
229    #[test]
230    fn errors_when_pushing_incompatible_coord() {
231        let mut builder = SeparatedCoordBufferBuilder::new(Dimension::XY);
232        builder
233            .try_push_coord(&Coord {
234                x: 0.0,
235                y: 0.0,
236                z: Some(0.0),
237                m: None,
238            })
239            .expect_err("Should err pushing XYZ to XY buffer");
240
241        let mut builder = SeparatedCoordBufferBuilder::new(Dimension::XYZ);
242        builder
243            .try_push_coord(&Coord {
244                x: 0.0,
245                y: 0.0,
246                z: None,
247                m: None,
248            })
249            .expect_err("Should err pushing XY to XYZ buffer");
250        builder
251            .try_push_coord(&Coord {
252                x: 0.0,
253                y: 0.0,
254                z: Some(0.0),
255                m: None,
256            })
257            .unwrap();
258    }
259}
geoarrow_array/builder/coord/separated.rs

geoarrow_array/builder/coord/
separated.rs
