geoarrow_array/capacity/
linestring.rs

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

geoarrow_array/capacity/
linestring.rs
