geoarrow_array/capacity/
multipoint.rs

1use std::ops::{Add, AddAssign};
2
3use geo_traits::{GeometryTrait, GeometryType, MultiPointTrait, PointTrait};
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 [`MultiPointArray`][crate::array::MultiPointArray].
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 MultiPointCapacity {
14    pub(crate) coord_capacity: usize,
15    pub(crate) geom_capacity: usize,
16}
17
18impl MultiPointCapacity {
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 the capacity of a point
38    #[inline]
39    pub fn add_point(&mut self, point: Option<&impl PointTrait>) {
40        self.geom_capacity += 1;
41        if let Some(point) = point {
42            self.add_valid_point(point)
43        }
44    }
45
46    #[inline]
47    fn add_valid_point(&mut self, _point: &impl PointTrait) {
48        self.coord_capacity += 1;
49    }
50
51    /// Add the capacity of the given MultiPoint
52    #[inline]
53    pub fn add_multi_point(&mut self, maybe_multi_point: Option<&impl MultiPointTrait>) {
54        self.geom_capacity += 1;
55
56        if let Some(multi_point) = maybe_multi_point {
57            self.add_valid_multi_point(multi_point);
58        }
59    }
60
61    #[inline]
62    fn add_valid_multi_point(&mut self, multi_point: &impl MultiPointTrait) {
63        self.coord_capacity += multi_point.num_points();
64    }
65
66    /// Add the capacity of the given Geometry
67    ///
68    /// The type of the geometry must be either Point or MultiPoint
69    #[inline]
70    pub fn add_geometry(&mut self, value: Option<&impl GeometryTrait>) -> GeoArrowResult<()> {
71        self.geom_capacity += 1;
72
73        if let Some(g) = value {
74            match g.as_type() {
75                GeometryType::Point(p) => self.add_valid_point(p),
76                GeometryType::MultiPoint(p) => self.add_valid_multi_point(p),
77                gt => {
78                    return Err(GeoArrowError::IncorrectGeometryType(format!(
79                        "Expected Point or MultiPoint, got {}",
80                        gt.name()
81                    )));
82                }
83            }
84        };
85        Ok(())
86    }
87
88    /// The coordinate buffer capacity
89    pub fn coord_capacity(&self) -> usize {
90        self.coord_capacity
91    }
92
93    /// The geometry offsets buffer capacity
94    pub fn geom_capacity(&self) -> usize {
95        self.geom_capacity
96    }
97
98    /// Construct a new counter pre-filled with the given MultiPoints
99    pub fn from_multi_points<'a>(
100        geoms: impl Iterator<Item = Option<&'a (impl MultiPointTrait + 'a)>>,
101    ) -> Self {
102        let mut counter = Self::new_empty();
103
104        for maybe_line_string in geoms.into_iter() {
105            counter.add_multi_point(maybe_line_string);
106        }
107
108        counter
109    }
110
111    /// Construct a new counter pre-filled with the given geometries
112    pub fn from_geometries<'a>(
113        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait + 'a)>>,
114    ) -> GeoArrowResult<Self> {
115        let mut counter = Self::new_empty();
116        for g in geoms.into_iter() {
117            counter.add_geometry(g)?;
118        }
119        Ok(counter)
120    }
121
122    /// The number of bytes an array with this capacity would occupy.
123    pub fn num_bytes(&self, dim: Dimension) -> usize {
124        let offsets_byte_width = 4;
125        let num_offsets = self.geom_capacity;
126        (offsets_byte_width * num_offsets) + (self.coord_capacity * dim.size() * 8)
127    }
128}
129
130impl Default for MultiPointCapacity {
131    fn default() -> Self {
132        Self::new_empty()
133    }
134}
135
136impl Add for MultiPointCapacity {
137    type Output = Self;
138
139    fn add(self, rhs: Self) -> Self::Output {
140        let coord_capacity = self.coord_capacity + rhs.coord_capacity;
141        let geom_capacity = self.geom_capacity + rhs.geom_capacity;
142        Self::new(coord_capacity, geom_capacity)
143    }
144}
145
146impl AddAssign for MultiPointCapacity {
147    fn add_assign(&mut self, rhs: Self) {
148        self.coord_capacity += rhs.coord_capacity;
149        self.geom_capacity += rhs.geom_capacity;
150    }
151}