pick smallest type

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

Author: 0ax1

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();

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

@@ -1,13 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use num_traits::AsPrimitive;
 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;
@@ -40,21 +41,34 @@ 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>();
+                let total_element_count = indices_slice
+                    .iter()
+                    .map(|idx| {
+                        let idx: usize = idx.as_();
+                        let diff: usize = (offsets_slice[idx + 1] - offsets_slice[idx]).as_();
+                        diff
+                    })
+                    .sum::<usize>();
+
+                match_smallest_offset_type!(total_element_count, |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 +78,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 +87,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 +108,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 +135,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 +157,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 +170,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 +182,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);
     }
@@ -200,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;
@@ -455,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
+        );
+    }
 }

vortex-dtype/src/ptype.rs

@@ -668,6 +668,34 @@ macro_rules! match_each_native_simd_ptype {
     }};
 }
 
+/// Macro to match the smallest offset type that can fit the given number of elements.
+///
+/// This macro selects u8, u16, u32, or u64 based on the size of `n_elements`:
+/// - u8 if n_elements <= 255
+/// - u16 if n_elements <= 65535
+/// - u32 if n_elements <= 4294967295
+/// - u64 otherwise
+#[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]