feat(arrow-buffer): add OverflowError and fallible offset constructors

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);
     }

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<Self, OverflowError> {
+        let capacity = self
+            .buffer
+            .len()
+            .checked_mul(count)
+            .ok_or_else(|| OverflowError::new::<usize>("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

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<O: ArrowNativeType> OffsetBuffer<O> {
     ///
     /// # Panics
     ///
-    /// Panics if `(len + 1) * size_of::<O>()` overflows `usize`
+    /// Panics if `(len + 1) * size_of::<O>()` 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::<O>()` overflows `usize`
+    pub fn try_new_zeroed(len: usize) -> Result<Self, OverflowError> {
         let len_bytes = len
             .checked_add(1)
             .and_then(|o| o.checked_mul(std::mem::size_of::<O>()))
-            .expect("overflow");
+            .ok_or_else(|| OverflowError::new::<usize>("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<O: ArrowNativeType> OffsetBuffer<O> {
     ///
     /// # Panics
     ///
-    /// Panics on overflow
+    /// Panics on overflow. Use [`Self::try_from_lengths`] for a fallible version.
     pub fn from_lengths<I>(lengths: I) -> Self
+    where
+        I: IntoIterator<Item = usize>,
+    {
+        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::<i32>::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<i32>`.
+    pub fn try_from_lengths<I>(lengths: I) -> Result<Self, OverflowError>
     where
         I: IntoIterator<Item = usize>,
     {
@@ -132,12 +161,14 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
 
         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::<usize>("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::<O>("offset").with_value(acc))?;
+        Ok(Self(out.into()))
     }
 
     /// Create a new [`OffsetBuffer`] where each slice has the same length
@@ -153,28 +184,45 @@ impl<O: ArrowNativeType> OffsetBuffer<O> {
     ///
     /// # 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::<i32>::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<Self, OverflowError> {
         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::<usize>("total length"))?;
 
-        // Check for overflow
-        O::from_usize(length * n).expect("offset overflow");
+        O::from_usize(total_length)
+            .ok_or_else(|| OverflowError::new::<O>("offset").with_value(total_length))?;
 
         let offsets = (0..=n)
             .map(|index| O::usize_as(index * length))
             .collect::<Vec<O>>();
 
-        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::<i32>::try_from_lengths(lengths).unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "offset overflow: 4294967296 does not fit in i32"
+        );
+        assert!(OffsetBuffer::<i64>::try_from_lengths(lengths).is_ok());
+
+        let err = OffsetBuffer::<i32>::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::<i32>::from_lengths(lengths)).is_err());
+    }
+
+    #[test]
+    fn try_from_repeated_length_overflow() {
+        assert_eq!(
+            OffsetBuffer::<i32>::try_from_repeated_length(2, usize::MAX)
+                .unwrap_err()
+                .to_string(),
+            "total length overflow: does not fit in usize"
+        );
+        assert_eq!(
+            OffsetBuffer::<i32>::try_from_repeated_length(u32::MAX as usize, 2)
+                .unwrap_err()
+                .to_string(),
+            "offset overflow: 8589934590 does not fit in i32"
+        );
+        assert_eq!(
+            OffsetBuffer::<i32>::try_from_repeated_length(4, 3)
+                .unwrap()
+                .as_ref(),
+            &[0, 4, 8, 12]
+        );
+    }
+
+    #[test]
+    fn try_new_zeroed_overflow() {
+        assert_eq!(
+            OffsetBuffer::<i64>::try_new_zeroed(usize::MAX)
+                .unwrap_err()
+                .to_string(),
+            "buffer length overflow: does not fit in usize"
+        );
+        assert_eq!(
+            OffsetBuffer::<i32>::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::<i32>::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::<i32>::from_repeated_length(usize::MAX, 2);
     }