wip

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

Author: 0ax1

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

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use num_traits::ToPrimitive;
 use vortex_buffer::BitBufferMut;
 use vortex_dtype::IntegerPType;
 use vortex_dtype::Nullability;
@@ -40,21 +41,44 @@ impl TakeKernel for ListVTable {
 
         match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| {
             match_each_integer_ptype!(indices.ptype(), |I| {
-                _take::<I, O>(
-                    array,
-                    offsets.as_slice::<O>(),
-                    &indices,
-                    array.validity_mask(),
-                    indices.validity_mask(),
-                )
+                let offsets_slice = offsets.as_slice::<O>();
+                let indices_slice: &[I] = indices.as_slice::<I>();
+
+                // Calculate total count to determine appropriate accumulation type
+                let total_count = indices_slice
+                    .iter()
+                    .map(|idx| {
+                        let idx = idx.to_usize().unwrap_or_else(|| {
+                            vortex_panic!("Failed to convert index to usize: {}", idx)
+                        });
+                        (offsets_slice[idx + 1] - offsets_slice[idx])
+                            .to_usize()
+                            .unwrap_or_else(|| {
+                                vortex_panic!(
+                                    "Failed to convert offset difference to usize: {}",
+                                    offsets_slice[idx + 1] - offsets_slice[idx]
+                                )
+                            })
+                    })
+                    .sum::<usize>();
+
+                match_smallest_offset_type!(total_count, |AccumType| {
+                    _take::<I, O, AccumType>(
+                        array,
+                        offsets_slice,
+                        &indices,
+                        array.validity_mask(),
+                        indices.validity_mask(),
+                    )
+                })
             })
         })
     }
 }
 
 register_kernel!(TakeKernelAdapter(ListVTable).lift());
 
-fn _take<I: IntegerPType, O: IntegerPType>(
+fn _take<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
     array: &ListArray,
     offsets: &[O],
     indices_array: &PrimitiveArray,
@@ -64,7 +88,7 @@ fn _take<I: IntegerPType, O: IntegerPType>(
     let indices: &[I] = indices_array.as_slice::<I>();
 
     if !indices_validity_mask.all_true() || !data_validity.all_true() {
-        return _take_nullable::<I, O>(
+        return _take_nullable::<I, O, AccumType>(
             array,
             offsets,
             indices,
@@ -74,24 +98,13 @@ fn _take<I: IntegerPType, O: IntegerPType>(
     }
 
     let mut new_offsets =
-        PrimitiveBuilder::<u64>::with_capacity(Nullability::NonNullable, indices.len());
+        PrimitiveBuilder::<AccumType>::with_capacity(Nullability::NonNullable, indices.len());
     let mut elements_to_take =
         PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len());
 
-    let mut current_offset = 0u64;
+    let mut current_offset = AccumType::zero();
     new_offsets.append_zero();
 
-    // Total element count.
-    let total_count = indices
-        .iter()
-        .map(|idx| {
-            let idx = idx
-                .to_usize()
-                .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", idx));
-            (offsets[idx + 1] - offsets[idx]).as_() as usize
-        })
-        .sum::<usize>();
-
     for &data_idx in indices {
         let data_idx = data_idx
             .to_usize()
@@ -113,7 +126,13 @@ 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).as_() as u64;
+        current_offset += AccumType::from_usize((stop - start).to_usize().unwrap_or_else(|| {
+            vortex_panic!(
+                "Failed to convert offset difference to usize: {}",
+                stop - start
+            )
+        }))
+        .vortex_expect("offset conversion");
         new_offsets.append_value(current_offset);
     }
 
@@ -133,15 +152,15 @@ fn _take<I: IntegerPType, O: IntegerPType>(
     .to_array())
 }
 
-fn _take_nullable<I: IntegerPType, O: IntegerPType>(
+fn _take_nullable<I: IntegerPType, O: IntegerPType, AccumType: IntegerPType>(
     array: &ListArray,
     offsets: &[O],
     indices: &[I],
     data_validity: Mask,
     indices_validity: Mask,
 ) -> VortexResult<ArrayRef> {
     let mut new_offsets =
-        PrimitiveBuilder::<u64>::with_capacity(Nullability::NonNullable, indices.len());
+        PrimitiveBuilder::<AccumType>::with_capacity(Nullability::NonNullable, indices.len());
 
     // This will be the indices we push down to the child array to call `take` with.
     //
@@ -153,7 +172,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 = 0u64;
+    let mut current_offset = AccumType::zero();
     new_offsets.append_zero();
 
     // Set all bits to invalid and selectively set which values are valid.
@@ -188,7 +207,13 @@ 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).as_() as u64;
+        current_offset += AccumType::from_usize((stop - start).to_usize().unwrap_or_else(|| {
+            vortex_panic!(
+                "Failed to convert offset difference to usize: {}",
+                stop - start
+            )
+        }))
+        .vortex_expect("offset conversion");
         new_offsets.append_value(current_offset);
         new_validity.set(idx);
     }