geoarrow_array/capacity/
polygon.rs

1use std::ops::Add;
2
3use geo_traits::{GeometryTrait, GeometryType, LineStringTrait, PolygonTrait, RectTrait};
4use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
5
6use crate::util::GeometryTypeName;
7
8/// A counter for the buffer sizes of a [`PolygonArray`][crate::array::PolygonArray].
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 PolygonCapacity {
13    pub(crate) coord_capacity: usize,
14    pub(crate) ring_capacity: usize,
15    pub(crate) geom_capacity: usize,
16}
17
18impl PolygonCapacity {
19    /// Create a new capacity with known sizes.
20    pub fn new(coord_capacity: usize, ring_capacity: usize, geom_capacity: usize) -> Self {
21        Self {
22            coord_capacity,
23            ring_capacity,
24            geom_capacity,
25        }
26    }
27
28    /// Create a new empty capacity.
29    pub fn new_empty() -> Self {
30        Self::new(0, 0, 0)
31    }
32
33    /// Return `true` if the capacity is empty.
34    pub fn is_empty(&self) -> bool {
35        self.coord_capacity == 0 && self.ring_capacity == 0 && self.geom_capacity == 0
36    }
37
38    /// The coordinate buffer capacity
39    pub fn coord_capacity(&self) -> usize {
40        self.coord_capacity
41    }
42
43    /// The ring offset buffer capacity
44    pub fn ring_capacity(&self) -> usize {
45        self.ring_capacity
46    }
47
48    /// The geometry offset buffer capacity
49    pub fn geom_capacity(&self) -> usize {
50        self.geom_capacity
51    }
52
53    /// Add the capacity of the given Polygon
54    #[inline]
55    pub fn add_polygon<'a>(&mut self, polygon: Option<&'a (impl PolygonTrait + 'a)>) {
56        self.geom_capacity += 1;
57        if let Some(polygon) = polygon {
58            // Total number of rings in this polygon
59            let num_interiors = polygon.num_interiors();
60            self.ring_capacity += num_interiors + 1;
61
62            // Number of coords for each ring
63            if let Some(exterior) = polygon.exterior() {
64                self.coord_capacity += exterior.num_coords();
65            }
66
67            for int_ring in polygon.interiors() {
68                self.coord_capacity += int_ring.num_coords();
69            }
70        }
71    }
72
73    /// Add the capacity of the given Rect
74    #[inline]
75    pub fn add_rect<'a>(&mut self, rect: Option<&'a (impl RectTrait + 'a)>) {
76        self.geom_capacity += 1;
77        if rect.is_some() {
78            // A rect is a simple polygon with only one ring
79            self.ring_capacity += 1;
80            // A rect is a closed polygon with 5 coordinates
81            self.coord_capacity += 5;
82        }
83    }
84
85    /// Add the capacity of the given Geometry
86    ///
87    /// The type of the geometry must be either Polygon or Rect
88    #[inline]
89    pub fn add_geometry(&mut self, value: Option<&impl GeometryTrait>) -> GeoArrowResult<()> {
90        if let Some(geom) = value {
91            match geom.as_type() {
92                GeometryType::Polygon(g) => self.add_polygon(Some(g)),
93                GeometryType::Rect(g) => self.add_rect(Some(g)),
94                gt => {
95                    return Err(GeoArrowError::IncorrectGeometryType(format!(
96                        "Expected polygon, got {}",
97                        gt.name()
98                    )));
99                }
100            }
101        } else {
102            self.geom_capacity += 1;
103        };
104        Ok(())
105    }
106
107    /// Construct a new counter pre-filled with the given Polygons
108    pub fn from_polygons<'a>(
109        geoms: impl Iterator<Item = Option<&'a (impl PolygonTrait + 'a)>>,
110    ) -> Self {
111        let mut counter = Self::new_empty();
112        for maybe_polygon in geoms.into_iter() {
113            counter.add_polygon(maybe_polygon);
114        }
115        counter
116    }
117
118    /// Construct a new counter pre-filled with the given Rects
119    pub fn from_rects<'a>(geoms: impl Iterator<Item = Option<&'a (impl RectTrait + 'a)>>) -> Self {
120        let mut counter = Self::new_empty();
121        for maybe_rect in geoms.into_iter() {
122            counter.add_rect(maybe_rect);
123        }
124        counter
125    }
126
127    /// Construct a new counter pre-filled with the given geometries
128    pub fn from_geometries<'a>(
129        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait + 'a)>>,
130    ) -> GeoArrowResult<Self> {
131        let mut counter = Self::new_empty();
132        for g in geoms.into_iter() {
133            counter.add_geometry(g)?;
134        }
135        Ok(counter)
136    }
137
138    /// The number of bytes an array with this capacity would occupy.
139    pub fn num_bytes(&self) -> usize {
140        let offsets_byte_width = 4;
141        let num_offsets = self.geom_capacity + self.ring_capacity;
142        (offsets_byte_width * num_offsets) + (self.coord_capacity * 2 * 8)
143    }
144}
145
146impl Default for PolygonCapacity {
147    fn default() -> Self {
148        Self::new_empty()
149    }
150}
151
152impl Add for PolygonCapacity {
153    type Output = Self;
154
155    fn add(self, rhs: Self) -> Self::Output {
156        let coord_capacity = self.coord_capacity + rhs.coord_capacity;
157        let ring_capacity = self.ring_capacity + rhs.ring_capacity;
158        let geom_capacity = self.geom_capacity + rhs.geom_capacity;
159        Self::new(coord_capacity, ring_capacity, geom_capacity)
160    }
161}
geoarrow_array/capacity/polygon.rs

geoarrow_array/capacity/
polygon.rs
