wip

spiraldb · Dec 20, 2024 · 0a2b50b · 0a2b50b
1 parent 2cf4f98
commit 0a2b50b
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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"),