From 0a2b50bab6684f8060a8a965c96193ef05cf41ef Mon Sep 17 00:00:00 2001
From: Joe Isaacs <joe.isaacs@live.co.uk>
Date: Fri, 20 Dec 2024 14:37:23 +0000
Subject: [PATCH] wip

---
 fuzz/src/slice.rs             |  4 ++--
 fuzz/src/take.rs              | 31 +++++++++++++++++++++----------
 vortex-dtype/src/arbitrary.rs |  2 +-
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/fuzz/src/slice.rs b/fuzz/src/slice.rs
index 3ca18b45e1..41614d59b5 100644
--- a/fuzz/src/slice.rs
+++ b/fuzz/src/slice.rs
@@ -94,11 +94,11 @@ fn shift_offsets<O: NativePType + ArrowNativeType>(offsets: PrimitiveArray) -> P
 }
 
 fn slice_primitive<T: NativePType + ArrowNativeType>(
-    prim: PrimitiveArray,
+    primitive_array: PrimitiveArray,
     validity: Validity,
     start: usize,
     stop: usize,
 ) -> ArrayData {
-    let vec_values = prim.into_maybe_null_slice::<T>();
+    let vec_values = primitive_array.into_maybe_null_slice::<T>();
     PrimitiveArray::from_vec(vec_values[start..stop].into(), validity).into_array()
 }
diff --git a/fuzz/src/take.rs b/fuzz/src/take.rs
index 65ac9cdcf1..d0da4df0f2 100644
--- a/fuzz/src/take.rs
+++ b/fuzz/src/take.rs
@@ -1,9 +1,10 @@
+use arrow_buffer::ArrowNativeType;
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::array::{BoolArray, PrimitiveArray, StructArray, VarBinViewArray};
 use vortex_array::validity::{ArrayValidity, Validity};
 use vortex_array::variants::StructArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData, IntoArrayVariant};
-use vortex_dtype::{match_each_native_ptype, DType};
+use vortex_dtype::{match_each_native_ptype, DType, NativePType};
 use vortex_error::VortexExpect;
 
 pub fn take_canonical_array(array: &ArrayData, indices: &[usize]) -> ArrayData {
@@ -30,16 +31,12 @@ pub fn take_canonical_array(array: &ArrayData, indices: &[usize]) -> ArrayData {
                 .vortex_expect("Validity length cannot mismatch")
                 .into_array()
         }
-        DType::Primitive(p, _) => match_each_native_ptype!(p, |$P| {
+        DType::Primitive(p, _) => {
             let primitive_array = array.clone().into_primitive().unwrap();
-            let vec_values = primitive_array
-                .maybe_null_slice::<$P>()
-                .iter()
-                .copied()
-                .collect::<Vec<_>>();
-            PrimitiveArray::from_vec(indices.iter().map(|i| vec_values[*i]).collect(),validity)
-                .into_array()
-        }),
+            match_each_native_ptype!(p, |$P| {
+                take_primitive::<$P>(primitive_array, validity, indices)
+            })
+        }
         DType::Utf8(_) | DType::Binary(_) => {
             let utf8 = array.clone().into_varbinview().unwrap();
             let values = utf8
@@ -70,3 +67,17 @@ pub fn take_canonical_array(array: &ArrayData, indices: &[usize]) -> ArrayData {
         _ => unreachable!("Not a canonical array"),
     }
 }
+
+fn take_primitive<T: NativePType + ArrowNativeType>(
+    primitive_array: PrimitiveArray,
+    validity: Validity,
+    indices: &[usize],
+) -> ArrayData {
+    let vec_values = primitive_array
+        .maybe_null_slice::<T>()
+        .iter()
+        .copied()
+        .collect::<Vec<_>>();
+    PrimitiveArray::from_vec(indices.iter().map(|i| vec_values[*i]).collect(), validity)
+        .into_array()
+}
diff --git a/vortex-dtype/src/arbitrary.rs b/vortex-dtype/src/arbitrary.rs
index 3cd3fe08ef..b04d4348c8 100644
--- a/vortex-dtype/src/arbitrary.rs
+++ b/vortex-dtype/src/arbitrary.rs
@@ -18,7 +18,7 @@ fn random_dtype(u: &mut Unstructured<'_>, depth: u8) -> Result<DType> {
         2 => DType::Utf8(u.arbitrary()?),
         3 => DType::Binary(u.arbitrary()?),
         4 => DType::Struct(random_struct_dtype(u, depth - 1)?, u.arbitrary()?),
-        5 => DType::List(Arc::new(u.arbitrary()?), u.arbitrary()?),
+        5 => DType::List(Arc::new(random_dtype(u, depth - 1)?), u.arbitrary()?),
         // Null,
         // Extension(ExtDType, Nullability),
         _ => unreachable!("Number out of range"),