wip

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-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]