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