pick smallest type

0ax1 · 0ax1 · commit bc019fb7e67d · 2025-12-11T13:52:22.000Z
Signed-off-by: Alexander Droste &lt;alexander.droste@protonmail.com&gt;
diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs
@@ -41,6 +41,7 @@ use vortex_dtype::StructFields;
 use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_dtype::match_each_native_ptype;
+use vortex_dtype::match_smallest_offset_type;
 use vortex_error::VortexError;
 use vortex_error::VortexExpect;
 use vortex_error::vortex_panic;
@@ -124,27 +125,6 @@ fn canonicalize_sparse_lists(
     values_dtype: Arc<DType>,
     nullability: Nullability,
 ) -> Canonical {
-    // TODO(connor): We should move this to `vortex-dtype` so that we can use this elsewhere.
-    macro_rules! match_smallest_offset_type {
-        ($n_elements:expr, | $offset_type:ident | $body:block) => {{
-            let n_elements = $n_elements;
-            if n_elements <= u8::MAX as usize {
-                type $offset_type = u8;
-                $body
-            } else if n_elements <= u16::MAX as usize {
-                type $offset_type = u16;
-                $body
-            } else if n_elements <= u32::MAX as usize {
-                type $offset_type = u32;
-                $body
-            } else {
-                assert!(u64::try_from(n_elements).is_ok());
-                type $offset_type = u64;
-                $body
-            }
-        }};
-    }
-
     let resolved_patches = array.resolved_patches();
 
     let indices = resolved_patches.indices().to_primitive();
diff --git a/vortex-array/src/arrays/list/compute/take.rs b/vortex-array/src/arrays/list/compute/take.rs
@@ -5,9 +5,9 @@ use vortex_buffer::BitBufferMut;
 use vortex_dtype::IntegerPType;
 use vortex_dtype::Nullability;
 use vortex_dtype::match_each_integer_ptype;
+use vortex_dtype::match_smallest_offset_type;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 
 use crate::Array;
@@ -34,27 +34,33 @@ use crate::vtable::ValidityHelper;
 /// that lists are stored contiguously and in-order (`offset[i+1] >= offset[i]`). Taking
 /// non-contiguous indices would violate this requirement.
 impl TakeKernel for ListVTable {
+    #[expect(clippy::cognitive_complexity)]
     fn take(&self, array: &ListArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
         let indices = indices.to_primitive();
         let offsets = array.offsets().to_primitive();
+        // This is an over-approximation of the total number of elements in the resulting array.
+        let total_approx = array.elements().len() * indices.len();
 
         match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| {
+            let offsets_slice = offsets.as_slice::<O>();
             match_each_integer_ptype!(indices.ptype(), |I| {
-                _take::<I, O>(
-                    array,
-                    offsets.as_slice::<O>(),
-                    &indices,
-                    array.validity_mask(),
-                    indices.validity_mask(),
-                )
+                match_smallest_offset_type!(total_approx, |OutputOffsetType| {
+                    _take::<I, O, OutputOffsetType>(
+                        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, OutputOffsetType: IntegerPType>(
     array: &ListArray,
     offsets: &[O],
     indices_array: &PrimitiveArray,
@@ -64,7 +70,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, OutputOffsetType>(
             array,
             offsets,
             indices,
@@ -73,18 +79,18 @@ fn _take<I: IntegerPType, O: IntegerPType>(
         );
     }
 
-    let mut new_offsets =
-        PrimitiveBuilder::<u64>::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 = 0u64;
+    let mut current_offset = OutputOffsetType::zero();
     new_offsets.append_zero();
 
     for &data_idx in indices {
-        let data_idx = data_idx
-            .to_usize()
-            .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx));
+        let data_idx: usize = data_idx.as_();
 
         let start = offsets[data_idx];
         let stop = offsets[data_idx + 1];
@@ -94,15 +100,14 @@ fn _take<I: IntegerPType, O: IntegerPType>(
         // We could convert start and end to usize, but that would impose a potentially
         // harder constraint - now we don't care if they fit into usize as long as their
         // difference does.
-        let additional = (stop - start).to_usize().unwrap_or_else(|| {
-            vortex_panic!("Failed to convert range length to usize: {}", stop - start)
-        });
+        let additional: usize = (stop - start).as_();
 
         elements_to_take.reserve_exact(additional);
         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 +=
+            OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
         new_offsets.append_value(current_offset);
     }
 
@@ -122,15 +127,17 @@ fn _take<I: IntegerPType, O: IntegerPType>(
     .to_array())
 }
 
-fn _take_nullable<I: IntegerPType, O: 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::<u64>::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.
     //
@@ -142,7 +149,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 = OutputOffsetType::zero();
     new_offsets.append_zero();
 
     // Set all bits to invalid and selectively set which values are valid.
@@ -155,9 +162,7 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType>(
             continue;
         }
 
-        let data_idx = data_idx
-            .to_usize()
-            .unwrap_or_else(|| vortex_panic!("Failed to convert index to usize: {}", data_idx));
+        let data_idx: usize = data_idx.as_();
 
         if !data_validity.value(data_idx) {
             new_offsets.append_value(current_offset);
@@ -169,15 +174,14 @@ fn _take_nullable<I: IntegerPType, O: IntegerPType>(
         let stop = offsets[data_idx + 1];
 
         // See the note it the `take` on the reasoning
-        let additional = (stop - start).to_usize().unwrap_or_else(|| {
-            vortex_panic!("Failed to convert range length to usize: {}", stop - start)
-        });
+        let additional: usize = (stop - start).as_();
 
         elements_to_take.reserve_exact(additional);
         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 +=
+            OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
         new_offsets.append_value(current_offset);
         new_validity.set(idx);
     }
@@ -201,6 +205,9 @@ mod test {
     use vortex_dtype::DType;
     use vortex_dtype::Nullability;
     use vortex_dtype::PType::I32;
+    use vortex_dtype::PType::U8;
+    use vortex_dtype::PType::U16;
+    use vortex_dtype::PType::U32;
     use vortex_scalar::Scalar;
 
     use crate::Array;
@@ -455,4 +462,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: u32,
+        #[case] expected_ptype: vortex_dtype::PType,
+    ) {
+        let elements: Vec<_> = (0..element_count).collect();
+        let elements_array = PrimitiveArray::from_iter(elements).to_array();
+
+        let mut offsets = Vec::with_capacity((element_count + 1) as usize);
+        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).collect();
+        let result = take(&list, &PrimitiveArray::from_iter(indices).to_array()).unwrap();
+
+        assert_eq!(
+            result.to_listview().offsets().dtype().as_ptype(),
+            expected_ptype
+        );
+    }
 }
diff --git a/vortex-dtype/src/ptype.rs b/vortex-dtype/src/ptype.rs
@@ -668,6 +668,28 @@ macro_rules! match_each_native_simd_ptype {
     }};
 }
 
+/// Macro to match the smallest offset type for a given value
+#[macro_export]
+macro_rules! match_smallest_offset_type {
+    ($n_elements:expr, | $offset_type:ident | $body:block) => {{
+        let n_elements = $n_elements;
+        if n_elements <= u8::MAX as usize {
+            type $offset_type = u8;
+            $body
+        } else if n_elements <= u16::MAX as usize {
+            type $offset_type = u16;
+            $body
+        } else if n_elements <= u32::MAX as usize {
+            type $offset_type = u32;
+            $body
+        } else {
+            assert!(u64::try_from(n_elements).is_ok());
+            type $offset_type = u64;
+            $body
+        }
+    }};
+}
+
 impl PType {
     /// Returns `true` iff this PType is an unsigned integer type
     #[inline]
diff --git a/vortex-io/src/file/driver.rs b/vortex-io/src/file/driver.rs
@@ -125,7 +125,7 @@ impl State {
         }
     }
 
-    #[allow(clippy::cognitive_complexity)]
+    #[expect(clippy::cognitive_complexity)]
     fn on_event(&mut self, event: ReadEvent) {
         tracing::debug!(?event, "Received ReadEvent");
         match event {

Original file line number	Diff line number	Diff line change
`@@ -125,7 +125,7 @@ impl State {`
`125`	`125`	`}`
`126`	`126`	`}`
`127`	`127`
`128`		`- #[allow(clippy::cognitive_complexity)]`
	`128`	`+ #[expect(clippy::cognitive_complexity)]`
`129`	`129`	`fn on_event(&mut self, event: ReadEvent) {`
`130`	`130`	`tracing::debug!(?event, "Received ReadEvent");`
`131`	`131`	`match event {`