Panic on out-of-bounds take and Interleave indices

vortex-array/src/arrays/chunked/compute/take.rs

@@ -7,9 +7,9 @@ use vortex_buffer::Buffer;
 use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure_eq;
 use vortex_error::vortex_err;
+use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 
 use crate::ArrayRef;
@@ -96,14 +96,11 @@ fn take_chunked_via_sort(
 
     if let Some(&(index, _)) = pairs.last() {
         let index = usize::try_from(index)?;
-        if index >= array.len() {
-            vortex_bail!(
-                "index {} out of bounds from {} to {}",
-                index,
-                0,
-                array.len()
-            );
-        }
+        let len = array.len();
+        assert!(
+            index < len,
+            "take index {index} out of bounds for length {len}"
+        );
     }
 
     let chunk_offsets = array.chunk_offset_values();
@@ -213,14 +210,11 @@ fn take_chunked(
         }
 
         let index = usize::try_from(index)?;
-        if index >= array.len() {
-            vortex_bail!(
-                "index {} out of bounds from {} to {}",
-                index,
-                0,
-                array.len()
-            );
-        }
+        let len = array.len();
+        assert!(
+            index < len,
+            "take index {index} out of bounds for length {len}"
+        );
 
         let still_sorted = last_index.is_none_or(|last_index| index >= last_index);
         let chunk_idx = if monotonic && still_sorted {
@@ -322,13 +316,10 @@ fn take_chunked_sorted(
         }
         cursor = range_end;
     }
-    if cursor != indices_values.len() {
-        vortex_bail!(
-            "index {} out of bounds from {} to {}",
-            usize::try_from(indices_values[cursor])?,
-            0,
-            array.as_ref().len()
-        );
+    // The chunk loop stops at the first index that no chunk holds.
+    if let Some(&index) = indices_values.get(cursor) {
+        let len = array.as_ref().len();
+        vortex_panic!("take index {index} out of bounds for length {len}");
     }
 
     if chunks.len() == 1 {
@@ -446,14 +437,11 @@ fn take_piecewise_chunked(
         let end = start
             .checked_add(length)
             .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?;
-        if end > array_len {
-            vortex_bail!(
-                "index {} out of bounds from {} to {}",
-                end - 1,
-                0,
-                array_len
-            );
-        }
+        assert!(
+            end <= array_len,
+            "take index {} out of bounds for length {array_len}",
+            end - 1
+        );
         total_len = total_len
             .checked_add(length)
             .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?;
@@ -579,9 +567,11 @@ fn contiguous_runs(starts: Vec<u64>, lengths: Vec<u64>, total: usize) -> ArrayRe
 mod tests {
     use std::sync::LazyLock;
 
+    use rstest::rstest;
     use vortex_buffer::Buffer;
     use vortex_buffer::bitbuffer;
     use vortex_buffer::buffer;
+    use vortex_error::VortexExpect;
     use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
@@ -743,17 +733,16 @@ mod tests {
     }
 
     #[test]
-    fn test_take_piecewise_out_of_bounds() -> VortexResult<()> {
+    #[should_panic(expected = "take index 16 out of bounds for length 15")]
+    fn test_take_piecewise_out_of_bounds() {
         let mut ctx = SESSION.create_execution_ctx();
-        let arr = chunked_i32()?;
-
-        let indices = contiguous_pieces(&[12], &[5])?;
-        let result = arr
-            .take(indices)
-            .and_then(|taken| taken.execute::<Canonical>(&mut ctx));
+        let arr = chunked_i32().vortex_expect("the chunks have the dtype of the array");
+        let indices = contiguous_pieces(&[12], &[5]).vortex_expect("the pieces are valid");
 
-        assert!(result.is_err());
-        Ok(())
+        drop(
+            arr.take(indices)
+                .and_then(|taken| taken.execute::<Canonical>(&mut ctx)),
+        );
     }
 
     #[test]
@@ -1351,19 +1340,19 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn test_take_out_of_bounds() -> VortexResult<()> {
+    /// Sorted indices take the sorted path, and unsorted indices take the generic path.
+    #[rstest]
+    #[case::sorted(buffer![0u64, 100])]
+    #[case::unsorted(buffer![100u64, 0])]
+    #[should_panic(expected = "take index 100 out of bounds for length 15")]
+    fn test_take_out_of_bounds(#[case] indices: Buffer<u64>) {
         let mut ctx = SESSION.create_execution_ctx();
-        let arr = chunked_i32()?;
+        let arr = chunked_i32().vortex_expect("the chunks have the dtype of the array");
 
-        // Sorted indices hit the fast path, unsorted the generic path; both must fail.
-        for indices in [buffer![0u64, 100], buffer![100u64, 0]] {
-            let result = arr
-                .take(indices.into_array())
-                .and_then(|taken| taken.execute::<Canonical>(&mut ctx));
-            assert!(result.is_err(), "expected out-of-bounds error");
-        }
-        Ok(())
+        drop(
+            arr.take(indices.into_array())
+                .and_then(|taken| taken.execute::<Canonical>(&mut ctx)),
+        );
     }
 
     #[test]

vortex-array/src/arrays/interleave/execute/bool.rs

@@ -8,8 +8,6 @@ use vortex_buffer::BitBuffer;
 use vortex_buffer::BitBufferMut;
 use vortex_buffer::BufferAllocatorRef;
 use vortex_error::VortexResult;
-use vortex_error::vortex_ensure;
-use vortex_error::vortex_ensure_eq;
 
 use super::super::Interleave;
 use super::super::InterleaveArrayExt;
@@ -59,7 +57,7 @@ pub(super) fn execute(
                 array_indices.as_slice::<A>(),
                 row_indices.as_slice::<R>(),
                 ctx.allocator(),
-            )?
+            )
         })
     });
 
@@ -76,45 +74,53 @@ fn gather<A: AsPrimitive<usize>, R: AsPrimitive<usize>>(
     branches: &[A],
     rows: &[R],
     allocator: &BufferAllocatorRef,
-) -> VortexResult<BitBufferMut> {
-    let len = validate_selectors(value_bits, branches, rows)?;
+) -> BitBufferMut {
+    let len = validate_selectors(value_bits, branches, rows);
 
     // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every
     // `i < len` that `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`.
-    Ok(unsafe { gather_bits(len, value_bits, branches, rows, allocator) })
+    unsafe { gather_bits(len, value_bits, branches, rows, allocator) }
 }
 
 /// Validates the per-row selector bounds, returning the output length (`branches.len()`).
 ///
-/// On success, `rows.len() == branches.len() == len` and, for every `i < len`,
-/// `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()` — exactly the
-/// preconditions of [`gather_bits`]. Errors (rather than panics) on any out-of-bounds selector.
+/// On return, `rows.len() == branches.len() == len` and, for every `i < len`,
+/// `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`. These are
+/// exactly the preconditions of [`gather_bits`].
+///
+/// # Panics
+///
+/// Panics if the selectors have different lengths, or if a selector names a value or a row that
+/// does not exist.
 fn validate_selectors<A: AsPrimitive<usize>, R: AsPrimitive<usize>>(
     value_bits: &[BitBuffer],
     branches: &[A],
     rows: &[R],
-) -> VortexResult<usize> {
+) -> usize {
     // The two selectors are validated to equal length at construction, which is the output length.
     let len = branches.len();
-    vortex_ensure_eq!(
+    assert_eq!(
         rows.len(),
         len,
         "interleave row_indices length does not match array_indices length",
     );
 
     for i in 0..len {
         let branch = branches[i].as_();
-        vortex_ensure!(
+        assert!(
             branch < value_bits.len(),
-            "interleave array index out of bounds"
+            "interleave array index {branch} out of bounds for {} values",
+            value_bits.len()
         );
-        vortex_ensure!(
-            rows[i].as_() < value_bits[branch].len(),
-            "interleave row index out of bounds"
+        let row = rows[i].as_();
+        assert!(
+            row < value_bits[branch].len(),
+            "interleave row index {row} out of bounds for length {}",
+            value_bits[branch].len()
         );
     }
 
-    Ok(len)
+    len
 }
 
 /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per

vortex-array/src/arrays/interleave/execute/decimal.rs

@@ -61,7 +61,7 @@ fn gather_values<T: NativeDecimalType>(
         .map(|i| prepare_source::<T>(array.value(i), dtype, ctx))
         .collect::<VortexResult<Vec<_>>>()?;
 
-    gather(array, &sources, ctx.allocator())
+    Ok(gather(array, &sources, ctx.allocator()))
 }
 
 /// Normalizes one decimal source while retaining a single payload for a constant.

vortex-array/src/arrays/interleave/execute/fixed_width.rs

@@ -11,10 +11,7 @@ use num_traits::AsPrimitive;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
-use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
-use vortex_error::vortex_ensure;
-use vortex_error::vortex_ensure_eq;
+use vortex_error::vortex_panic;
 
 use super::super::Interleave;
 use super::super::InterleaveArrayExt;
@@ -61,15 +58,15 @@ impl<T: Copy> Source<T> {
 /// `sources` holds one [`Source`] for each value of `array`, in order. Both selectors of `array`
 /// **must** be [`Primitive`] arrays, which the kernels ensure with `require_child!`.
 ///
-/// # Errors
+/// # Panics
 ///
-/// Returns an error if the selectors have different lengths, or if a selector names a value or a
-/// row that does not exist.
+/// Panics if the selectors have different lengths, or if a selector names a value or a row that
+/// does not exist.
 pub(super) fn gather<T: Copy>(
     array: &Array<Interleave>,
     sources: &[Source<T>],
     allocator: &BufferAllocatorRef,
-) -> VortexResult<Buffer<T>> {
+) -> Buffer<T> {
     let branches = array.array_indices().as_::<Primitive>();
     let rows = array.row_indices().as_::<Primitive>();
 
@@ -90,14 +87,14 @@ fn gather_rows<T, A, R>(
     branches: &[A],
     rows: &[R],
     allocator: &BufferAllocatorRef,
-) -> VortexResult<Buffer<T>>
+) -> Buffer<T>
 where
     T: Copy,
     A: AsPrimitive<usize>,
     R: AsPrimitive<usize>,
 {
     // `zip` truncates to the shorter input.
-    vortex_ensure_eq!(
+    assert_eq!(
         rows.len(),
         branches.len(),
         "interleave row_indices length does not match array_indices length",
@@ -110,48 +107,61 @@ where
     let mut output = BufferMut::with_capacity_in(len, allocator.clone());
     let slots = &mut output.spare_capacity_mut()[..len];
     for ((slot, branch), row) in slots.iter_mut().zip(branches).zip(rows) {
-        let Some(source) = sources.get((*branch).as_()) else {
-            vortex_bail!("interleave array index out of bounds");
+        let branch = (*branch).as_();
+        let Some(source) = sources.get(branch) else {
+            vortex_panic!(
+                "interleave array index {branch} out of bounds for {} values",
+                sources.len()
+            );
         };
+        // NB: `row & row_mask` is always in bounds for a constant, so this check is the only one
+        // that rejects a row past the end of a constant.
         let row = (*row).as_();
-        vortex_ensure!(row < source.len, "interleave row index out of bounds");
+        assert!(
+            row < source.len,
+            "interleave row index {row} out of bounds for length {}",
+            source.len
+        );
         slot.write(source.data[row & source.row_mask]);
     }
-    // SAFETY: the loop returns on any error, so reaching here means it wrote all `len` slots:
+    // SAFETY: the loop panics on a bad selector, so reaching here means it wrote all `len` slots:
     // `slots`, `branches`, and `rows` all have length `len`.
     unsafe { output.set_len(len) };
-    Ok(output.freeze())
+    output.freeze()
 }
 
 #[cfg(test)]
 mod tests {
     use rstest::rstest;
     use vortex_buffer::buffer;
-    use vortex_error::VortexResult;
 
     use super::Source;
     use super::gather_rows;
     use crate::memory::BufferAllocatorRef;
 
-    #[rstest]
-    #[case::array_index(2, 0, "array index out of bounds")]
-    #[case::constant_row_index(0, 1, "row index out of bounds")]
-    #[case::decoded_row_index(1, 1, "row index out of bounds")]
-    fn rejects_out_of_bounds_selectors(
-        #[case] branch: u8,
-        #[case] row: u8,
-        #[case] message: &str,
-    ) -> VortexResult<()> {
+    /// Gathers row `row` of value `branch`. Value 0 is a one-row constant, and value 1 is a one-row
+    /// decoded buffer.
+    fn gather_one(branch: u8, row: u8) {
         let allocator = BufferAllocatorRef::statically_allocated();
         let sources = [
             Source::constant(1u32, 1, &allocator),
             Source::decoded(buffer![2u32]),
         ];
 
-        let error = gather_rows(&sources, &[branch], &[row], &allocator)
-            .expect_err("every case names a value or a row that does not exist");
-        assert!(error.to_string().contains(message), "{error}");
+        drop(gather_rows(&sources, &[branch], &[row], &allocator));
+    }
 
-        Ok(())
+    #[test]
+    #[should_panic(expected = "interleave array index 2 out of bounds for 2 values")]
+    fn panics_on_out_of_bounds_array_index() {
+        gather_one(2, 0);
+    }
+
+    #[rstest]
+    #[case::constant(0)]
+    #[case::decoded(1)]
+    #[should_panic(expected = "interleave row index 1 out of bounds for length 1")]
+    fn panics_on_out_of_bounds_row_index(#[case] branch: u8) {
+        gather_one(branch, 1);
     }
 }