chore[fuzzer]: nullable take indices and compare scalar (#3582)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

fuzz/src/filter.rs

@@ -9,7 +9,7 @@ use vortex_dtype::{DType, match_each_native_ptype};
 use vortex_error::VortexResult;
 use vortex_scalar::match_each_decimal_value_type;
 
-use crate::take::take_canonical_array;
+use crate::take::take_canonical_array_non_nullable_indices;
 
 pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResult<ArrayRef> {
     let validity = if array.dtype().is_nullable() {
@@ -103,7 +103,7 @@ pub fn filter_canonical_array(array: &dyn Array, filter: &[bool]) -> VortexResul
                     indices.push(idx);
                 }
             }
-            take_canonical_array(array, &indices)
+            take_canonical_array_non_nullable_indices(array, indices.as_slice())
         }
         d @ (DType::Null | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")

fuzz/src/lib.rs

@@ -15,13 +15,13 @@ use std::ops::{Range, RangeInclusive};
 use libfuzzer_sys::arbitrary::Error::EmptyChoose;
 use libfuzzer_sys::arbitrary::{Arbitrary, Result, Unstructured};
 pub use sort::sort_canonical_array;
+use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::arbitrary::ArbitraryArray;
 use vortex_array::compute::Operator;
 use vortex_array::search_sorted::{SearchResult, SearchSortedSide};
 use vortex_array::{Array, ArrayRef, IntoArray};
 use vortex_btrblocks::BtrBlocksCompressor;
-use vortex_buffer::Buffer;
-use vortex_dtype::DType;
+use vortex_dtype::{DType, Nullability};
 use vortex_error::{VortexUnwrap, vortex_panic};
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
@@ -117,11 +117,11 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
 
                     let indices = random_vec_in_range(u, 0, current_array.len() - 1)?;
                     current_array = take_canonical_array(&current_array, &indices).vortex_unwrap();
-                    let indices_array = indices
-                        .iter()
-                        .map(|i| *i as u64)
-                        .collect::<Buffer<u64>>()
-                        .into_array();
+                    let indices_array = PrimitiveArray::from_option_iter(
+                        indices.iter().map(|i| i.map(|i| i as u64)),
+                    )
+                    .into_array();
+
                     let compressed = BtrBlocksCompressor.compress(&indices_array).vortex_unwrap();
                     (
                         Action::Take(compressed),
@@ -175,7 +175,9 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
                             .scalar_at(u.choose_index(current_array.len())?)
                             .vortex_unwrap()
                     } else {
-                        random_scalar(u, current_array.dtype())?
+                        // We can compare arrays with different nullability
+                        let null: Nullability = u.arbitrary()?;
+                        random_scalar(u, &current_array.dtype().union_nullability(null))?
                     };
 
                     let op = u.arbitrary()?;
@@ -194,11 +196,19 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
     }
 }
 
-fn random_vec_in_range(u: &mut Unstructured<'_>, min: usize, max: usize) -> Result<Vec<usize>> {
+fn random_vec_in_range(
+    u: &mut Unstructured<'_>,
+    min: usize,
+    max: usize,
+) -> Result<Vec<Option<usize>>> {
     iter::from_fn(|| {
-        u.arbitrary()
-            .unwrap_or(false)
-            .then(|| u.int_in_range(min..=max))
+        u.arbitrary().unwrap_or(false).then(|| {
+            if u.arbitrary()? {
+                Ok(None)
+            } else {
+                Ok(Some(u.int_in_range(min..=max)?))
+            }
+        })
     })
     .collect::<Result<Vec<_>>>()
 }

fuzz/src/sort.rs

@@ -7,7 +7,7 @@ use vortex_dtype::{DType, NativePType, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult, VortexUnwrap};
 use vortex_scalar::match_each_decimal_value_type;
 
-use crate::take::take_canonical_array;
+use crate::take::take_canonical_array_non_nullable_indices;
 
 pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
     match array.dtype() {
@@ -67,7 +67,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
                     .partial_cmp(&array.scalar_at(*b).vortex_unwrap())
                     .vortex_expect("must be a valid comparison")
             });
-            take_canonical_array(array, &sort_indices)
+            take_canonical_array_non_nullable_indices(array, &sort_indices)
         }
         DType::List(..) => {
             let mut sort_indices = (0..array.len()).collect::<Vec<_>>();
@@ -78,7 +78,7 @@ pub fn sort_canonical_array(array: &dyn Array) -> VortexResult<ArrayRef> {
                     .partial_cmp(&array.scalar_at(*b).vortex_unwrap())
                     .vortex_expect("must be a valid comparison")
             });
-            take_canonical_array(array, &sort_indices)
+            take_canonical_array_non_nullable_indices(array, &sort_indices)
         }
         d @ (DType::Null | DType::Extension(_)) => {
             unreachable!("DType {d} not supported for fuzzing")

fuzz/src/take.rs

@@ -9,15 +9,39 @@ use vortex_dtype::{DType, DecimalDType, NativePType, match_each_native_ptype};
 use vortex_error::VortexResult;
 use vortex_scalar::{NativeDecimalType, match_each_decimal_value_type};
 
-pub fn take_canonical_array(array: &dyn Array, indices: &[usize]) -> VortexResult<ArrayRef> {
+pub fn take_canonical_array_non_nullable_indices(
+    array: &dyn Array,
+    indices: &[usize],
+) -> VortexResult<ArrayRef> {
+    take_canonical_array(
+        array,
+        indices
+            .iter()
+            .map(|i| Some(*i))
+            .collect::<Vec<_>>()
+            .as_slice(),
+    )
+}
+
+pub fn take_canonical_array(
+    array: &dyn Array,
+    indices: &[Option<usize>],
+) -> VortexResult<ArrayRef> {
     let validity = if array.dtype().is_nullable() {
         let validity_idx = array.validity_mask()?.to_boolean_buffer();
 
-        Validity::from_iter(indices.iter().map(|i| validity_idx.value(*i)))
+        Validity::from_iter(
+            indices
+                .iter()
+                .map(|i| i.is_some_and(|i| validity_idx.value(i))),
+        )
     } else {
         Validity::NonNullable
     };
 
+    let indices = indices.iter().map(|i| i.unwrap_or(0)).collect::<Vec<_>>();
+    let indices = indices.as_slice();
+
     match array.dtype() {
         DType::Bool(_) => {
             let bool_array = array.to_bool()?;
@@ -55,7 +79,7 @@ pub fn take_canonical_array(array: &dyn Array, indices: &[usize]) -> VortexResul
             let taken_children = struct_array
                 .fields()
                 .iter()
-                .map(|c| take_canonical_array(c, indices))
+                .map(|c| take_canonical_array_non_nullable_indices(c, indices))
                 .collect::<VortexResult<Vec<_>>>()?;
 
             StructArray::try_new(