diff --git a/arrow-array/src/array/byte_array.rs b/arrow-array/src/array/byte_array.rs index 74b9d0353bab..5901b8a16a08 100644 --- a/arrow-array/src/array/byte_array.rs +++ b/arrow-array/src/array/byte_array.rs @@ -711,13 +711,13 @@ mod tests { } #[test] - #[should_panic(expected = "usize overflow")] + #[should_panic(expected = "total length overflow: does not fit in usize")] fn create_repeated_usize_overflow_1() { let _arr = BinaryArray::new_repeated(b"hello", (usize::MAX / "hello".len()) + 1); } #[test] - #[should_panic(expected = "usize overflow")] + #[should_panic(expected = "total length overflow: does not fit in usize")] fn create_repeated_usize_overflow_2() { let _arr = BinaryArray::new_repeated(b"hello", usize::MAX); } diff --git a/arrow-buffer/src/buffer/null.rs b/arrow-buffer/src/buffer/null.rs index 3303d8c350ce..003585d91bff 100644 --- a/arrow-buffer/src/buffer/null.rs +++ b/arrow-buffer/src/buffer/null.rs @@ -17,7 +17,7 @@ use crate::bit_iterator::{BitIndexIterator, BitIterator, BitSliceIterator}; use crate::buffer::BooleanBuffer; -use crate::{Buffer, MutableBuffer}; +use crate::{Buffer, MutableBuffer, OverflowError}; /// A [`BooleanBuffer`] used to encode validity (null values) for Arrow arrays /// @@ -121,9 +121,25 @@ impl NullBuffer { /// /// # Panics /// - /// Panics if `self.len() * count` overflows `usize` + /// Panics if `self.len() * count` overflows `usize`. + /// Use [`Self::try_expand`] for a fallible version. pub fn expand(&self, count: usize) -> Self { - let capacity = self.buffer.len().checked_mul(count).unwrap(); + self.try_expand(count).unwrap_or_else(|err| panic!("{err}")) + } + + /// Returns a new [`NullBuffer`] where each bit in the current null buffer + /// is repeated `count` times. This is useful for masking the nulls of + /// the child of a FixedSizeListArray based on its parent + /// + /// # Errors + /// + /// Errors if `self.len() * count` overflows `usize` + pub fn try_expand(&self, count: usize) -> Result { + let capacity = self + .buffer + .len() + .checked_mul(count) + .ok_or_else(|| OverflowError::new::("buffer length"))?; let mut buffer = MutableBuffer::new_null(capacity); // Expand each bit within `null_mask` into `element_len` @@ -136,10 +152,10 @@ impl NullBuffer { crate::bit_util::set_bit(buffer.as_mut(), i * count + j) } } - Self { + Ok(Self { buffer: BooleanBuffer::new(buffer.into(), 0, capacity), null_count: self.null_count * count, - } + }) } /// Returns the length of this [`NullBuffer`] in bits diff --git a/arrow-buffer/src/buffer/offset.rs b/arrow-buffer/src/buffer/offset.rs index f1dcdfd5aae7..ff3340c29151 100644 --- a/arrow-buffer/src/buffer/offset.rs +++ b/arrow-buffer/src/buffer/offset.rs @@ -16,7 +16,7 @@ // under the License. use crate::buffer::ScalarBuffer; -use crate::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBufferBuilder}; +use crate::{ArrowNativeType, MutableBuffer, NullBuffer, OffsetBufferBuilder, OverflowError}; use std::ops::Deref; /// A non-empty buffer of monotonically increasing, positive integers. @@ -98,14 +98,24 @@ impl OffsetBuffer { /// /// # Panics /// - /// Panics if `(len + 1) * size_of::()` overflows `usize` + /// Panics if `(len + 1) * size_of::()` overflows `usize`. + /// Use [`Self::try_new_zeroed`] for a fallible version. pub fn new_zeroed(len: usize) -> Self { + Self::try_new_zeroed(len).unwrap_or_else(|err| panic!("{err}")) + } + + /// Create a new [`OffsetBuffer`] containing `len + 1` `0` values + /// + /// # Errors + /// + /// Errors if `(len + 1) * size_of::()` overflows `usize` + pub fn try_new_zeroed(len: usize) -> Result { let len_bytes = len .checked_add(1) .and_then(|o| o.checked_mul(std::mem::size_of::())) - .expect("overflow"); + .ok_or_else(|| OverflowError::new::("buffer length"))?; let buffer = MutableBuffer::from_len_zeroed(len_bytes); - Self(buffer.into_buffer().into()) + Ok(Self(buffer.into_buffer().into())) } /// Create a new [`OffsetBuffer`] from the iterator of slice lengths @@ -121,8 +131,27 @@ impl OffsetBuffer { /// /// # Panics /// - /// Panics on overflow + /// Panics on overflow. Use [`Self::try_from_lengths`] for a fallible version. pub fn from_lengths(lengths: I) -> Self + where + I: IntoIterator, + { + Self::try_from_lengths(lengths).unwrap_or_else(|err| panic!("{err}")) + } + + /// Create a new [`OffsetBuffer`] from the iterator of slice lengths + /// + /// ``` + /// # use arrow_buffer::OffsetBuffer; + /// let offsets = OffsetBuffer::::try_from_lengths([1, 3, 5]).unwrap(); + /// assert_eq!(offsets.as_ref(), &[0, 1, 4, 9]); + /// ``` + /// + /// # Errors + /// + /// Errors if the total length overflows `usize` or `O`, e.g. if the lengths + /// add up to more than `i32::MAX` for a `OffsetBuffer`. + pub fn try_from_lengths(lengths: I) -> Result where I: IntoIterator, { @@ -132,12 +161,14 @@ impl OffsetBuffer { let mut acc = 0_usize; for length in iter { - acc = acc.checked_add(length).expect("usize overflow"); + acc = acc + .checked_add(length) + .ok_or_else(|| OverflowError::new::("total length"))?; out.push(O::usize_as(acc)) } // Check for overflow - O::from_usize(acc).expect("offset overflow"); - Self(out.into()) + O::from_usize(acc).ok_or_else(|| OverflowError::new::("offset").with_value(acc))?; + Ok(Self(out.into())) } /// Create a new [`OffsetBuffer`] where each slice has the same length @@ -153,28 +184,45 @@ impl OffsetBuffer { /// /// # Panics /// - /// Panics on overflow + /// Panics on overflow. Use [`Self::try_from_repeated_length`] for a fallible version. pub fn from_repeated_length(length: usize, n: usize) -> Self { + Self::try_from_repeated_length(length, n).unwrap_or_else(|err| panic!("{err}")) + } + + /// Create a new [`OffsetBuffer`] where each slice has the same length + /// `length`, repeated `n` times. + /// + /// ``` + /// # use arrow_buffer::OffsetBuffer; + /// let offsets = OffsetBuffer::::try_from_repeated_length(4, 3).unwrap(); + /// assert_eq!(offsets.as_ref(), &[0, 4, 8, 12]); + /// ``` + /// + /// # Errors + /// + /// Errors if `length * n` overflows `usize` or `O`. + pub fn try_from_repeated_length(length: usize, n: usize) -> Result { if n == 0 { - return Self::new_empty(); + return Ok(Self::new_empty()); } if length == 0 { - return Self::new_zeroed(n); + return Self::try_new_zeroed(n); } - // Check for overflow // Making sure we don't overflow usize or O when calculating the total length - length.checked_mul(n).expect("usize overflow"); + let total_length = length + .checked_mul(n) + .ok_or_else(|| OverflowError::new::("total length"))?; - // Check for overflow - O::from_usize(length * n).expect("offset overflow"); + O::from_usize(total_length) + .ok_or_else(|| OverflowError::new::("offset").with_value(total_length))?; let offsets = (0..=n) .map(|index| O::usize_as(index * length)) .collect::>(); - Self(ScalarBuffer::from(offsets)) + Ok(Self(ScalarBuffer::from(offsets))) } /// Get an Iterator over the lengths of this [`OffsetBuffer`] @@ -445,6 +493,64 @@ mod tests { OffsetBuffer::new(vec![-1, 0, 1].into()); } + #[test] + fn try_from_lengths_overflow() { + // Fits in an i64, but not in an i32: + let lengths = [u32::MAX as usize, 1]; + + let err = OffsetBuffer::::try_from_lengths(lengths).unwrap_err(); + assert_eq!( + err.to_string(), + "offset overflow: 4294967296 does not fit in i32" + ); + assert!(OffsetBuffer::::try_from_lengths(lengths).is_ok()); + + let err = OffsetBuffer::::try_from_lengths([usize::MAX, 1]).unwrap_err(); + assert_eq!( + err.to_string(), + "total length overflow: does not fit in usize" + ); + + // The panicking version agrees: + assert!(std::panic::catch_unwind(|| OffsetBuffer::::from_lengths(lengths)).is_err()); + } + + #[test] + fn try_from_repeated_length_overflow() { + assert_eq!( + OffsetBuffer::::try_from_repeated_length(2, usize::MAX) + .unwrap_err() + .to_string(), + "total length overflow: does not fit in usize" + ); + assert_eq!( + OffsetBuffer::::try_from_repeated_length(u32::MAX as usize, 2) + .unwrap_err() + .to_string(), + "offset overflow: 8589934590 does not fit in i32" + ); + assert_eq!( + OffsetBuffer::::try_from_repeated_length(4, 3) + .unwrap() + .as_ref(), + &[0, 4, 8, 12] + ); + } + + #[test] + fn try_new_zeroed_overflow() { + assert_eq!( + OffsetBuffer::::try_new_zeroed(usize::MAX) + .unwrap_err() + .to_string(), + "buffer length overflow: does not fit in usize" + ); + assert_eq!( + OffsetBuffer::::try_new_zeroed(3).unwrap().as_ref(), + &[0; 4] + ); + } + #[test] fn offsets() { OffsetBuffer::new(vec![0, 1, 2, 3].into()); @@ -485,7 +591,7 @@ mod tests { } #[test] - #[should_panic(expected = "usize overflow")] + #[should_panic(expected = "total length overflow: does not fit in usize")] fn from_lengths_usize_overflow() { OffsetBuffer::::from_lengths([usize::MAX, 1]); } @@ -509,7 +615,7 @@ mod tests { } #[test] - #[should_panic(expected = "usize overflow")] + #[should_panic(expected = "total length overflow: does not fit in usize")] fn from_repeated_lengths_usize_length_usize_overflow() { OffsetBuffer::::from_repeated_length(usize::MAX, 2); } diff --git a/arrow-buffer/src/builder/offset.rs b/arrow-buffer/src/builder/offset.rs index a51ca5f01d76..ee5ba02ed991 100644 --- a/arrow-buffer/src/builder/offset.rs +++ b/arrow-buffer/src/builder/offset.rs @@ -17,7 +17,7 @@ use std::ops::Deref; -use crate::{ArrowNativeType, OffsetBuffer}; +use crate::{ArrowNativeType, OffsetBuffer, OverflowError}; /// Builder of [`OffsetBuffer`] #[derive(Debug)] @@ -41,11 +41,27 @@ impl OffsetBufferBuilder { /// /// # Panics /// - /// Panics if adding `length` would overflow `usize` + /// Panics if adding `length` would overflow `usize`. + /// Use [`Self::try_push_length`] for a fallible version. #[inline] pub fn push_length(&mut self, length: usize) { - self.last_offset = self.last_offset.checked_add(length).expect("overflow"); - self.offsets.push(O::usize_as(self.last_offset)) + self.try_push_length(length) + .unwrap_or_else(|err| panic!("{err}")) + } + + /// Push a slice of `length` bytes + /// + /// # Errors + /// + /// Errors if adding `length` would overflow `usize`. The builder is left unchanged. + #[inline] + pub fn try_push_length(&mut self, length: usize) -> Result<(), OverflowError> { + self.last_offset = self + .last_offset + .checked_add(length) + .ok_or_else(|| OverflowError::new::("total length"))?; + self.offsets.push(O::usize_as(self.last_offset)); + Ok(()) } /// Reserve space for at least `additional` further offsets @@ -58,23 +74,44 @@ impl OffsetBufferBuilder { /// /// # Panics /// - /// Panics if offsets overflow `O` + /// Panics if offsets overflow `O`. Use [`Self::try_finish`] for a fallible version. pub fn finish(self) -> OffsetBuffer { - O::from_usize(self.last_offset).expect("overflow"); - unsafe { OffsetBuffer::new_unchecked(self.offsets.into()) } + self.try_finish().unwrap_or_else(|err| panic!("{err}")) + } + + /// Takes the builder itself and returns an [`OffsetBuffer`] + /// + /// # Errors + /// + /// Errors if offsets overflow `O`, e.g. if they add up to more than `i32::MAX` + /// for a `OffsetBufferBuilder`. + pub fn try_finish(self) -> Result, OverflowError> { + O::from_usize(self.last_offset) + .ok_or_else(|| OverflowError::new::("offset").with_value(self.last_offset))?; + Ok(unsafe { OffsetBuffer::new_unchecked(self.offsets.into()) }) } /// Builds the [OffsetBuffer] without resetting the builder. /// /// # Panics /// - /// Panics if offsets overflow `O` + /// Panics if offsets overflow `O`. Use [`Self::try_finish_cloned`] for a fallible version. pub fn finish_cloned(&self) -> OffsetBuffer { + self.try_finish_cloned() + .unwrap_or_else(|err| panic!("{err}")) + } + + /// Builds the [OffsetBuffer] without resetting the builder. + /// + /// # Errors + /// + /// Errors for the same reasons as [`Self::try_finish`]. + pub fn try_finish_cloned(&self) -> Result, OverflowError> { let cloned = Self { offsets: self.offsets.clone(), last_offset: self.last_offset, }; - cloned.finish() + cloned.try_finish() } } @@ -88,6 +125,27 @@ impl Deref for OffsetBufferBuilder { #[cfg(test)] mod tests { + + #[test] + fn try_finish_overflow() { + let mut builder = OffsetBufferBuilder::::new(2); + builder.try_push_length(u32::MAX as usize).unwrap(); + let expected = "offset overflow: 4294967295 does not fit in i32"; + assert_eq!( + builder.try_finish_cloned().unwrap_err().to_string(), + expected + ); + assert_eq!(builder.try_finish().unwrap_err().to_string(), expected); + + let mut builder = OffsetBufferBuilder::::new(2); + builder.try_push_length(usize::MAX).unwrap(); + // The builder is unchanged by a failed push: + assert_eq!( + builder.try_push_length(1).unwrap_err().to_string(), + "total length overflow: does not fit in usize" + ); + assert_eq!(builder.len(), 2); + } use crate::OffsetBufferBuilder; #[test] diff --git a/arrow-buffer/src/error.rs b/arrow-buffer/src/error.rs new file mode 100644 index 000000000000..ad769848f51f --- /dev/null +++ b/arrow-buffer/src/error.rs @@ -0,0 +1,89 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Errors returned by `arrow-buffer`. + +/// An arithmetic overflow. +/// +/// Returned by the `try_` alternatives to the functions that panic on overflow, +/// for instance [`OffsetBuffer::try_from_lengths`](crate::OffsetBuffer::try_from_lengths). +/// +/// ``` +/// # use arrow_buffer::OffsetBuffer; +/// // 32 bit offsets cannot describe more than 2 GiB of data: +/// let err = OffsetBuffer::::try_from_lengths([u32::MAX as usize]).unwrap_err(); +/// assert_eq!(err.to_string(), "offset overflow: 4294967295 does not fit in i32"); +/// +/// // 64 bit offsets can: +/// assert!(OffsetBuffer::::try_from_lengths([u32::MAX as usize]).is_ok()); +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct OverflowError { + what: &'static str, + value: Option, + type_name: &'static str, +} + +impl OverflowError { + /// `what` names what overflowed, for instance `"offset"`, + /// and `T` is the type it did not fit in, for instance `i32`. + pub fn new(what: &'static str) -> Self { + Self { + what, + value: None, + type_name: std::any::type_name::(), + } + } + + /// The value that did not fit. + pub const fn with_value(mut self, value: usize) -> Self { + self.value = Some(value); + self + } + + /// What overflowed, for instance `"offset"`. + pub const fn what(&self) -> &'static str { + self.what + } + + /// The value that did not fit, if known. + pub const fn value(&self) -> Option { + self.value + } + + /// The name of the type that the value did not fit in, for instance `"i32"`. + pub const fn type_name(&self) -> &'static str { + self.type_name + } +} + +impl std::fmt::Display for OverflowError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let Self { + what, + value, + type_name, + } = self; + write!(f, "{what} overflow: ")?; + match value { + Some(value) => write!(f, "{value} does not fit in {type_name}"), + None => write!(f, "does not fit in {type_name}"), + } + } +} + +impl std::error::Error for OverflowError {} diff --git a/arrow-buffer/src/lib.rs b/arrow-buffer/src/lib.rs index 230747b8b84a..3e4fc828f63c 100644 --- a/arrow-buffer/src/lib.rs +++ b/arrow-buffer/src/lib.rs @@ -46,6 +46,9 @@ pub use buffer::*; pub mod builder; pub use builder::*; +mod error; +pub use error::OverflowError; + mod bigint; pub use bigint::i256;