geoarrow_array/capacity/
multipoint.rs

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