geoarrow_array/builder/
offsets.rs

1//! This was originally copied from arrow2.
2
3use std::hint::unreachable_unchecked;
4
5use arrow_array::OffsetSizeTrait;
6use arrow_buffer::OffsetBuffer;
7use geoarrow_schema::error::{GeoArrowError, GeoArrowResult};
8
9/// A wrapper type of [`Vec<O>`] representing the invariants of Arrow's offsets.
10/// It is guaranteed to (sound to assume that):
11/// * every element is `>= 0`
12/// * element at position `i` is >= than element at position `i-1`.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct OffsetsBuilder<O: OffsetSizeTrait>(Vec<O>);
15
16impl<O: OffsetSizeTrait> Default for OffsetsBuilder<O> {
17    #[inline]
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl<O: OffsetSizeTrait> OffsetsBuilder<O> {
24    /// Returns an empty [`OffsetsBuilder`] (i.e. with a single element, the zero)
25    #[inline]
26    pub(crate) fn new() -> Self {
27        Self(vec![O::zero()])
28    }
29
30    /// Returns a new [`OffsetsBuilder`] with a capacity, allocating at least `capacity + 1`
31    /// entries.
32    pub(crate) fn with_capacity(capacity: usize) -> Self {
33        let mut offsets = Vec::with_capacity(capacity + 1);
34        offsets.push(O::zero());
35        Self(offsets)
36    }
37
38    /// Reserves `additional` entries.
39    pub(crate) fn reserve(&mut self, additional: usize) {
40        self.0.reserve(additional);
41    }
42
43    /// Reserves exactly `additional` entries.
44    pub(crate) fn reserve_exact(&mut self, additional: usize) {
45        self.0.reserve_exact(additional);
46    }
47
48    /// Shrinks the capacity of self to fit.
49    pub(crate) fn shrink_to_fit(&mut self) {
50        self.0.shrink_to_fit();
51    }
52
53    /// Pushes a new element with a given length.
54    /// # Error
55    /// This function errors iff the new last item is larger than what `O` supports.
56    /// # Implementation
57    /// This function:
58    /// * checks that this length does not overflow
59    #[inline]
60    pub(crate) fn try_push_usize(&mut self, length: usize) -> GeoArrowResult<()> {
61        let length = O::usize_as(length);
62
63        let old_length = self.last();
64        // let new_length = old_length.checked_add(&length).ok_or(Error::Overflow)?;
65        let new_length = *old_length + length;
66        self.0.push(new_length);
67        Ok(())
68    }
69
70    /// Returns the last offset of this container.
71    #[inline]
72    pub(crate) fn last(&self) -> &O {
73        match self.0.last() {
74            Some(element) => element,
75            None => unsafe { unreachable_unchecked() },
76        }
77    }
78
79    /// Returns the length an array with these offsets would be.
80    #[inline]
81    pub(crate) fn len_proxy(&self) -> usize {
82        self.0.len() - 1
83    }
84
85    #[inline]
86    /// Returns the number of offsets in this container.
87    pub(crate) fn len(&self) -> usize {
88        self.0.len()
89    }
90
91    /// Returns the byte slice stored in this buffer
92    #[inline]
93    pub(crate) fn as_slice(&self) -> &[O] {
94        self.0.as_slice()
95    }
96
97    /// Extends itself with `additional` elements equal to the last offset.
98    /// This is useful to extend offsets with empty values, e.g. for null slots.
99    #[inline]
100    pub(crate) fn extend_constant(&mut self, additional: usize) {
101        let offset = *self.last();
102        if additional == 1 {
103            self.0.push(offset)
104        } else {
105            self.0.resize(self.len() + additional, offset)
106        }
107    }
108
109    pub(crate) fn finish(self) -> OffsetBuffer<O> {
110        OffsetBuffer::new(self.0.into())
111    }
112}
113
114impl From<OffsetsBuilder<i32>> for OffsetsBuilder<i64> {
115    fn from(offsets: OffsetsBuilder<i32>) -> Self {
116        // this conversion is lossless and uphelds all invariants
117        Self(
118            offsets
119                .as_slice()
120                .iter()
121                .map(|x| *x as i64)
122                .collect::<Vec<_>>(),
123        )
124    }
125}
126
127impl TryFrom<OffsetsBuilder<i64>> for OffsetsBuilder<i32> {
128    type Error = GeoArrowError;
129
130    fn try_from(offsets: OffsetsBuilder<i64>) -> GeoArrowResult<Self> {
131        i32::try_from(*offsets.last()).map_err(|_| GeoArrowError::Overflow)?;
132
133        // this conversion is lossless and uphelds all invariants
134        Ok(Self(
135            offsets
136                .as_slice()
137                .iter()
138                .map(|x| *x as i32)
139                .collect::<Vec<_>>(),
140        ))
141    }
142}
geoarrow_array/builder/offsets.rs

geoarrow_array/builder/
offsets.rs
