Repository: spiraldb/vortex
Pull request #2776: fix: return nullable type from between primitive comput fn
Author: login: joseph-isaacs
Created: 2025-03-24T15:46:36Z | Merged: 2025-03-25T11:14:04Z



Diff:
diff --git a/encodings/dict/src/compute/compare.rs b/encodings/dict/src/compute/compare.rs
index 19079ddc30d..c727c7542d4 100644
--- a/encodings/dict/src/compute/compare.rs
+++ b/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();
diff --git a/vortex-array/src/arrays/primitive/compute/between.rs b/vortex-array/src/arrays/primitive/compute/between.rs
index d3d3f81899b..3740e34639b 100644
--- a/vortex-array/src/arrays/primitive/compute/between.rs
+++ b/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()
 }
diff --git a/vortex-array/src/compute/between.rs b/vortex-array/src/compute/between.rs
index 482d8abc5c9..3a1a1ab6f73 100644
--- a/vortex-array/src/compute/between.rs
+++ b/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());
diff --git a/vortex-array/src/validity.rs b/vortex-array/src/validity.rs
index b4ffb402c23..fa9d0d180e1 100644
--- a/vortex-array/src/validity.rs
+++ b/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,


Reviews:
user: login: robert3005
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-24T18:39:10Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T09:45:57Z
user: login: robert3005
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T10:46:28Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T10:59:47Z
user: login: robert3005
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T11:03:38Z

Review comments:
diff_hunk: @@ -65,7 +90,7 @@ where
             let i = unsafe { *slice.get_unchecked(idx) };
             lower_fn(lower, i) & upper_fn(i, upper)
         }),
-        arr.validity().clone(),
+        arr.validity().clone().merge_nullability(nullability),
path: vortex-array/src/arrays/primitive/compute/between.rs
user: login: robert3005
body: if nullability == Nullability::NonNullable { arr.validity().clone() } else { arr.validity().clone().into_nullable() }
created_at: 2025-03-24T18:39:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -65,7 +90,7 @@ where
             let i = unsafe { *slice.get_unchecked(idx) };
             lower_fn(lower, i) & upper_fn(i, upper)
         }),
-        arr.validity().clone(),
+        arr.validity().clone().merge_nullability(nullability),
path: vortex-array/src/arrays/primitive/compute/between.rs
user: login: joseph-isaacs
body: So you are saying you don't want to add a new method?
created_at: 2025-03-25T09:45:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -22,6 +23,16 @@ impl BetweenFn<&PrimitiveArray> for PrimitiveEncoding {
         let nullability =
             arr.dtype.nullability() | lower.dtype().nullability() | upper.dtype().nullability();
 
+        if lower.is_null() || upper.is_null() {
path: vortex-array/src/arrays/primitive/compute/between.rs
user: login: robert3005
body: Shouldn’t this be in entry point function, seems pretty universal?
created_at: 2025-03-25T10:46:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -22,6 +23,16 @@ impl BetweenFn<&PrimitiveArray> for PrimitiveEncoding {
         let nullability =
             arr.dtype.nullability() | lower.dtype().nullability() | upper.dtype().nullability();
 
+        if lower.is_null() || upper.is_null() {
path: vortex-array/src/arrays/primitive/compute/between.rs
user: login: joseph-isaacs
body: sounds good
created_at: 2025-03-25T10:59:46Z
author_association: CONTRIBUTOR

PR comments:
user: login: joseph-isaacs
created_at: 2025-03-25T09:59:28Z
body: I think that this pattern exist in many places in the code base
author_association: CONTRIBUTOR
user: login: robert3005
created_at: 2025-03-25T10:04:11Z
body: If this method is useful then the pr that adds it should have more than use of it. If this pattern happens more often in the code base then it should be relatively easy to replace other occurrences. I would start with my suggestion and generalise from there
author_association: CONTRIBUTOR
user: login: joseph-isaacs
created_at: 2025-03-25T10:13:34Z
body: So would you like me to pull in those examples in this PR? Also see this stacked https://github.com/spiraldb/vortex/pull/2777
author_association: CONTRIBUTOR
user: login: robert3005
created_at: 2025-03-25T10:27:06Z
body: I see, didn't notice #2777. We can add the method then but need to think about the name. There's something unintuitive about `merge_nullability`. Don't have an immediate suggestion.
author_association: CONTRIBUTOR