chore: Fix fuzzers baseline nullability for compare and take

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

fuzz/src/array/compare.rs

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::fmt::Debug;
-use std::ops::Deref;
 
 use vortex_array::accessor::ArrayAccessor;
 use vortex_array::arrays::BoolArray;
@@ -11,7 +10,7 @@ use vortex_array::validity::Validity;
 use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
 use vortex_buffer::BitBuffer;
 use vortex_dtype::{
-    DType, NativeDecimalType, NativePType, match_each_decimal_value_type, match_each_native_ptype,
+    DType, NativePType, Nullability, match_each_decimal_value_type, match_each_native_ptype,
 };
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 use vortex_scalar::Scalar;
@@ -29,6 +28,8 @@ pub fn compare_canonical_array(
         .into_array());
     }
 
+    let result_nullability = array.dtype().nullability() | value.dtype().nullability();
+
     match array.dtype() {
         DType::Bool(_) => {
             let bool = value
@@ -44,6 +45,7 @@ pub fn compare_canonical_array(
                     .map(|(b, v)| v.then_some(b)),
                 bool,
                 operator,
+                result_nullability,
             ))
         }
         DType::Primitive(p, _) => {
@@ -62,6 +64,7 @@ pub fn compare_canonical_array(
                         .map(|(b, v)| v.then_some(b)),
                     pval,
                     operator,
+                    result_nullability,
                 ))
             })
         }
@@ -75,14 +78,15 @@ pub fn compare_canonical_array(
                     .cast::<D>()
                     .ok_or_else(|| vortex_err!("todo: handle upcast of decimal array"))?;
                 let buf = decimal_array.buffer::<D>();
-                Ok(compare_native_decimal_type(
+                Ok(compare_to(
                     buf.as_slice()
                         .iter()
                         .copied()
                         .zip(array.validity_mask().to_bit_buffer().iter())
                         .map(|(b, v)| v.then_some(b)),
                     dval,
                     operator,
+                    result_nullability,
                 ))
             })
         }
@@ -93,8 +97,9 @@ pub fn compare_canonical_array(
                 .vortex_expect("nulls handled before");
             compare_to(
                 iter.map(|v| v.map(|b| unsafe { str::from_utf8_unchecked(b) })),
-                utf8_value.deref(),
+                &utf8_value,
                 operator,
+                result_nullability,
             )
         }),
         DType::Binary(_) => array.to_varbinview().with_iterator(|iter| {
@@ -106,8 +111,9 @@ pub fn compare_canonical_array(
                 // Don't understand the lifetime problem here but identity map makes it go away
                 #[allow(clippy::map_identity)]
                 iter.map(|v| v),
-                binary_value.deref(),
+                &binary_value,
                 operator,
+                result_nullability,
             )
         }),
         DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => {
@@ -125,56 +131,48 @@ pub fn compare_canonical_array(
     }
 }
 
+#[allow(clippy::unwrap_used)]
 fn compare_to<T: PartialOrd + PartialEq + Debug>(
     values: impl Iterator<Item = Option<T>>,
     cmp_value: T,
     operator: Operator,
+    nullability: Nullability,
 ) -> ArrayRef {
-    BoolArray::from_iter(values.map(|val| {
-        val.map(|v| match operator {
-            Operator::Eq => v == cmp_value,
-            Operator::NotEq => v != cmp_value,
-            Operator::Gt => v > cmp_value,
-            Operator::Gte => v >= cmp_value,
-            Operator::Lt => v < cmp_value,
-            Operator::Lte => v <= cmp_value,
-        })
-    }))
-    .into_array()
+    let eval_fn = |v| match operator {
+        Operator::Eq => v == cmp_value,
+        Operator::NotEq => v != cmp_value,
+        Operator::Gt => v > cmp_value,
+        Operator::Gte => v >= cmp_value,
+        Operator::Lt => v < cmp_value,
+        Operator::Lte => v <= cmp_value,
+    };
+
+    if !nullability.is_nullable() {
+        BoolArray::from_iter(values.map(|val| val.unwrap()).map(eval_fn)).into_array()
+    } else {
+        BoolArray::from_iter(values.map(|val| val.map(eval_fn))).into_array()
+    }
 }
 
+#[allow(clippy::unwrap_used)]
 fn compare_native_ptype<T: NativePType>(
     values: impl Iterator<Item = Option<T>>,
     cmp_value: T,
     operator: Operator,
+    nullability: Nullability,
 ) -> ArrayRef {
-    BoolArray::from_iter(values.map(|val| {
-        val.map(|v| match operator {
-            Operator::Eq => v.is_eq(cmp_value),
-            Operator::NotEq => !v.is_eq(cmp_value),
-            Operator::Gt => v.is_gt(cmp_value),
-            Operator::Gte => v.is_ge(cmp_value),
-            Operator::Lt => v.is_lt(cmp_value),
-            Operator::Lte => v.is_le(cmp_value),
-        })
-    }))
-    .into_array()
-}
+    let eval_fn = |v: T| match operator {
+        Operator::Eq => v.is_eq(cmp_value),
+        Operator::NotEq => !v.is_eq(cmp_value),
+        Operator::Gt => v.is_gt(cmp_value),
+        Operator::Gte => v.is_ge(cmp_value),
+        Operator::Lt => v.is_lt(cmp_value),
+        Operator::Lte => v.is_le(cmp_value),
+    };
 
-fn compare_native_decimal_type<D: NativeDecimalType>(
-    values: impl Iterator<Item = Option<D>>,
-    cmp_value: D,
-    operator: Operator,
-) -> ArrayRef {
-    BoolArray::from_iter(values.map(|val| {
-        val.map(|v| match operator {
-            Operator::Eq => v == cmp_value,
-            Operator::NotEq => v != cmp_value,
-            Operator::Gt => v > cmp_value,
-            Operator::Gte => v >= cmp_value,
-            Operator::Lt => v < cmp_value,
-            Operator::Lte => v <= cmp_value,
-        })
-    }))
-    .into_array()
+    if !nullability.is_nullable() {
+        BoolArray::from_iter(values.map(|val| val.unwrap()).map(eval_fn)).into_array()
+    } else {
+        BoolArray::from_iter(values.map(|val| val.map(eval_fn))).into_array()
+    }
 }

fuzz/src/array/mod.rs

@@ -180,11 +180,23 @@ impl<'a> Arbitrary<'a> for FuzzArrayAction {
                     }
 
                     let indices = random_vec_in_range(u, 0, current_array.len() - 1)?;
+                    let nullable = indices.contains(&None);
+
                     current_array = take_canonical_array(&current_array, &indices).vortex_unwrap();
-                    let indices_array = PrimitiveArray::from_option_iter(
-                        indices.iter().map(|i| i.map(|i| i as u64)),
-                    )
-                    .into_array();
+                    let indices_array = if nullable {
+                        PrimitiveArray::from_option_iter(
+                            indices.iter().map(|i| i.map(|i| i as u64)),
+                        )
+                        .into_array()
+                    } else {
+                        PrimitiveArray::from_iter(
+                            indices
+                                .iter()
+                                .map(|i| i.vortex_expect("must be present"))
+                                .map(|i| i as u64),
+                        )
+                        .into_array()
+                    };
 
                     let compressed = BtrBlocksCompressor::default()
                         .compress(&indices_array)