geoarrow_array/capacity/
polygon.rs

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

geoarrow_array/capacity/
polygon.rs
