geoarrow_array/capacity/
multilinestring.rs

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

geoarrow_array/capacity/
multilinestring.rs
