fix[vortex-array]: fix overflow on FSL element take indices

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

@@ -17,6 +17,7 @@ use crate::arrays::dict::TakeExecute;
 use crate::dtype::IntegerPType;
 use crate::executor::ExecutionCtx;
 use crate::match_each_integer_ptype;
+use crate::match_smallest_offset_type;
 use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 
@@ -26,20 +27,24 @@ use crate::vtable::ValidityHelper;
 /// that elements start at offset 0 and be perfectly packed without gaps. We expand list indices
 /// into element indices and push them down to the child elements array.
 impl TakeExecute for FixedSizeList {
+    #[expect(clippy::cognitive_complexity)]
     fn take(
         array: &FixedSizeListArray,
         indices: &ArrayRef,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        let max_element_idx = array.elements().len();
         match_each_integer_ptype!(indices.dtype().as_ptype(), |I| {
-            take_with_indices::<I>(array, indices, ctx)
+            match_smallest_offset_type!(max_element_idx, |E| {
+                take_with_indices::<I, E>(array, indices, ctx)
+            })
         })
         .map(Some)
     }
 }
 
 /// Dispatches to the appropriate take implementation based on list size and nullability.
-fn take_with_indices<I: IntegerPType>(
+fn take_with_indices<I: IntegerPType, E: IntegerPType>(
     array: &FixedSizeListArray,
     indices: &ArrayRef,
     ctx: &mut ExecutionCtx,
@@ -74,15 +79,15 @@ fn take_with_indices<I: IntegerPType>(
     } else {
         // The result's nullability is the union of the input nullabilities.
         if array.dtype().is_nullable() || indices_array.dtype().is_nullable() {
-            take_nullable_fsl::<I>(array, &indices_array)
+            take_nullable_fsl::<I, E>(array, &indices_array)
         } else {
-            take_non_nullable_fsl::<I>(array, &indices_array)
+            take_non_nullable_fsl::<I, E>(array, &indices_array)
         }
     }
 }
 
 /// Takes from an array when both the array and indices are non-nullable.
-fn take_non_nullable_fsl<I: IntegerPType>(
+fn take_non_nullable_fsl<I: IntegerPType, E: IntegerPType>(
     array: &FixedSizeListArray,
     indices_array: &PrimitiveArray,
 ) -> VortexResult<ArrayRef> {
@@ -91,7 +96,7 @@ fn take_non_nullable_fsl<I: IntegerPType>(
     let new_len = indices.len();
 
     // Build the element indices directly without validity tracking.
-    let mut elements_indices = BufferMut::<I>::with_capacity(new_len * list_size);
+    let mut elements_indices = BufferMut::<E>::with_capacity(new_len * list_size);
 
     // Build the element indices for each list.
     for data_idx in indices {
@@ -106,7 +111,7 @@ fn take_non_nullable_fsl<I: IntegerPType>(
         for i in list_start..list_end {
             // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity.
             unsafe {
-                elements_indices.push_unchecked(I::from_usize(i).vortex_expect("i < list_end"))
+                elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end"))
             };
         }
     }
@@ -131,7 +136,7 @@ fn take_non_nullable_fsl<I: IntegerPType>(
 }
 
 /// Takes from an array when either the array or indices are nullable.
-fn take_nullable_fsl<I: IntegerPType>(
+fn take_nullable_fsl<I: IntegerPType, E: IntegerPType>(
     array: &FixedSizeListArray,
     indices_array: &PrimitiveArray,
 ) -> VortexResult<ArrayRef> {
@@ -144,7 +149,7 @@ fn take_nullable_fsl<I: IntegerPType>(
 
     // We must use placeholder zeros for null lists to maintain the array length without
     // propagating nullability to the element array's take operation.
-    let mut elements_indices = BufferMut::<I>::with_capacity(new_len * list_size);
+    let mut elements_indices = BufferMut::<E>::with_capacity(new_len * list_size);
     let mut new_validity_builder = BitBufferMut::with_capacity(new_len);
 
     // Build the element indices while tracking which lists are null.
@@ -159,7 +164,7 @@ fn take_nullable_fsl<I: IntegerPType>(
         if !is_index_valid || !array_validity.value(data_idx) {
             // Append placeholder zeros for null lists. These will be masked by the validity array.
             // We cannot use append_nulls here as explained above.
-            unsafe { elements_indices.push_n_unchecked(I::zero(), list_size) };
+            unsafe { elements_indices.push_n_unchecked(E::zero(), list_size) };
             new_validity_builder.append(false);
         } else {
             // Append the actual element indices for this list.
@@ -170,7 +175,7 @@ fn take_nullable_fsl<I: IntegerPType>(
             for i in list_start..list_end {
                 // SAFETY: We've allocated enough space for enough indices for all `new_len` lists (that each consist of `list_size = list_end - list_start` elements), so we know we have enough capacity.
                 unsafe {
-                    elements_indices.push_unchecked(I::from_usize(i).vortex_expect("i < list_end"))
+                    elements_indices.push_unchecked(E::from_usize(i).vortex_expect("i < list_end"))
                 };
             }
 

vortex-array/src/arrays/fixed_size_list/tests/take.rs

@@ -9,6 +9,7 @@ use super::common::create_empty_fsl;
 use super::common::create_large_fsl;
 use super::common::create_nullable_fsl;
 use super::common::create_single_element_fsl;
+use crate::ArrayRef;
 use crate::DynArray;
 use crate::IntoArray;
 use crate::arrays::FixedSizeListArray;
@@ -130,6 +131,38 @@ fn test_take_fsl_with_null_indices_preserves_elements() {
     assert_arrays_eq!(expected, result);
 }
 
+// Element index overflow: with u8 indices and list_size=16, data_idx=16 produces element index
+// 16*16=256 which overflows u8. The take kernel must widen the element index type.
+#[rstest]
+#[case::non_nullable(
+    FixedSizeListArray::new(
+        buffer![0u8; 320].into_array(), 16, Validity::NonNullable, 20,
+    ),
+    buffer![0u8, 16, 19].into_array(),
+    FixedSizeListArray::new(
+        buffer![0u8; 48].into_array(), 16, Validity::NonNullable, 3,
+    ),
+)]
+#[case::nullable(
+    FixedSizeListArray::new(
+        buffer![0u8; 320].into_array(), 16,
+        Validity::from_iter((0..20).map(|i| i != 5)), 20,
+    ),
+    buffer![0u8, 16, 5].into_array(),
+    FixedSizeListArray::new(
+        buffer![0u8; 48].into_array(), 16,
+        Validity::from_iter([true, true, false]), 3,
+    ),
+)]
+fn test_element_index_overflow(
+    #[case] fsl: FixedSizeListArray,
+    #[case] indices: ArrayRef,
+    #[case] expected: FixedSizeListArray,
+) {
+    let result = fsl.take(indices.to_array()).unwrap();
+    assert_arrays_eq!(expected, result);
+}
+
 // Parameterized test for nullable array scenarios that are specific to FSL's implementation.
 #[rstest]
 #[case::nullable_mixed_elements(