geoarrow_array/capacity/
linestring.rs

1use std::ops::Add;
2
3use geo_traits::{GeometryTrait, GeometryType, LineStringTrait};
4use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
5
6use crate::util::GeometryTypeName;
7
8/// A counter for the buffer sizes of a [`LineStringArray`][crate::array::LineStringArray].
9///
10/// This can be used to reduce allocations by allocating once for exactly the array size you need.
11#[derive(Debug, Clone, Copy)]
12pub struct LineStringCapacity {
13    pub(crate) coord_capacity: usize,
14    pub(crate) geom_capacity: usize,
15}
16
17impl LineStringCapacity {
18    /// Create a new capacity with known sizes.
19    pub fn new(coord_capacity: usize, geom_capacity: usize) -> Self {
20        Self {
21            coord_capacity,
22            geom_capacity,
23        }
24    }
25
26    /// Create a new empty capacity.
27    pub fn new_empty() -> Self {
28        Self::new(0, 0)
29    }
30
31    /// Return `true` if the capacity is empty.
32    pub fn is_empty(&self) -> bool {
33        self.coord_capacity == 0 && self.geom_capacity == 0
34    }
35
36    /// Add a LineString to this capacity counter.
37    #[inline]
38    pub fn add_line_string(&mut self, maybe_line_string: Option<&impl LineStringTrait>) {
39        self.geom_capacity += 1;
40        if let Some(line_string) = maybe_line_string {
41            self.add_valid_line_string(line_string);
42        }
43    }
44
45    #[inline]
46    fn add_valid_line_string(&mut self, line_string: &impl LineStringTrait) {
47        self.coord_capacity += line_string.num_coords();
48    }
49
50    /// Add the capacity of the given Geometry
51    ///
52    /// The type of the geometry must be LineString
53    #[inline]
54    pub fn add_geometry(&mut self, value: Option<&impl GeometryTrait>) -> GeoArrowResult<()> {
55        self.geom_capacity += 1;
56
57        if let Some(g) = value {
58            match g.as_type() {
59                GeometryType::LineString(p) => self.add_valid_line_string(p),
60                gt => {
61                    return Err(GeoArrowError::IncorrectGeometryType(format!(
62                        "Expected LineString, got {}",
63                        gt.name()
64                    )));
65                }
66            }
67        };
68        Ok(())
69    }
70
71    /// The coordinate buffer capacity
72    pub fn coord_capacity(&self) -> usize {
73        self.coord_capacity
74    }
75
76    /// The geometry offset buffer capacity
77    pub fn geom_capacity(&self) -> usize {
78        self.geom_capacity
79    }
80
81    /// Create a capacity counter from an iterator of LineStrings.
82    pub fn from_line_strings<'a>(
83        geoms: impl Iterator<Item = Option<&'a (impl LineStringTrait + 'a)>>,
84    ) -> Self {
85        let mut counter = Self::new_empty();
86
87        for maybe_line_string in geoms.into_iter() {
88            counter.add_line_string(maybe_line_string);
89        }
90
91        counter
92    }
93
94    /// Construct a new counter pre-filled with the given geometries
95    pub fn from_geometries<'a>(
96        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait + 'a)>>,
97    ) -> GeoArrowResult<Self> {
98        let mut counter = Self::new_empty();
99        for g in geoms.into_iter() {
100            counter.add_geometry(g)?;
101        }
102        Ok(counter)
103    }
104
105    /// The number of bytes an array with this capacity would occupy.
106    pub fn num_bytes(&self) -> usize {
107        let offsets_byte_width = 4;
108        let num_offsets = self.geom_capacity;
109        (offsets_byte_width * num_offsets) + (self.coord_capacity * 2 * 8)
110    }
111}
112
113impl Default for LineStringCapacity {
114    fn default() -> Self {
115        Self::new_empty()
116    }
117}
118
119impl Add for LineStringCapacity {
120    type Output = Self;
121
122    fn add(self, rhs: Self) -> Self::Output {
123        let coord_capacity = self.coord_capacity + rhs.coord_capacity;
124        let geom_capacity = self.geom_capacity + rhs.geom_capacity;
125        Self::new(coord_capacity, geom_capacity)
126    }
127}