fix: list take offset type

Signed-off-by: Alexander Droste <[email protected]>

Author: 0ax1

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

@@ -73,11 +73,12 @@ fn _take<I: IntegerPType, O: IntegerPType>(
         );
     }
 
-    let mut new_offsets = PrimitiveBuilder::with_capacity(Nullability::NonNullable, indices.len());
+    let mut new_offsets =
+        PrimitiveBuilder::<u64>::with_capacity(Nullability::NonNullable, indices.len());
     let mut elements_to_take =
         PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len());
 
-    let mut current_offset = O::zero();
+    let mut current_offset = 0u64;
     new_offsets.append_zero();
 
     for &data_idx in indices {
@@ -101,7 +102,7 @@ fn _take<I: IntegerPType, O: IntegerPType>(
         for i in 0..additional {
             elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional"));
         }
-        current_offset += stop - start;
+        current_offset += (stop - start).as_() as u64;
         new_offsets.append_value(current_offset);
     }
 
@@ -128,7 +129,8 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType>(
     data_validity: Mask,
     indices_validity: Mask,
 ) -> VortexResult<ArrayRef> {
-    let mut new_offsets = PrimitiveBuilder::with_capacity(Nullability::NonNullable, indices.len());
+    let mut new_offsets =
+        PrimitiveBuilder::<u64>::with_capacity(Nullability::NonNullable, indices.len());
 
     // This will be the indices we push down to the child array to call `take` with.
     //
@@ -140,7 +142,7 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType>(
     let mut elements_to_take =
         PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, 2 * indices.len());
 
-    let mut current_offset = O::zero();
+    let mut current_offset = 0u64;
     new_offsets.append_zero();
 
     // Set all bits to invalid and selectively set which values are valid.
@@ -175,7 +177,7 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType>(
         for i in 0..additional {
             elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional"));
         }
-        current_offset += stop - start;
+        current_offset += (stop - start).as_() as u64;
         new_offsets.append_value(current_offset);
         new_validity.set(idx);
     }
@@ -411,4 +413,46 @@ mod test {
     fn test_take_list_conformance(#[case] list: ListArray) {
         test_take_conformance(list.as_ref());
     }
+
+    #[test]
+    fn test_u64_offset_accumulation_non_nullable() {
+        let elements = buffer![0i32; 200].into_array();
+        let offsets = buffer![0u8, 200].into_array();
+        let list = ListArray::try_new(elements, offsets, Validity::NonNullable)
+            .unwrap()
+            .to_array();
+
+        // Take the same large list twice - would overflow u8 but works with u64.
+        let idx = buffer![0u8, 0].into_array();
+        let result = take(&list, &idx).unwrap();
+
+        assert_eq!(result.len(), 2);
+
+        let result_view = result.to_listview();
+        assert_eq!(result_view.len(), 2);
+        assert!(result_view.is_valid(0));
+        assert!(result_view.is_valid(1));
+    }
+
+    #[test]
+    fn test_u64_offset_accumulation_nullable() {
+        let elements = buffer![0i32; 150].into_array();
+        let offsets = buffer![0u8, 150, 150].into_array();
+        let validity = BoolArray::from_iter(vec![true, false]).to_array();
+        let list = ListArray::try_new(elements, offsets, Validity::Array(validity))
+            .unwrap()
+            .to_array();
+
+        // Take the same large list twice - would overflow u8 but works with u64.
+        let idx = PrimitiveArray::from_option_iter(vec![Some(0u8), None, Some(0u8)]).to_array();
+        let result = take(&list, &idx).unwrap();
+
+        assert_eq!(result.len(), 3);
+
+        let result_view = result.to_listview();
+        assert_eq!(result_view.len(), 3);
+        assert!(result_view.is_valid(0));
+        assert!(result_view.is_invalid(1));
+        assert!(result_view.is_valid(2));
+    }
 }