fix: return nullable type from between primitive comput fn (#2776)

Author: joseph-isaacs

encodings/dict/src/compute/compare.rs

@@ -24,7 +24,6 @@ impl CompareFn<&DictArray> for DictEncoding {
                 &ConstantArray::new(rhs, lhs.values().len()),
                 operator,
             )?;
-
             return if operator == Operator::Eq {
                 let result_nullability =
                     compare_result.dtype().nullability() | lhs.dtype().nullability();

vortex-array/src/arrays/primitive/compute/between.rs

@@ -1,5 +1,5 @@
 use arrow_buffer::BooleanBuffer;
-use vortex_dtype::{NativePType, match_each_native_ptype};
+use vortex_dtype::{NativePType, Nullability, match_each_native_ptype};
 use vortex_error::VortexResult;
 
 use crate::arrays::{BoolArray, PrimitiveArray, PrimitiveEncoding};
@@ -19,8 +19,14 @@ impl BetweenFn<&PrimitiveArray> for PrimitiveEncoding {
             return Ok(None);
         };
 
+        // Note, we know that have checked before that the lower and upper bounds are not constant
+        // null values
+
+        let nullability =
+            arr.dtype.nullability() | lower.dtype().nullability() | upper.dtype().nullability();
+
         Ok(Some(match_each_native_ptype!(arr.ptype(), |$P| {
-            between_impl::<$P>(arr, $P::try_from(lower)?, $P::try_from(upper)?, options)
+            between_impl::<$P>(arr, $P::try_from(lower)?, $P::try_from(upper)?, nullability, options)
         })))
     }
 }
@@ -29,22 +35,43 @@ fn between_impl<T: NativePType + Copy>(
     arr: &PrimitiveArray,
     lower: T,
     upper: T,
+    nullability: Nullability,
     options: &BetweenOptions,
 ) -> ArrayRef {
     match (options.lower_strict, options.upper_strict) {
         // Note: these comparisons are explicitly passed in to allow function impl inlining
-        (StrictComparison::Strict, StrictComparison::Strict) => {
-            between_impl_(arr, lower, NativePType::is_lt, upper, NativePType::is_lt)
-        }
-        (StrictComparison::Strict, StrictComparison::NonStrict) => {
-            between_impl_(arr, lower, NativePType::is_lt, upper, NativePType::is_le)
-        }
-        (StrictComparison::NonStrict, StrictComparison::Strict) => {
-            between_impl_(arr, lower, NativePType::is_le, upper, NativePType::is_lt)
-        }
-        (StrictComparison::NonStrict, StrictComparison::NonStrict) => {
-            between_impl_(arr, lower, NativePType::is_le, upper, NativePType::is_le)
-        }
+        (StrictComparison::Strict, StrictComparison::Strict) => between_impl_(
+            arr,
+            lower,
+            NativePType::is_lt,
+            upper,
+            NativePType::is_lt,
+            nullability,
+        ),
+        (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl_(
+            arr,
+            lower,
+            NativePType::is_lt,
+            upper,
+            NativePType::is_le,
+            nullability,
+        ),
+        (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl_(
+            arr,
+            lower,
+            NativePType::is_le,
+            upper,
+            NativePType::is_lt,
+            nullability,
+        ),
+        (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl_(
+            arr,
+            lower,
+            NativePType::is_le,
+            upper,
+            NativePType::is_le,
+            nullability,
+        ),
     }
 }
 
@@ -54,6 +81,7 @@ fn between_impl_<T>(
     lower_fn: impl Fn(T, T) -> bool,
     upper: T,
     upper_fn: impl Fn(T, T) -> bool,
+    nullability: Nullability,
 ) -> ArrayRef
 where
     T: NativePType + Copy,
@@ -65,7 +93,7 @@ where
             let i = unsafe { *slice.get_unchecked(idx) };
             lower_fn(lower, i) & upper_fn(i, upper)
         }),
-        arr.validity().clone(),
+        arr.validity().clone().union_nullability(nullability),
     )
     .into_array()
 }

vortex-array/src/compute/between.rs

@@ -1,6 +1,8 @@
 use vortex_dtype::{DType, Nullability};
 use vortex_error::{VortexExpect, VortexResult};
+use vortex_scalar::Scalar;
 
+use crate::arrays::ConstantArray;
 use crate::compute::{BinaryOperator, Operator, binary_boolean, compare};
 use crate::{Array, ArrayRef, Canonical, Encoding, IntoArray};
 
@@ -91,6 +93,21 @@ pub fn between(
     debug_assert_eq!(arr.len(), lower.len());
     debug_assert_eq!(arr.len(), upper.len());
 
+    // A quick check to see if either array might is a null constant array.
+    if lower.is_invalid(0)? || upper.is_invalid(0)? {
+        if let (Some(c_lower), Some(c_upper)) = (lower.as_constant(), upper.as_constant()) {
+            if c_lower.is_null() || c_upper.is_null() {
+                return Ok(ConstantArray::new(
+                    Scalar::null(arr.dtype().with_nullability(
+                        lower.dtype().nullability() | upper.dtype().nullability(),
+                    )),
+                    arr.len(),
+                )
+                .to_array());
+            }
+        }
+    }
+
     let result = between_impl(arr, lower, upper, options)?;
 
     debug_assert_eq!(result.len(), arr.len());

vortex-array/src/validity.rs

@@ -76,6 +76,14 @@ impl Validity {
         }
     }
 
+    /// The union nullability and validity.
+    pub fn union_nullability(self, nullability: Nullability) -> Self {
+        match nullability {
+            Nullability::NonNullable => self,
+            Nullability::Nullable => self.into_nullable(),
+        }
+    }
+
     pub fn all_valid(&self) -> VortexResult<bool> {
         Ok(match self {
             Validity::NonNullable | Validity::AllValid => true,