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    /// Returns the total number of coordinates the vector can hold without reallocating.
81    pub fn capacity(&self) -> usize {
82        self.buffers[0].capacity()
83    }
84
85    /// The number of coordinates in this builder
86    pub fn len(&self) -> usize {
87        self.buffers[0].len()
88    }
89
90    /// Whether this builder is empty
91    pub fn is_empty(&self) -> bool {
92        self.len() == 0
93    }
94
95    /// Push a new coord onto the end of this coordinate buffer
96    ///
97    /// ## Panics
98    ///
99    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
100    pub fn push_coord(&mut self, coord: &impl CoordTrait<T = f64>) {
101        self.try_push_coord(coord).unwrap()
102    }
103
104    /// Push a new coord onto the end of this coordinate buffer
105    ///
106    /// ## Errors
107    ///
108    /// - If the added coordinate does not have the same dimension as the coordinate buffer.
109    pub fn try_push_coord(&mut self, coord: &impl CoordTrait<T = f64>) -> GeoArrowResult<()> {
110        // Note duplicated across buffer types; consider refactoring
111        match self.dim {
112            Dimension::XY => match coord.dim() {
113                geo_traits::Dimensions::Xy | geo_traits::Dimensions::Unknown(2) => {}
114                d => {
115                    return Err(GeoArrowError::IncorrectGeometryType(format!(
116                        "coord dimension must be XY for this buffer; got {d:?}."
117                    )));
118                }
119            },
120            Dimension::XYZ => match coord.dim() {
121                geo_traits::Dimensions::Xyz | geo_traits::Dimensions::Unknown(3) => {}
122                d => {
123                    return Err(GeoArrowError::IncorrectGeometryType(format!(
124                        "coord dimension must be XYZ for this buffer; got {d:?}."
125                    )));
126                }
127            },
128            Dimension::XYM => match coord.dim() {
129                geo_traits::Dimensions::Xym | geo_traits::Dimensions::Unknown(3) => {}
130                d => {
131                    return Err(GeoArrowError::IncorrectGeometryType(format!(
132                        "coord dimension must be XYM for this buffer; got {d:?}."
133                    )));
134                }
135            },
136            Dimension::XYZM => match coord.dim() {
137                geo_traits::Dimensions::Xyzm | geo_traits::Dimensions::Unknown(4) => {}
138                d => {
139                    return Err(GeoArrowError::IncorrectGeometryType(format!(
140                        "coord dimension must be XYZM for this buffer; got {d:?}."
141                    )));
142                }
143            },
144        }
145
146        self.buffers[0].push(coord.x());
147        self.buffers[1].push(coord.y());
148        if let Some(z) = coord.nth(2) {
149            self.buffers[2].push(z);
150        };
151        if let Some(m) = coord.nth(3) {
152            self.buffers[3].push(m);
153        };
154        Ok(())
155    }
156
157    /// Push a valid coordinate with the given constant value
158    ///
159    /// Used in the case of point and rect arrays, where a `null` array value still needs to have
160    /// space allocated for it.
161    pub(crate) fn push_constant(&mut self, value: f64) {
162        for i in 0..self.dim.size() {
163            self.buffers[i].push(value);
164        }
165    }
166
167    /// Push a new point onto the end of this coordinate buffer
168    ///
169    /// ## Panics
170    ///
171    /// - If the added point does not have the same dimension as the coordinate buffer.
172    pub(crate) fn push_point(&mut self, point: &impl PointTrait<T = f64>) {
173        self.try_push_point(point).unwrap()
174    }
175
176    /// Push a new point onto the end of this coordinate buffer
177    ///
178    /// ## Errors
179    ///
180    /// - If the added point does not have the same dimension as the coordinate buffer.
181    pub(crate) fn try_push_point(
182        &mut self,
183        point: &impl PointTrait<T = f64>,
184    ) -> GeoArrowResult<()> {
185        if let Some(coord) = point.coord() {
186            self.try_push_coord(&coord)?;
187        } else {
188            self.push_constant(f64::NAN);
189        };
190        Ok(())
191    }
192
193    /// Construct a new builder and pre-fill it with coordinates from the provided iterator
194    pub fn from_coords<'a>(
195        coords: impl ExactSizeIterator<Item = &'a (impl CoordTrait<T = f64> + 'a)>,
196        dim: Dimension,
197    ) -> GeoArrowResult<Self> {
198        let mut buffer = SeparatedCoordBufferBuilder::with_capacity(coords.len(), dim);
199        for coord in coords {
200            buffer.try_push_coord(coord)?;
201        }
202        Ok(buffer)
203    }
204
205    /// Consume the builder and convert to an immutable [`SeparatedCoordBuffer`]
206    pub fn finish(self) -> SeparatedCoordBuffer {
207        // Initialize buffers with empty array, then mutate into it
208        let mut buffers = core::array::from_fn(|_| vec![].into());
209        for (i, buffer) in self.buffers.into_iter().enumerate() {
210            buffers[i] = buffer.into();
211        }
212        SeparatedCoordBuffer::from_array(buffers, self.dim).unwrap()
213    }
214}
215
216#[cfg(test)]
217mod test {
218    use wkt::types::Coord;
219
220    use super::*;
221
222    #[test]
223    fn errors_when_pushing_incompatible_coord() {
224        let mut builder = SeparatedCoordBufferBuilder::new(Dimension::XY);
225        builder
226            .try_push_coord(&Coord {
227                x: 0.0,
228                y: 0.0,
229                z: Some(0.0),
230                m: None,
231            })
232            .expect_err("Should err pushing XYZ to XY buffer");
233
234        let mut builder = SeparatedCoordBufferBuilder::new(Dimension::XYZ);
235        builder
236            .try_push_coord(&Coord {
237                x: 0.0,
238                y: 0.0,
239                z: None,
240                m: None,
241            })
242            .expect_err("Should err pushing XY to XYZ buffer");
243        builder
244            .try_push_coord(&Coord {
245                x: 0.0,
246                y: 0.0,
247                z: Some(0.0),
248                m: None,
249            })
250            .unwrap();
251    }
252}
geoarrow_array/builder/coord/separated.rs

geoarrow_array/builder/coord/
separated.rs
