arrow_buffer/builder/
offset.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::ops::Deref;
19
20use crate::{ArrowNativeType, OffsetBuffer, OverflowError};
21
22/// Builder of [`OffsetBuffer`]
23#[derive(Debug)]
24pub struct OffsetBufferBuilder<O: ArrowNativeType> {
25    offsets: Vec<O>,
26    last_offset: usize,
27}
28
29impl<O: ArrowNativeType> OffsetBufferBuilder<O> {
30    /// Create a new builder with space for `capacity + 1` offsets
31    pub fn new(capacity: usize) -> Self {
32        let mut offsets = Vec::with_capacity(capacity + 1);
33        offsets.push(O::usize_as(0));
34        Self {
35            offsets,
36            last_offset: 0,
37        }
38    }
39
40    /// Push a slice of `length` bytes
41    ///
42    /// # Panics
43    ///
44    /// Panics if adding `length` would overflow `usize`.
45    /// Use [`Self::try_push_length`] for a fallible version.
46    #[inline]
47    pub fn push_length(&mut self, length: usize) {
48        self.try_push_length(length)
49            .unwrap_or_else(|err| panic!("{err}"))
50    }
51
52    /// Push a slice of `length` bytes
53    ///
54    /// # Errors
55    ///
56    /// Errors if adding `length` would overflow `usize`. The builder is left unchanged.
57    #[inline]
58    pub fn try_push_length(&mut self, length: usize) -> Result<(), OverflowError> {
59        self.last_offset = self
60            .last_offset
61            .checked_add(length)
62            .ok_or_else(|| OverflowError::new::<usize>("total length"))?;
63        self.offsets.push(O::usize_as(self.last_offset));
64        Ok(())
65    }
66
67    /// Reserve space for at least `additional` further offsets
68    #[inline]
69    pub fn reserve(&mut self, additional: usize) {
70        self.offsets.reserve(additional);
71    }
72
73    /// Takes the builder itself and returns an [`OffsetBuffer`]
74    ///
75    /// # Panics
76    ///
77    /// Panics if offsets overflow `O`. Use [`Self::try_finish`] for a fallible version.
78    pub fn finish(self) -> OffsetBuffer<O> {
79        self.try_finish().unwrap_or_else(|err| panic!("{err}"))
80    }
81
82    /// Takes the builder itself and returns an [`OffsetBuffer`]
83    ///
84    /// # Errors
85    ///
86    /// Errors if offsets overflow `O`, e.g. if they add up to more than `i32::MAX`
87    /// for a `OffsetBufferBuilder<i32>`.
88    pub fn try_finish(self) -> Result<OffsetBuffer<O>, OverflowError> {
89        O::from_usize(self.last_offset)
90            .ok_or_else(|| OverflowError::new::<O>("offset").with_value(self.last_offset))?;
91        Ok(unsafe { OffsetBuffer::new_unchecked(self.offsets.into()) })
92    }
93
94    /// Builds the [OffsetBuffer] without resetting the builder.
95    ///
96    /// # Panics
97    ///
98    /// Panics if offsets overflow `O`. Use [`Self::try_finish_cloned`] for a fallible version.
99    pub fn finish_cloned(&self) -> OffsetBuffer<O> {
100        self.try_finish_cloned()
101            .unwrap_or_else(|err| panic!("{err}"))
102    }
103
104    /// Builds the [OffsetBuffer] without resetting the builder.
105    ///
106    /// # Errors
107    ///
108    /// Errors for the same reasons as [`Self::try_finish`].
109    pub fn try_finish_cloned(&self) -> Result<OffsetBuffer<O>, OverflowError> {
110        let cloned = Self {
111            offsets: self.offsets.clone(),
112            last_offset: self.last_offset,
113        };
114        cloned.try_finish()
115    }
116}
117
118impl<O: ArrowNativeType> Deref for OffsetBufferBuilder<O> {
119    type Target = [O];
120
121    fn deref(&self) -> &Self::Target {
122        self.offsets.as_ref()
123    }
124}
125
126#[cfg(test)]
127mod tests {
128
129    #[test]
130    fn try_finish_overflow() {
131        let mut builder = OffsetBufferBuilder::<i32>::new(2);
132        builder.try_push_length(u32::MAX as usize).unwrap();
133        let expected = "offset overflow: 4294967295 does not fit in i32";
134        assert_eq!(
135            builder.try_finish_cloned().unwrap_err().to_string(),
136            expected
137        );
138        assert_eq!(builder.try_finish().unwrap_err().to_string(), expected);
139
140        let mut builder = OffsetBufferBuilder::<i32>::new(2);
141        builder.try_push_length(usize::MAX).unwrap();
142        // The builder is unchanged by a failed push:
143        assert_eq!(
144            builder.try_push_length(1).unwrap_err().to_string(),
145            "total length overflow: does not fit in usize"
146        );
147        assert_eq!(builder.len(), 2);
148    }
149    use crate::OffsetBufferBuilder;
150
151    #[test]
152    fn test_basic() {
153        let mut builder = OffsetBufferBuilder::<i32>::new(5);
154        assert_eq!(builder.len(), 1);
155        assert_eq!(&*builder, &[0]);
156        let finished = builder.finish_cloned();
157        assert_eq!(finished.len(), 1);
158        assert_eq!(&*finished, &[0]);
159
160        builder.push_length(2);
161        builder.push_length(6);
162        builder.push_length(0);
163        builder.push_length(13);
164
165        let finished = builder.finish();
166        assert_eq!(&*finished, &[0, 2, 8, 8, 21]);
167    }
168
169    #[test]
170    #[should_panic(expected = "overflow")]
171    fn test_usize_overflow() {
172        let mut builder = OffsetBufferBuilder::<i32>::new(5);
173        builder.push_length(1);
174        builder.push_length(usize::MAX);
175        builder.finish();
176    }
177
178    #[test]
179    #[should_panic(expected = "overflow")]
180    fn test_i32_overflow() {
181        let mut builder = OffsetBufferBuilder::<i32>::new(5);
182        builder.push_length(1);
183        builder.push_length(i32::MAX as usize);
184        builder.finish();
185    }
186}
arrow_buffer/builder/offset.rs

arrow_buffer/builder/
offset.rs
