geoarrow_array/capacity/
multilinestring.rs

1use std::ops::{Add, AddAssign};
2
3use geo_traits::{GeometryTrait, GeometryType, LineStringTrait, MultiLineStringTrait};
4use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
5
6use crate::capacity::LineStringCapacity;
7use crate::util::GeometryTypeName;
8
9/// A counter for the buffer sizes of a
10/// [`MultiLineStringArray`][crate::array::MultiLineStringArray].
11///
12/// This can be used to reduce allocations by allocating once for exactly the array size you need.
13#[derive(Debug, Clone, Copy)]
14pub struct MultiLineStringCapacity {
15    pub(crate) coord_capacity: usize,
16    pub(crate) ring_capacity: usize,
17    pub(crate) geom_capacity: usize,
18}
19
20impl MultiLineStringCapacity {
21    /// Create a new capacity with known sizes.
22    pub fn new(coord_capacity: usize, ring_capacity: usize, geom_capacity: usize) -> Self {
23        Self {
24            coord_capacity,
25            ring_capacity,
26            geom_capacity,
27        }
28    }
29
30    /// Create a new empty capacity.
31    pub fn new_empty() -> Self {
32        Self::new(0, 0, 0)
33    }
34
35    /// Return `true` if the capacity is empty.
36    pub fn is_empty(&self) -> bool {
37        self.coord_capacity == 0 && self.ring_capacity == 0 && self.geom_capacity == 0
38    }
39
40    /// The coordinate buffer capacity
41    pub fn coord_capacity(&self) -> usize {
42        self.coord_capacity
43    }
44
45    /// The ring offset buffer capacity
46    pub fn ring_capacity(&self) -> usize {
47        self.ring_capacity
48    }
49
50    /// The geometry offset buffer capacity
51    pub fn geom_capacity(&self) -> usize {
52        self.geom_capacity
53    }
54
55    /// Add the capacity of the given LineString
56    #[inline]
57    pub fn add_line_string(&mut self, maybe_line_string: Option<&impl LineStringTrait>) {
58        self.geom_capacity += 1;
59        if let Some(line_string) = maybe_line_string {
60            // A single line string
61            self.ring_capacity += 1;
62            self.coord_capacity += line_string.num_coords();
63        }
64    }
65
66    /// Add the capacity of the given MultiLineString
67    #[inline]
68    pub fn add_multi_line_string(&mut self, multi_line_string: Option<&impl MultiLineStringTrait>) {
69        self.geom_capacity += 1;
70        if let Some(multi_line_string) = multi_line_string {
71            // Total number of rings in this polygon
72            let num_line_strings = multi_line_string.num_line_strings();
73            self.ring_capacity += num_line_strings;
74
75            for line_string in multi_line_string.line_strings() {
76                self.coord_capacity += line_string.num_coords();
77            }
78        }
79    }
80
81    /// Add the capacity of the given Geometry
82    ///
83    /// The type of the geometry must be either LineString or MultiLineString
84    #[inline]
85    pub fn add_geometry(&mut self, value: Option<&impl GeometryTrait>) -> GeoArrowResult<()> {
86        if let Some(geom) = value {
87            match geom.as_type() {
88                GeometryType::LineString(g) => self.add_line_string(Some(g)),
89                GeometryType::MultiLineString(g) => self.add_multi_line_string(Some(g)),
90                gt => {
91                    return Err(GeoArrowError::IncorrectGeometryType(format!(
92                        "Expected LineString or MultiLineString, got {}",
93                        gt.name()
94                    )));
95                }
96            }
97        } else {
98            self.geom_capacity += 1;
99        };
100        Ok(())
101    }
102
103    /// Construct a new counter pre-filled with the given MultiLineStrings
104    pub fn from_multi_line_strings<'a>(
105        geoms: impl Iterator<Item = Option<&'a (impl MultiLineStringTrait + 'a)>>,
106    ) -> Self {
107        let mut counter = Self::new_empty();
108        for maybe_multi_line_string in geoms.into_iter() {
109            counter.add_multi_line_string(maybe_multi_line_string);
110        }
111        counter
112    }
113
114    /// Construct a new counter pre-filled with the given geometries
115    pub fn from_geometries<'a>(
116        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait + 'a)>>,
117    ) -> GeoArrowResult<Self> {
118        let mut counter = Self::new_empty();
119        for g in geoms.into_iter() {
120            counter.add_geometry(g)?;
121        }
122        Ok(counter)
123    }
124
125    /// The number of bytes an array with this capacity would occupy.
126    pub fn num_bytes(&self) -> usize {
127        let offsets_byte_width = 4;
128        let num_offsets = self.geom_capacity + self.ring_capacity;
129        (offsets_byte_width * num_offsets) + (self.coord_capacity * 2 * 8)
130    }
131}
132
133impl Default for MultiLineStringCapacity {
134    fn default() -> Self {
135        Self::new_empty()
136    }
137}
138
139impl Add for MultiLineStringCapacity {
140    type Output = Self;
141
142    fn add(self, rhs: Self) -> Self::Output {
143        let coord_capacity = self.coord_capacity + rhs.coord_capacity;
144        let ring_capacity = self.ring_capacity + rhs.ring_capacity;
145        let geom_capacity = self.geom_capacity + rhs.geom_capacity;
146        Self::new(coord_capacity, ring_capacity, geom_capacity)
147    }
148}
149
150impl AddAssign for MultiLineStringCapacity {
151    fn add_assign(&mut self, rhs: Self) {
152        self.coord_capacity += rhs.coord_capacity;
153        self.ring_capacity += rhs.ring_capacity;
154        self.geom_capacity += rhs.geom_capacity;
155    }
156}
157
158impl AddAssign<LineStringCapacity> for MultiLineStringCapacity {
159    fn add_assign(&mut self, rhs: LineStringCapacity) {
160        self.coord_capacity += rhs.coord_capacity();
161        self.ring_capacity += rhs.geom_capacity();
162        self.geom_capacity += rhs.geom_capacity();
163    }
164}
geoarrow_array/capacity/multilinestring.rs

geoarrow_array/capacity/
multilinestring.rs
