wip

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

Author: 0ax1

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

@@ -43,18 +43,17 @@ impl TakeKernel for ListVTable {
             match_each_integer_ptype!(indices.ptype(), |I| {
                 let offsets_slice = offsets.as_slice::<O>();
                 let indices_slice: &[I] = indices.as_slice::<I>();
-                let approx_total_count = indices_slice
+                let total_element_count = indices_slice
                     .iter()
                     .map(|idx| {
                         let idx: usize = idx.as_();
-                        let length: usize = (offsets_slice[idx + 1] - offsets_slice[idx]).as_();
-                        length
+                        let diff: usize = (offsets_slice[idx + 1] - offsets_slice[idx]).as_();
+                        diff
                     })
-                    .max()
-                    .unwrap_or(0);
+                    .sum::<usize>();
 
-                match_smallest_offset_type!(approx_total_count, |AccumType| {
-                    _take::<I, O, AccumType>(
+                match_smallest_offset_type!(total_element_count, |OutputOffsetType| {
+                    _take::<I, O, OutputOffsetType>(
                         array,
                         offsets_slice,
                         &indices,
@@ -69,7 +68,7 @@ impl TakeKernel for ListVTable {
 
 register_kernel!(TakeKernelAdapter(ListVTable).lift());
 
-fn _take<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
+fn _take<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
     array: &ListArray,
     offsets: &[O],
     indices_array: &PrimitiveArray,
@@ -79,7 +78,7 @@ fn _take<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
     let indices: &[I] = indices_array.as_slice::<I>();
 
     if !indices_validity_mask.all_true() || !data_validity.all_true() {
-        return _take_nullable::<I, O, AccumType>(
+        return _take_nullable::<I, O, OutputOffsetType>(
             array,
             offsets,
             indices,
@@ -88,12 +87,14 @@ fn _take<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
         );
     }
 
-    let mut new_offsets =
-        PrimitiveBuilder::<AccumType>::with_capacity(Nullability::NonNullable, indices.len());
+    let mut new_offsets = PrimitiveBuilder::<OutputOffsetType>::with_capacity(
+        Nullability::NonNullable,
+        indices.len(),
+    );
     let mut elements_to_take =
         PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len());
 
-    let mut current_offset = AccumType::zero();
+    let mut current_offset = OutputOffsetType::zero();
     new_offsets.append_zero();
 
     for &data_idx in indices {
@@ -114,7 +115,7 @@ fn _take<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
             elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional"));
         }
         current_offset +=
-            AccumType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
+            OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
         new_offsets.append_value(current_offset);
     }
 
@@ -134,15 +135,17 @@ fn _take<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
     .to_array())
 }
 
-fn _take_nullable<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
+fn _take_nullable<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
     array: &ListArray,
     offsets: &[O],
     indices: &[I],
     data_validity: Mask,
     indices_validity: Mask,
 ) -> VortexResult<ArrayRef> {
-    let mut new_offsets =
-        PrimitiveBuilder::<AccumType>::with_capacity(Nullability::NonNullable, indices.len());
+    let mut new_offsets = PrimitiveBuilder::<OutputOffsetType>::with_capacity(
+        Nullability::NonNullable,
+        indices.len(),
+    );
 
     // This will be the indices we push down to the child array to call `take` with.
     //
@@ -154,7 +157,7 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
     let mut elements_to_take =
         PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, 2 * indices.len());
 
-    let mut current_offset = AccumType::zero();
+    let mut current_offset = OutputOffsetType::zero();
     new_offsets.append_zero();
 
     // Set all bits to invalid and selectively set which values are valid.
@@ -186,7 +189,7 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
             elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional"));
         }
         current_offset +=
-            AccumType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
+            OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
         new_offsets.append_value(current_offset);
         new_validity.set(idx);
     }
@@ -209,7 +212,7 @@ mod test {
     use vortex_buffer::buffer;
     use vortex_dtype::DType;
     use vortex_dtype::Nullability;
-    use vortex_dtype::PType::I32;
+    use vortex_dtype::PType::{I32, U8, U16, U32};
     use vortex_scalar::Scalar;
 
     use crate::Array;
@@ -464,4 +467,34 @@ mod test {
         assert!(result_view.is_invalid(1));
         assert!(result_view.is_valid(2));
     }
+
+    #[rstest]
+    #[case(10, U8)]
+    #[case(300, U16)]
+    #[case(70000, U32)]
+    fn test_output_offset_type_selection(
+        #[case] element_count: usize,
+        #[case] expected_ptype: vortex_dtype::PType,
+    ) {
+        let elements: Vec<i32> = (0..element_count as i32).collect();
+        let elements_array = PrimitiveArray::from_iter(elements).to_array();
+
+        let mut offsets = Vec::with_capacity(element_count + 1);
+        for idx in 0..element_count {
+            offsets.push(idx as u64);
+        }
+        let offsets_array = PrimitiveArray::from_iter(offsets).to_array();
+
+        let list = ListArray::try_new(elements_array, offsets_array, Validity::NonNullable)
+            .unwrap()
+            .to_array();
+
+        let indices: Vec<u32> = (0..element_count as u32).collect();
+        let result = take(&list, &PrimitiveArray::from_iter(indices).to_array()).unwrap();
+
+        assert_eq!(
+            result.to_listview().offsets().dtype().as_ptype(),
+            expected_ptype
+        );
+    }
 }