include validities in comparisons

src/compute/comparison/binary.rs

@@ -1,4 +1,5 @@
 //! Comparison functions for [`BinaryArray`]
+use crate::compute::comparison::{eq_validities, neq_validities};
 use crate::{
     array::{BinaryArray, BooleanArray, Offset},
     bitmap::Bitmap,
@@ -49,23 +50,67 @@ pub fn eq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray
     compare_op(lhs, rhs, |a, b| a == b)
 }
 
+/// Perform `lhs == rhs` operation on [`BinaryArray`] and include validities in comparison.
+/// # Panic
+/// iff the arrays do not have the same length.
+pub fn eq_and_validity<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| a == b);
+
+    eq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `lhs == rhs` operation on [`BinaryArray`] and a scalar.
 pub fn eq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
     compare_op_scalar(lhs, rhs, |a, b| a == b)
 }
 
+/// Perform `lhs == rhs` operation on [`BinaryArray`] and a scalar and include validities in comparison.
+pub fn eq_scalar_and_validity<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = compare_op_scalar(&lhs, rhs, |a, b| a == b);
+
+    eq_validities(out, validity, None)
+}
+
 /// Perform `lhs != rhs` operation on [`BinaryArray`].
 /// # Panic
 /// iff the arrays do not have the same length.
 pub fn neq<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| a != b)
 }
 
+/// Perform `lhs != rhs` operation on [`BinaryArray`].
+/// # Panic
+/// iff the arrays do not have the same length and include validities in comparison.
+pub fn neq_and_validity<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+
+    let out = compare_op(&lhs, &rhs, |a, b| a != b);
+    neq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `lhs != rhs` operation on [`BinaryArray`] and a scalar.
 pub fn neq_scalar<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
     compare_op_scalar(lhs, rhs, |a, b| a != b)
 }
 
+/// Perform `lhs != rhs` operation on [`BinaryArray`] and a scalar and include validities in comparison.
+pub fn neq_scalar_and_validity<O: Offset>(lhs: &BinaryArray<O>, rhs: &[u8]) -> BooleanArray {
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = compare_op_scalar(&lhs, rhs, |a, b| a != b);
+
+    neq_validities(out, validity, None)
+}
+
 /// Perform `lhs < rhs` operation on [`BinaryArray`].
 pub fn lt<O: Offset>(lhs: &BinaryArray<O>, rhs: &BinaryArray<O>) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| a < b)

src/compute/comparison/boolean.rs

@@ -1,4 +1,5 @@
 //! Comparison functions for [`BooleanArray`]
+use crate::compute::comparison::{eq_validities, neq_validities};
 use crate::{
     array::BooleanArray,
     bitmap::{binary, unary, Bitmap},
@@ -38,6 +39,17 @@ pub fn eq(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| !(a ^ b))
 }
 
+/// Perform `lhs == rhs` operation on two [`BooleanArray`]s and include validities in comparison.
+pub fn eq_and_validity(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| !(a ^ b));
+
+    eq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `lhs == rhs` operation on a [`BooleanArray`] and a scalar value.
 pub fn eq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
     if rhs {
@@ -47,16 +59,50 @@ pub fn eq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
     }
 }
 
+/// Perform `lhs == rhs` operation on a [`BooleanArray`] and a scalar value and include validities in comparison.
+pub fn eq_scalar_and_validity(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    if rhs {
+        eq_validities(lhs.clone(), validity, None)
+    } else {
+        let lhs = lhs.with_validity(None);
+
+        let out = compare_op_scalar(&lhs, rhs, |a, _| !a);
+
+        eq_validities(out, validity, None)
+    }
+}
+
 /// `lhs != rhs` for [`BooleanArray`]
 pub fn neq(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| a ^ b)
 }
 
+/// `lhs != rhs` for [`BooleanArray`] and include validities in comparison.
+pub fn neq_and_validity(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| a ^ b);
+
+    neq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `left != right` operation on an array and a scalar value.
 pub fn neq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
     eq_scalar(lhs, !rhs)
 }
 
+/// Perform `left != right` operation on an array and a scalar value.
+pub fn neq_scalar_and_validity(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = eq_scalar(&lhs, !rhs);
+    neq_validities(out, validity, None)
+}
+
 /// Perform `left < right` operation on two arrays.
 pub fn lt(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| !a & b)

src/compute/comparison/mod.rs

@@ -57,6 +57,8 @@ mod simd;
 pub use simd::{Simd8, Simd8Lanes, Simd8PartialEq, Simd8PartialOrd};
 
 use super::take::take_boolean;
+use crate::bitmap::Bitmap;
+use crate::compute;
 pub(crate) use primitive::{
     compare_values_op as primitive_compare_values_op,
     compare_values_op_scalar as primitive_compare_values_op_scalar,
@@ -457,3 +459,54 @@ fn can_partial_eq_scalar(data_type: &DataType) -> bool {
                 | DataType::Interval(IntervalUnit::MonthDayNano)
         )
 }
+
+fn eq_validities(
+    output_without_validities: BooleanArray,
+    validity_lhs: Option<Bitmap>,
+    validity_rhs: Option<Bitmap>,
+) -> BooleanArray {
+    match (validity_lhs, validity_rhs) {
+        (None, None) => output_without_validities,
+        (Some(lhs), None) => compute::boolean::and(
+            &BooleanArray::from_data(DataType::Boolean, lhs, None),
+            &output_without_validities,
+        )
+        .unwrap(),
+        (None, Some(rhs)) => compute::boolean::and(
+            &output_without_validities,
+            &BooleanArray::from_data(DataType::Boolean, rhs, None),
+        )
+        .unwrap(),
+        (Some(lhs), Some(rhs)) => {
+            let lhs = BooleanArray::from_data(DataType::Boolean, lhs, None);
+            let rhs = BooleanArray::from_data(DataType::Boolean, rhs, None);
+            let eq_validities = compute::comparison::boolean::eq(&lhs, &rhs);
+            compute::boolean::and(&output_without_validities, &eq_validities).unwrap()
+        }
+    }
+}
+fn neq_validities(
+    output_without_validities: BooleanArray,
+    validity_lhs: Option<Bitmap>,
+    validity_rhs: Option<Bitmap>,
+) -> BooleanArray {
+    match (validity_lhs, validity_rhs) {
+        (None, None) => output_without_validities,
+        (Some(lhs), None) => {
+            let lhs_negated =
+                compute::boolean::not(&BooleanArray::from_data(DataType::Boolean, lhs, None));
+            compute::boolean::or(&lhs_negated, &output_without_validities).unwrap()
+        }
+        (None, Some(rhs)) => {
+            let rhs_negated =
+                compute::boolean::not(&BooleanArray::from_data(DataType::Boolean, rhs, None));
+            compute::boolean::or(&output_without_validities, &rhs_negated).unwrap()
+        }
+        (Some(lhs), Some(rhs)) => {
+            let lhs = BooleanArray::from_data(DataType::Boolean, lhs, None);
+            let rhs = BooleanArray::from_data(DataType::Boolean, rhs, None);
+            let neq_validities = compute::comparison::boolean::neq(&lhs, &rhs);
+            compute::boolean::or(&output_without_validities, &neq_validities).unwrap()
+        }
+    }
+}

src/compute/comparison/primitive.rs

@@ -1,4 +1,5 @@
 //! Comparison functions for [`PrimitiveArray`]
+use crate::compute::comparison::{eq_validities, neq_validities};
 use crate::{
     array::{BooleanArray, PrimitiveArray},
     bitmap::MutableBitmap,
@@ -99,6 +100,21 @@ where
     compare_op(lhs, rhs, |a, b| a.eq(b))
 }
 
+/// Perform `lhs == rhs` operation on two arrays and include validities in comparison.
+pub fn eq_and_validity<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> BooleanArray
+where
+    T: NativeType + Simd8,
+    T::Simd: Simd8PartialEq,
+{
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| a.eq(b));
+
+    eq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `left == right` operation on an array and a scalar value.
 pub fn eq_scalar<T>(lhs: &PrimitiveArray<T>, rhs: T) -> BooleanArray
 where
@@ -108,6 +124,19 @@ where
     compare_op_scalar(lhs, rhs, |a, b| a.eq(b))
 }
 
+/// Perform `left == right` operation on an array and a scalar value and include validities in comparison.
+pub fn eq_scalar_and_validity<T>(lhs: &PrimitiveArray<T>, rhs: T) -> BooleanArray
+where
+    T: NativeType + Simd8,
+    T::Simd: Simd8PartialEq,
+{
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = compare_op_scalar(&lhs, rhs, |a, b| a.eq(b));
+
+    eq_validities(out, validity, None)
+}
+
 /// Perform `left != right` operation on two arrays.
 pub fn neq<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> BooleanArray
 where
@@ -117,6 +146,21 @@ where
     compare_op(lhs, rhs, |a, b| a.neq(b))
 }
 
+/// Perform `left != right` operation on two arrays and include validities in comparison.
+pub fn neq_and_validity<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> BooleanArray
+where
+    T: NativeType + Simd8,
+    T::Simd: Simd8PartialEq,
+{
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| a.neq(b));
+
+    neq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `left != right` operation on an array and a scalar value.
 pub fn neq_scalar<T>(lhs: &PrimitiveArray<T>, rhs: T) -> BooleanArray
 where
@@ -126,6 +170,19 @@ where
     compare_op_scalar(lhs, rhs, |a, b| a.neq(b))
 }
 
+/// Perform `left != right` operation on an array and a scalar value and include validities in comparison.
+pub fn neq_scalar_and_validity<T>(lhs: &PrimitiveArray<T>, rhs: T) -> BooleanArray
+where
+    T: NativeType + Simd8,
+    T::Simd: Simd8PartialEq,
+{
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = compare_op_scalar(&lhs, rhs, |a, b| a.neq(b));
+
+    neq_validities(out, validity, None)
+}
+
 /// Perform `left < right` operation on two arrays.
 pub fn lt<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> BooleanArray
 where

src/compute/comparison/utf8.rs

@@ -1,4 +1,5 @@
 //! Comparison functions for [`Utf8Array`]
+use crate::compute::comparison::{eq_validities, neq_validities};
 use crate::{
     array::{BooleanArray, Offset, Utf8Array},
     bitmap::Bitmap,
@@ -46,11 +47,51 @@ pub fn eq<O: Offset>(lhs: &Utf8Array<O>, rhs: &Utf8Array<O>) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| a == b)
 }
 
+/// Perform `lhs == rhs` operation on [`Utf8Array`] and include validities in comparison.
+pub fn eq_and_validity<O: Offset>(lhs: &Utf8Array<O>, rhs: &Utf8Array<O>) -> BooleanArray {
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| a == b);
+
+    eq_validities(out, validity_lhs, validity_rhs)
+}
+
+/// Perform `lhs != rhs` operation on [`Utf8Array`] and include validities in comparison.
+pub fn neq_and_validity<O: Offset>(lhs: &Utf8Array<O>, rhs: &Utf8Array<O>) -> BooleanArray {
+    let validity_lhs = lhs.validity().cloned();
+    let validity_rhs = rhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let rhs = rhs.with_validity(None);
+    let out = compare_op(&lhs, &rhs, |a, b| a != b);
+
+    neq_validities(out, validity_lhs, validity_rhs)
+}
+
 /// Perform `lhs == rhs` operation on [`Utf8Array`] and a scalar.
 pub fn eq_scalar<O: Offset>(lhs: &Utf8Array<O>, rhs: &str) -> BooleanArray {
     compare_op_scalar(lhs, rhs, |a, b| a == b)
 }
 
+/// Perform `lhs == rhs` operation on [`Utf8Array`] and a scalar. Also includes null values in comparisson.
+pub fn eq_scalar_and_validity<O: Offset>(lhs: &Utf8Array<O>, rhs: &str) -> BooleanArray {
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = compare_op_scalar(&lhs, rhs, |a, b| a == b);
+
+    eq_validities(out, validity, None)
+}
+
+/// Perform `lhs != rhs` operation on [`Utf8Array`] and a scalar. Also includes null values in comparisson.
+pub fn neq_scalar_and_validity<O: Offset>(lhs: &Utf8Array<O>, rhs: &str) -> BooleanArray {
+    let validity = lhs.validity().cloned();
+    let lhs = lhs.with_validity(None);
+    let out = compare_op_scalar(&lhs, rhs, |a, b| a != b);
+
+    neq_validities(out, validity, None)
+}
+
 /// Perform `lhs != rhs` operation on [`Utf8Array`].
 pub fn neq<O: Offset>(lhs: &Utf8Array<O>, rhs: &Utf8Array<O>) -> BooleanArray {
     compare_op(lhs, rhs, |a, b| a != b)