Clarify decimal Interleave source preparation

Prepare each decimal source as one typed gather input, with an early return for constants. Keep buffer normalization separate from selector dispatch. Parameterize bounds-check cases and report assertion failures at their callers.

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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

@@ -16,6 +16,7 @@ use super::super::InterleaveSlots;
 use super::fixed_width::Source;
 use super::fixed_width::gather;
 use crate::AnyColumnar;
+use crate::ArrayRef;
 use crate::array::Array;
 use crate::arrays::Constant;
 use crate::arrays::Decimal;
@@ -59,40 +60,7 @@ fn gather_values<T: NativeDecimalType>(
 ) -> VortexResult<Buffer<T>> {
     let allocator = ctx.allocator().clone();
     let values = (0..array.num_values())
-        .map(|i| {
-            let value = array.value(i);
-            let len = value.len();
-            let (data, row_mask) = if let Some(constant) = value.as_opt::<Constant>() {
-                let payload = match constant.scalar().as_decimal().decimal_value() {
-                    Some(value) => value.cast::<T>().ok_or_else(|| {
-                        vortex_err!("interleave decimal constant must fit {dtype}, got {value}")
-                    })?,
-                    None => T::default(),
-                };
-                (Buffer::full_in(payload, 1, allocator.clone()), 0)
-            } else {
-                let decimal = value.as_::<Decimal>();
-                let data = if decimal.values_type() == T::DECIMAL_TYPE {
-                    decimal.buffer::<T>()
-                } else {
-                    // A decimal source can have noncanonical storage. Normalize before indexing
-                    // rows so that every source has the same physical width. A narrowing copy
-                    // checks the precision and zeros arbitrary null payloads.
-                    let mask = decimal.validity()?.execute_mask(len, ctx)?;
-                    let copy = match_each_decimal_value_type!(decimal.values_type(), |S| {
-                        copy_to_canonical_width(decimal.buffer::<S>().as_slice(), &mask, dtype)?
-                    });
-                    Buffer::<T>::from_byte_buffer(copy.unwrap_host())
-                };
-                (data, usize::MAX)
-            };
-
-            Ok(Source {
-                data,
-                len,
-                row_mask,
-            })
-        })
+        .map(|i| prepare_source::<T>(array.value(i), dtype, ctx))
         .collect::<VortexResult<Vec<_>>>()?;
     let branches = array.array_indices().as_::<Primitive>();
     let rows = array.row_indices().as_::<Primitive>();
@@ -109,5 +77,47 @@ fn gather_values<T: NativeDecimalType>(
     })
 }
 
+/// Normalizes one decimal source while retaining a single payload for a constant.
+fn prepare_source<T: NativeDecimalType>(
+    value: &ArrayRef,
+    dtype: DecimalDType,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Source<T>> {
+    let len = value.len();
+    if let Some(constant) = value.as_opt::<Constant>() {
+        let payload = match constant.scalar().as_decimal().decimal_value() {
+            Some(value) => value.cast::<T>().ok_or_else(|| {
+                vortex_err!("interleave decimal constant must fit {dtype}, got {value}")
+            })?,
+            None => T::default(),
+        };
+
+        return Ok(Source {
+            data: Buffer::full_in(payload, 1, ctx.allocator().clone()),
+            len,
+            row_mask: 0,
+        });
+    }
+
+    let decimal = value.as_::<Decimal>();
+    let data = if decimal.values_type() == T::DECIMAL_TYPE {
+        decimal.buffer::<T>()
+    } else {
+        // Decimal sources can still have noncanonical storage. A narrowing copy must check the
+        // precision and zero null slots before the gather reads the buffer.
+        let mask = decimal.validity()?.execute_mask(len, ctx)?;
+        let copy = match_each_decimal_value_type!(decimal.values_type(), |S| {
+            copy_to_canonical_width(decimal.buffer::<S>().as_slice(), &mask, dtype)?
+        });
+        Buffer::<T>::from_byte_buffer(copy.unwrap_host())
+    };
+
+    Ok(Source {
+        data,
+        len,
+        row_mask: usize::MAX,
+    })
+}
+
 #[cfg(test)]
 mod tests;

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

@@ -30,7 +30,12 @@ use crate::scalar::Scalar;
 use crate::validity::Validity;
 
 /// Executes the gather, checks its physical width and dtype, and compares it with a canonical oracle.
-fn check(values: Vec<ArrayRef>, indices: &[(u8, u64)], expected: DecimalArray) -> VortexResult<()> {
+#[track_caller]
+fn assert_interleaved(
+    values: Vec<ArrayRef>,
+    indices: &[(u8, u64)],
+    expected: DecimalArray,
+) -> VortexResult<()> {
     let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch));
     let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row));
     let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?;
@@ -87,7 +92,7 @@ fn test_decimal_interleave(
             dtype,
             output_validity,
         );
-        check(
+        assert_interleaved(
             vec![first.into_array(), second.into_array()],
             &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)],
             expected,
@@ -133,7 +138,7 @@ fn test_decimal_interleave_constants(
                 Validity::NonNullable
             },
         );
-        check(
+        assert_interleaved(
             vec![constant, column],
             &[(0, 2), (1, 0), (0, 1), (0, 2)],
             expected,
@@ -167,7 +172,7 @@ fn test_decimal_interleave_mixed_stored_widths(#[case] precision: u8) -> VortexR
             dtype,
             Validity::from_iter([true, true, false, true, false]),
         );
-        check(
+        assert_interleaved(
             vec![first, second],
             &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)],
             expected,
@@ -196,15 +201,15 @@ fn test_decimal_interleave_without_valid_values(#[case] indices: &[(u8, u64)]) -
         dtype,
         Validity::AllInvalid,
     );
-    check(vec![first, second], indices, expected)
+    assert_interleaved(vec![first, second], indices, expected)
 }
 
 #[test]
 fn test_decimal_interleave_empty_sources() -> VortexResult<()> {
     let dtype = DecimalDType::new(9, 0);
     let empty = DecimalArray::new(Buffer::<i8>::empty(), dtype, Validity::NonNullable).into_array();
     let expected = DecimalArray::new(Buffer::<i32>::empty(), dtype, Validity::NonNullable);
-    check(vec![empty.clone(), empty], &[], expected)
+    assert_interleaved(vec![empty.clone(), empty], &[], expected)
 }
 
 #[test]

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

@@ -65,27 +65,32 @@ where
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use rstest::rstest;
+    use vortex_buffer::Buffer;
+    use vortex_error::VortexResult;
+
+    use super::Source;
+    use super::gather;
     use crate::memory::BufferAllocatorRef;
 
-    #[test]
-    fn rejects_out_of_bounds_selectors() {
+    #[rstest]
+    #[case::array_index(1, 0, "array index out of bounds")]
+    #[case::row_index(0, 1, "row index out of bounds")]
+    fn rejects_out_of_bounds_selectors(
+        #[case] branch: u8,
+        #[case] row: u8,
+        #[case] message: &str,
+    ) -> VortexResult<()> {
         let values = [Source {
             data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()),
             len: 1,
             row_mask: 0,
         }];
 
         let allocator = BufferAllocatorRef::statically_allocated();
-        let error = gather(&values, &[1u8], &[0u8], &allocator).unwrap_err();
-        assert!(
-            error.to_string().contains("array index out of bounds"),
-            "{error}"
-        );
-        let error = gather(&values, &[0u8], &[1u8], &allocator).unwrap_err();
-        assert!(
-            error.to_string().contains("row index out of bounds"),
-            "{error}"
-        );
+        let error = gather(&values, &[branch], &[row], &allocator).unwrap_err();
+        assert!(error.to_string().contains(message), "{error}");
+
+        Ok(())
     }
 }