wip

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

Author: 0ax1

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

@@ -4,7 +4,7 @@
 use vortex_buffer::BitBufferMut;
 use vortex_dtype::IntegerPType;
 use vortex_dtype::Nullability;
-use vortex_dtype::match_each_integer_ptype;
+use vortex_dtype::{match_each_integer_ptype, match_smallest_offset_type};
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_panic;
@@ -81,18 +81,16 @@ fn _take<I: IntegerPType, O: IntegerPType>(
     let mut current_offset = 0u64;
     new_offsets.append_zero();
 
-    // Approximate the total element count.
-    let approx_total_count = indices
+    // 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
         })
-        .max()
-        .unwrap_or(0)
-        * array.len();
+        .sum::<usize>();
 
     for &data_idx in indices {
         let data_idx = data_idx