This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new d282f80763 feat/chore: introduce fallible alternatives for
MutableBuffer (#10317)
d282f80763 is described below
commit d282f807637acf262b04463b60cd3ea3b80fbf80
Author: RIchard Baah <[email protected]>
AuthorDate: Tue Aug 4 19:59:22 2026 -0400
feat/chore: introduce fallible alternatives for MutableBuffer (#10317)
# Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax.
-->
- works towards #9843.
- a follow PR will be made to wire up the fallible API into other
portions of the code base.
# Rationale for this change
see #9843.
`MutableBuffer` is a low-level building block used throughout the Arrow
implementation. Every growth operation, `reserve`, `resize`, `push`,
`extend_from_slice`, etc, can panic on arithmetic overflow or an invalid
allocation layout. In library code that needs to surface errors
gracefully. This PR adds methods to return errors instead of panicking.
I made new methods instead of adding conditional checks/wrappers for
each method because its easier to slowly roll these changes out to call
sites instead of swapping.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
# What changes are included in this PR?
```
┌──────────────────────────┬──────────────────────┐
│ Fallible │ Infallible │
├──────────────────────────┼──────────────────────┤
│ try_reserve │ reserve │
├──────────────────────────┼──────────────────────┤
│ try_repeat_slice_n_times │ repeat_slice_n_times │
├──────────────────────────┼──────────────────────┤
│ try_reallocate (private) │ reallocate │
├──────────────────────────┼──────────────────────┤
│ try_resize │ resize │
├──────────────────────────┼──────────────────────┤
│ try_shrink_to_fit │ shrink_to_fit │
├──────────────────────────┼──────────────────────┤
│ try_extend_from_slice │ extend_from_slice │
├──────────────────────────┼──────────────────────┤
│ try_push │ push │
├──────────────────────────┼──────────────────────┤
│ try_extend_zeros │ extend_zeros │
└──────────────────────────┴──────────────────────┘
```
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
# Are these changes tested?
Currently these methods aren't wired up to anything in the code base so
I didn't include any test. Id be happy to include some. The
implementations of the fallible API's are 1-1 to their non-fallible
counter parts so i'm not sure if this would be needed.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
If this PR claims a performance improvement, please include evidence
such as benchmark results.
-->
# Are there any user-facing changes?
yes, users can now work with errors instead of having their programs
panic.
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please call them out.
-->
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-buffer/src/buffer/mutable.rs | 322 +++++++++++++++++++++++--------------
1 file changed, 204 insertions(+), 118 deletions(-)
diff --git a/arrow-buffer/src/buffer/mutable.rs
b/arrow-buffer/src/buffer/mutable.rs
index b6e6a70c6c..b30a932b35 100644
--- a/arrow-buffer/src/buffer/mutable.rs
+++ b/arrow-buffer/src/buffer/mutable.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use std::alloc::{Layout, handle_alloc_error};
+use std::alloc::Layout;
use std::mem;
use std::ptr::NonNull;
@@ -33,6 +33,31 @@ use std::sync::Mutex;
use super::Buffer;
+/// Error returned by fallible [`MutableBuffer`] operations.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum MutableBufferError {
+ /// Arithmetic overflow when computing the required buffer length or
capacity.
+ LengthOverflow,
+ /// The requested capacity cannot be represented as a valid allocation
layout.
+ LayoutError,
+ /// An allocation failed due to insufficient memory.
+ AllocationError(Layout),
+}
+
+impl std::fmt::Display for MutableBufferError {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ match self {
+ Self::LengthOverflow => write!(f, "buffer length overflow"),
+ Self::LayoutError => write!(f, "invalid allocation layout for
requested capacity"),
+ Self::AllocationError(layout) => {
+ write!(f, "failed to allocate memory for layout {layout:?}")
+ }
+ }
+ }
+}
+
+impl std::error::Error for MutableBufferError {}
+
/// A [`MutableBuffer`] is a wrapper over memory regions, used to build
/// [`Buffer`]s out of items or slices of items.
///
@@ -118,7 +143,7 @@ impl MutableBuffer {
/// See [`MutableBuffer::with_capacity`].
#[inline]
pub fn new(capacity: usize) -> Self {
- Self::with_capacity(capacity)
+ Self::try_with_capacity(capacity).unwrap_or_else(|e| panic!("{e}"))
}
/// Allocate a new [MutableBuffer] with initial capacity to be at least
`capacity`.
@@ -129,24 +154,35 @@ impl MutableBuffer {
/// then `isize::MAX`, then this function will panic.
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
- let capacity = bit_util::round_upto_multiple_of_64(capacity);
+ Self::try_with_capacity(capacity).unwrap_or_else(|e| panic!("{e}"))
+ }
+
+ /// Fallible version of [`MutableBuffer::with_capacity`].
+ #[inline]
+ pub fn try_with_capacity(capacity: usize) -> Result<Self,
MutableBufferError> {
+ let capacity = capacity
+ .checked_next_multiple_of(64)
+ .ok_or(MutableBufferError::LayoutError)?;
let layout = Layout::from_size_align(capacity, ALIGNMENT)
- .expect("failed to create layout for MutableBuffer");
+ .map_err(|_| MutableBufferError::LayoutError)?;
let data = match layout.size() {
0 => dangling_ptr(),
_ => {
// Safety: Verified size != 0
let raw_ptr = unsafe { std::alloc::alloc(layout) };
- NonNull::new(raw_ptr).unwrap_or_else(||
handle_alloc_error(layout))
+ match NonNull::new(raw_ptr) {
+ Some(ptr) => ptr,
+ None => return
Err(MutableBufferError::AllocationError(layout)),
+ }
}
};
- Self {
+ Ok(Self {
data,
len: 0,
layout,
#[cfg(feature = "pool")]
reservation: std::sync::Mutex::new(None),
- }
+ })
}
/// Allocates a new [MutableBuffer] with `len` and capacity to be at least
`len` where
@@ -165,22 +201,31 @@ impl MutableBuffer {
///
/// Panics if `len` is too large to construct a valid allocation [`Layout`]
pub fn from_len_zeroed(len: usize) -> Self {
- let layout = Layout::from_size_align(len, ALIGNMENT).unwrap();
+ Self::try_from_len_zeroed(len).unwrap_or_else(|e| panic!("{e}"))
+ }
+
+ /// Fallible version of [`MutableBuffer::from_len_zeroed`].
+ pub fn try_from_len_zeroed(len: usize) -> Result<Self, MutableBufferError>
{
+ let layout =
+ Layout::from_size_align(len, ALIGNMENT).map_err(|_|
MutableBufferError::LayoutError)?;
let data = match layout.size() {
0 => dangling_ptr(),
_ => {
// Safety: Verified size != 0
let raw_ptr = unsafe { std::alloc::alloc_zeroed(layout) };
- NonNull::new(raw_ptr).unwrap_or_else(||
handle_alloc_error(layout))
+ match NonNull::new(raw_ptr) {
+ Some(ptr) => ptr,
+ None => return
Err(MutableBufferError::AllocationError(layout)),
+ }
}
};
- Self {
+ Ok(Self {
data,
len,
layout,
#[cfg(feature = "pool")]
reservation: std::sync::Mutex::new(None),
- }
+ })
}
/// Allocates a new [MutableBuffer] from given `Bytes`.
@@ -259,6 +304,22 @@ impl MutableBuffer {
}
}
+ /// Fallible version of [`MutableBuffer::reserve`].
+ #[inline]
+ pub fn try_reserve(&mut self, additional: usize) -> Result<(),
MutableBufferError> {
+ let required_cap = self
+ .len
+ .checked_add(additional)
+ .ok_or(MutableBufferError::LengthOverflow)?;
+ if required_cap > self.layout.size() {
+ let new_capacity = required_cap
+ .checked_next_multiple_of(64)
+ .ok_or(MutableBufferError::LayoutError)?;
+ let new_capacity = std::cmp::max(new_capacity,
self.layout.size().saturating_mul(2));
+ self.try_reallocate(new_capacity)?;
+ }
+ Ok(())
+ }
/// Ensures that this buffer has at least `self.len + additional` bytes.
This re-allocates iff
/// `self.len + additional > capacity`.
/// # Example
@@ -279,108 +340,98 @@ impl MutableBuffer {
// exits.
#[inline(always)]
pub fn reserve(&mut self, additional: usize) {
- let required_cap = self
- .len
- .checked_add(additional)
- .expect("buffer length overflow");
- if required_cap > self.layout.size() {
- let new_capacity =
bit_util::round_upto_multiple_of_64(required_cap);
- let new_capacity = std::cmp::max(new_capacity, self.layout.size()
* 2);
- self.reallocate(new_capacity)
- }
+ self.try_reserve(additional)
+ .unwrap_or_else(|e| panic!("{e}"))
}
- /// Adding to this mutable buffer `slice_to_repeat` repeated
`repeat_count` times.
- ///
- /// # Example
- ///
- /// ## Repeat the same string bytes multiple times
- /// ```
- /// # use arrow_buffer::buffer::MutableBuffer;
- /// let mut buffer = MutableBuffer::new(0);
- /// let bytes_to_repeat = b"ab";
- /// buffer.repeat_slice_n_times(bytes_to_repeat, 3);
- /// assert_eq!(buffer.as_slice(), b"ababab");
- /// ```
- ///
- /// # Panics
- ///
- /// Panics if the repeated slice byte length overflows `usize`, if the
resulting buffer
- /// length overflows `usize`, or if reserving the required capacity fails
for the same
- /// reasons as [`MutableBuffer::reserve`].
- pub fn repeat_slice_n_times<T: ArrowNativeType>(
+ /// Fallible version of [`MutableBuffer::repeat_slice_n_times`].
+ pub fn try_repeat_slice_n_times<T: ArrowNativeType>(
&mut self,
slice_to_repeat: &[T],
repeat_count: usize,
- ) {
+ ) -> Result<(), MutableBufferError> {
if repeat_count == 0 || slice_to_repeat.is_empty() {
- return;
+ return Ok(());
}
-
- let bytes_to_repeat = size_of_val(slice_to_repeat);
- let repeated_bytes = repeat_count
- .checked_mul(bytes_to_repeat)
- .expect("repeated slice byte length overflow");
+ let bytes_per_copy = size_of_val(slice_to_repeat);
+ let total_bytes = repeat_count
+ .checked_mul(bytes_per_copy)
+ .ok_or(MutableBufferError::LengthOverflow)?;
self.len
- .checked_add(repeated_bytes)
- .expect("mutable buffer length overflow");
+ .checked_add(total_bytes)
+ .ok_or(MutableBufferError::LengthOverflow)?;
// Ensure capacity
- self.reserve(repeated_bytes);
+ self.try_reserve(total_bytes)?;
// Save the length before we do all the copies to know where to start
from
let length_before = self.len;
// Copy the initial slice once so we can use doubling strategy on it
- self.extend_from_slice(slice_to_repeat);
-
- // This tracks how much bytes we have added by repeating so far
- let added_repeats_length = bytes_to_repeat;
- assert_eq!(
- self.len - length_before,
- added_repeats_length,
- "should copy exactly the same number of bytes"
- );
+ self.try_extend_from_slice(slice_to_repeat)?;
// Number of times the slice was repeated
- let mut already_repeated_times = 1;
+ let mut already_repeated = 1usize;
// We will use doubling strategy to fill the buffer in
log(repeat_count) steps
- while already_repeated_times < repeat_count {
+ while already_repeated < repeat_count {
// How many slices can we copy in this iteration
// (either double what we have, or just the remaining ones)
- let number_of_slices_to_copy =
- already_repeated_times.min(repeat_count -
already_repeated_times);
- let number_of_bytes_to_copy = number_of_slices_to_copy *
bytes_to_repeat;
-
+ let to_copy = already_repeated.min(repeat_count -
already_repeated);
+ let byte_count = to_copy * bytes_per_copy;
unsafe {
// Get to the start of the data before we started copying
anything
let src = self.data.as_ptr().add(length_before) as *const u8;
-
// Go to the current location to copy to (end of current data)
let dst = self.data.as_ptr().add(self.len);
-
- // SAFETY: the pointers are not overlapping as there is
`number_of_bytes_to_copy` or less between them
- std::ptr::copy_nonoverlapping(src, dst,
number_of_bytes_to_copy)
+ // SAFETY: the pointers are not overlapping as there is
`byte_count` or less between them
+ std::ptr::copy_nonoverlapping(src, dst, byte_count);
}
-
// Advance the length by the amount of data we just copied
(doubled)
- self.len += number_of_bytes_to_copy;
-
- already_repeated_times += number_of_slices_to_copy;
+ self.len += byte_count;
+ already_repeated += to_copy;
}
+ Ok(())
+ }
+ /// Adding to this mutable buffer `slice_to_repeat` repeated
`repeat_count` times.
+ ///
+ /// # Example
+ ///
+ /// ## Repeat the same string bytes multiple times
+ /// ```
+ /// # use arrow_buffer::buffer::MutableBuffer;
+ /// let mut buffer = MutableBuffer::new(0);
+ /// let bytes_to_repeat = b"ab";
+ /// buffer.repeat_slice_n_times(bytes_to_repeat, 3);
+ /// assert_eq!(buffer.as_slice(), b"ababab");
+ /// ```
+ ///
+ /// # Panics
+ ///
+ /// Panics if the repeated slice byte length overflows `usize`, if the
resulting buffer
+ /// length overflows `usize`, or if reserving the required capacity fails
for the same
+ /// reasons as [`MutableBuffer::reserve`].
+ pub fn repeat_slice_n_times<T: ArrowNativeType>(
+ &mut self,
+ slice_to_repeat: &[T],
+ repeat_count: usize,
+ ) {
+ self.try_repeat_slice_n_times(slice_to_repeat, repeat_count)
+ .unwrap_or_else(|e| panic!("{e}"))
}
#[cold]
- fn reallocate(&mut self, capacity: usize) {
- let new_layout = Layout::from_size_align(capacity,
self.layout.align()).unwrap();
+ fn try_reallocate(&mut self, capacity: usize) -> Result<(),
MutableBufferError> {
+ let new_layout = Layout::from_size_align(capacity, self.layout.align())
+ .map_err(|_| MutableBufferError::LayoutError)?;
+
if new_layout.size() == 0 {
if self.layout.size() != 0 {
// Safety: data was allocated with layout
unsafe { std::alloc::dealloc(self.as_mut_ptr(), self.layout) };
- self.layout = new_layout
+ self.layout = new_layout;
}
- return;
+ return Ok(());
}
let data = match self.layout.size() {
@@ -389,7 +440,10 @@ impl MutableBuffer {
// Safety: verified new layout is valid and not empty
_ => unsafe { std::alloc::realloc(self.as_mut_ptr(), self.layout,
capacity) },
};
- self.data = NonNull::new(data).unwrap_or_else(||
handle_alloc_error(new_layout));
+ self.data = match NonNull::new(data) {
+ Some(ptr) => ptr,
+ None => return
Err(MutableBufferError::AllocationError(new_layout)),
+ };
self.layout = new_layout;
#[cfg(feature = "pool")]
{
@@ -397,8 +451,8 @@ impl MutableBuffer {
reservation.resize(self.layout.size());
}
}
+ Ok(())
}
-
/// Truncates this buffer to `len` bytes
///
/// If `len` is greater than the buffer's current length, this has no
effect
@@ -416,6 +470,26 @@ impl MutableBuffer {
}
}
+ /// Fallible version of [`MutableBuffer::resize`].
+ #[inline]
+ pub fn try_resize(&mut self, new_len: usize, value: u8) -> Result<(),
MutableBufferError> {
+ if new_len > self.len {
+ let diff = new_len - self.len;
+ self.try_reserve(diff)?;
+ // Safety: try_reserve ensured capacity >= new_len.
+ // write the value
+ unsafe { self.data.as_ptr().add(self.len).write_bytes(value, diff)
};
+ }
+ // this truncates the buffer when new_len < self.len
+ self.len = new_len;
+ #[cfg(feature = "pool")]
+ {
+ if let Some(reservation) =
self.reservation.lock().unwrap().as_mut() {
+ reservation.resize(self.len);
+ }
+ }
+ Ok(())
+ }
/// Resizes the buffer, either truncating its contents (with no change in
capacity), or
/// growing it (potentially reallocating it) and writing `value` in the
newly available bytes.
/// # Example
@@ -434,22 +508,21 @@ impl MutableBuffer {
// exits.
#[inline(always)]
pub fn resize(&mut self, new_len: usize, value: u8) {
- if new_len > self.len {
- let diff = new_len - self.len;
- self.reserve(diff);
- // write the value
- unsafe { self.data.as_ptr().add(self.len).write_bytes(value, diff)
};
- }
- // this truncates the buffer when new_len < self.len
- self.len = new_len;
- #[cfg(feature = "pool")]
- {
- if let Some(reservation) =
self.reservation.lock().unwrap().as_mut() {
- reservation.resize(self.len);
- }
- }
+ self.try_resize(new_len, value)
+ .unwrap_or_else(|e| panic!("{e}"))
}
+ /// Fallible version of [`MutableBuffer::shrink_to_fit`].
+ pub fn try_shrink_to_fit(&mut self) -> Result<(), MutableBufferError> {
+ let new_capacity = self
+ .len
+ .checked_next_multiple_of(64)
+ .ok_or(MutableBufferError::LayoutError)?;
+ if new_capacity < self.layout.size() {
+ self.try_reallocate(new_capacity)?;
+ }
+ Ok(())
+ }
/// Shrinks the capacity of the buffer as much as possible.
/// The new capacity will aligned to the nearest 64 bit alignment.
///
@@ -471,10 +544,7 @@ impl MutableBuffer {
/// Panics if the current length is too large to round up to the next
64-byte boundary and
/// construct a valid allocation layout.
pub fn shrink_to_fit(&mut self) {
- let new_capacity = bit_util::round_upto_multiple_of_64(self.len);
- if new_capacity < self.layout.size() {
- self.reallocate(new_capacity)
- }
+ self.try_shrink_to_fit().unwrap_or_else(|e| panic!("{e}"))
}
/// Returns whether this buffer is empty or not.
@@ -575,6 +645,25 @@ impl MutableBuffer {
offsets
}
+ /// Fallible version of [`MutableBuffer::extend_from_slice`].
+ #[inline]
+ pub fn try_extend_from_slice<T: ArrowNativeType>(
+ &mut self,
+ items: &[T],
+ ) -> Result<(), MutableBufferError> {
+ let additional = mem::size_of_val(items);
+ self.try_reserve(additional)?;
+ unsafe {
+ // this assumes that `[ToByteSlice]` can be copied directly
+ // without calling `to_byte_slice` for each element,
+ // which is correct for all ArrowNativeType implementations.
+ let src = items.as_ptr() as *const u8;
+ let dst = self.data.as_ptr().add(self.len);
+ std::ptr::copy_nonoverlapping(src, dst, additional);
+ }
+ self.len += additional;
+ Ok(())
+ }
/// Extends this buffer from a slice of items that can be represented in
bytes, increasing its capacity if needed.
/// # Example
/// ```
@@ -590,17 +679,8 @@ impl MutableBuffer {
/// reasons as [`MutableBuffer::reserve`].
#[inline]
pub fn extend_from_slice<T: ArrowNativeType>(&mut self, items: &[T]) {
- let additional = mem::size_of_val(items);
- self.reserve(additional);
- unsafe {
- // this assumes that `[ToByteSlice]` can be copied directly
- // without calling `to_byte_slice` for each element,
- // which is correct for all ArrowNativeType implementations.
- let src = items.as_ptr() as *const u8;
- let dst = self.data.as_ptr().add(self.len);
- std::ptr::copy_nonoverlapping(src, dst, additional)
- }
- self.len += additional;
+ self.try_extend_from_slice(items)
+ .unwrap_or_else(|e| panic!("{e}"))
}
/// Extends the buffer with a new item, increasing its capacity if needed.
@@ -640,6 +720,15 @@ impl MutableBuffer {
self.len += additional;
}
+ /// Fallible version of [`MutableBuffer::extend_zeros`].
+ #[inline]
+ pub fn try_extend_zeros(&mut self, additional: usize) -> Result<(),
MutableBufferError> {
+ let new_len = self
+ .len
+ .checked_add(additional)
+ .ok_or(MutableBufferError::LengthOverflow)?;
+ self.try_resize(new_len, 0)
+ }
/// Extends the buffer by `additional` bytes equal to `0u8`, incrementing
its capacity if needed.
///
/// # Panics
@@ -648,11 +737,8 @@ impl MutableBuffer {
/// reserving a capacity that fails for the same reasons as
[`MutableBuffer::reserve`].
#[inline]
pub fn extend_zeros(&mut self, additional: usize) {
- let new_len = self
- .len
- .checked_add(additional)
- .expect("buffer length overflow");
- self.resize(new_len, 0);
+ self.try_extend_zeros(additional)
+ .unwrap_or_else(|e| panic!("{e}"))
}
/// # Safety
@@ -1424,7 +1510,7 @@ mod tests {
}
#[test]
- #[should_panic(expected = "failed to create layout for MutableBuffer:
LayoutError")]
+ #[should_panic(expected = "invalid allocation layout for requested
capacity")]
fn test_with_capacity_panics_above_max_capacity() {
let max_capacity = isize::MAX as usize - (isize::MAX as usize %
ALIGNMENT);
let _ = MutableBuffer::with_capacity(max_capacity + 1);
@@ -1446,14 +1532,14 @@ mod tests {
assert_eq!(pool.used(), 128);
// Reallocate to a larger size
- buffer.reallocate(200);
+ buffer.try_reallocate(200).unwrap();
// The capacity is exactly the requested size, not rounded up
assert_eq!(buffer.capacity(), 200);
assert_eq!(pool.used(), 200);
// Reallocate to a smaller size
- buffer.reallocate(50);
+ buffer.try_reallocate(50).unwrap();
// The capacity is exactly the requested size, not rounded up
assert_eq!(buffer.capacity(), 50);
@@ -1576,14 +1662,14 @@ mod tests {
}
#[test]
- #[should_panic(expected = "repeated slice byte length overflow")]
+ #[should_panic(expected = "buffer length overflow")]
fn test_repeat_slice_count_multiply_overflow() {
let mut buffer = MutableBuffer::new(0);
buffer.repeat_slice_n_times(&[0_u64], usize::MAX /
mem::size_of::<u64>() + 1);
}
#[test]
- #[should_panic(expected = "mutable buffer length overflow")]
+ #[should_panic(expected = "buffer length overflow")]
fn test_repeat_slice_count_len_overflow() {
let mut buffer = MutableBuffer::new(0);
buffer.push(0_u8);
@@ -1671,7 +1757,7 @@ mod tests {
}
#[test]
- #[should_panic(expected = "failed to round upto multiple of 64")]
+ #[should_panic(expected = "invalid allocation layout for requested
capacity")]
fn test_mutable_new_capacity_overflow() {
// Tests overflow during initial allocation
let _ = MutableBuffer::new(usize::MAX - 10);
