geoarrow_array/builder/coord/
interleaved.rs

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