geoarrow_array/capacity/
wkb.rs

1use std::ops::Add;
2
3use arrow_array::OffsetSizeTrait;
4use geo_traits::GeometryTrait;
5use wkb::writer::geometry_wkb_size;
6
7/// A counter for the buffer sizes of a [`GenericWkbArray`][crate::array::GenericWkbArray].
8///
9/// This can be used to reduce allocations by allocating once for exactly the array size you need.
10#[derive(Debug, Clone, Copy)]
11pub struct WkbCapacity {
12    pub(crate) buffer_capacity: usize,
13    pub(crate) offsets_capacity: usize,
14}
15
16impl WkbCapacity {
17    /// Create a new capacity with known sizes.
18    pub fn new(buffer_capacity: usize, offsets_capacity: usize) -> Self {
19        Self {
20            buffer_capacity,
21            offsets_capacity,
22        }
23    }
24
25    /// Create a new empty capacity.
26    pub fn new_empty() -> Self {
27        Self::new(0, 0)
28    }
29
30    /// Return `true` if the capacity is empty.
31    pub fn is_empty(&self) -> bool {
32        self.buffer_capacity == 0 && self.offsets_capacity == 0
33    }
34
35    /// The capacity of the underlying data buffer
36    pub fn buffer_capacity(&self) -> usize {
37        self.buffer_capacity
38    }
39
40    /// The capacity of the underlying offsets buffer
41    pub fn offsets_capacity(&self) -> usize {
42        self.offsets_capacity
43    }
44
45    /// Add a Geometry to this capacity counter.
46    #[inline]
47    pub fn add_geometry<'a>(&mut self, geom: Option<&'a (impl GeometryTrait<T = f64> + 'a)>) {
48        if let Some(geom) = geom {
49            self.buffer_capacity += geometry_wkb_size(geom);
50        }
51        self.offsets_capacity += 1;
52    }
53
54    /// Create a capacity counter from an iterator of Geometries.
55    pub fn from_geometries<'a>(
56        geoms: impl Iterator<Item = Option<&'a (impl GeometryTrait<T = f64> + 'a)>>,
57    ) -> Self {
58        let mut counter = Self::new_empty();
59        for maybe_geom in geoms.into_iter() {
60            counter.add_geometry(maybe_geom);
61        }
62        counter
63    }
64
65    /// The number of bytes an array with this capacity would occupy.
66    pub fn num_bytes<O: OffsetSizeTrait>(&self) -> usize {
67        let offsets_byte_width = if O::IS_LARGE { 8 } else { 4 };
68        let num_offsets = self.offsets_capacity;
69        (offsets_byte_width * num_offsets) + self.buffer_capacity
70    }
71}
72
73impl Default for WkbCapacity {
74    fn default() -> Self {
75        Self::new_empty()
76    }
77}
78
79impl Add for WkbCapacity {
80    type Output = Self;
81
82    fn add(self, rhs: Self) -> Self::Output {
83        let buffer_capacity = self.buffer_capacity + rhs.buffer_capacity;
84        let offsets_capacity = self.offsets_capacity + rhs.offsets_capacity;
85
86        Self::new(buffer_capacity, offsets_capacity)
87    }
88}
geoarrow_array/capacity/wkb.rs

geoarrow_array/capacity/
wkb.rs
