Repository: jorgecarleitao/arrow2
Pull request #1243 (2022-09-05 07:55:14)

fixed comparisson and validity kernels
When the values masked out by the validity bits were different the `comparison_and_validty` kernels were incorrect. This PR fixes that.

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1243](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (032d740) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/a7428e0d386f37cdf85d4aa8f765b7d34f05eb37?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (a7428e0) will **increase** coverage by `0.00%`.
> The diff coverage is `91.20%`.

```diff
@@           Coverage Diff           @@
##             main    #1243   +/-   ##
=======================================
  Coverage   83.32%   83.32%           
=======================================
  Files         359      359           
  Lines       37854    37939   +85     
=======================================
+ Hits        31541    31612   +71     
- Misses       6313     6327   +14     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/compute/boolean.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2NvbXB1dGUvYm9vbGVhbi5ycw==) | `92.59% <85.00%> (-4.63%)` | :arrow_down: |
| [src/compute/comparison/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2NvbXB1dGUvY29tcGFyaXNvbi9tb2QucnM=) | `65.21% <96.07%> (+11.14%)` | :arrow_up: |
| [src/io/ipc/read/stream\_async.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3N0cmVhbV9hc3luYy5ycw==) | `74.65% <0.00%> (-2.06%)` | :arrow_down: |
| [src/io/ipc/read/file\_async.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL2ZpbGVfYXN5bmMucnM=) | `60.44% <0.00%> (-0.75%)` | :arrow_down: |
| [src/io/ipc/read/file.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1243/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL2ZpbGUucnM=) | `96.87% <0.00%> (-0.45%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)


Patch:
diff --git a/src/compute/boolean.rs b/src/compute/boolean.rs
index 46d7266954..8a17eeaba3 100644
--- a/src/compute/boolean.rs
+++ b/src/compute/boolean.rs
@@ -6,17 +6,24 @@ use crate::scalar::BooleanScalar;
 
 use super::utils::combine_validities;
 
-/// Helper function to implement binary kernels
-fn binary_boolean_kernel<F>(lhs: &BooleanArray, rhs: &BooleanArray, op: F) -> BooleanArray
-where
-    F: Fn(&Bitmap, &Bitmap) -> Bitmap,
-{
+fn assert_lengths(lhs: &BooleanArray, rhs: &BooleanArray) {
     assert_eq!(
         lhs.len(),
         rhs.len(),
         "lhs and rhs must have the same length"
     );
+}
 
+/// Helper function to implement binary kernels
+pub(crate) fn binary_boolean_kernel<F>(
+    lhs: &BooleanArray,
+    rhs: &BooleanArray,
+    op: F,
+) -> BooleanArray
+where
+    F: Fn(&Bitmap, &Bitmap) -> Bitmap,
+{
+    assert_lengths(lhs, rhs);
     let validity = combine_validities(lhs.validity(), rhs.validity());
 
     let left_buffer = lhs.values();
@@ -41,6 +48,31 @@ where
 /// assert_eq!(and_ab, BooleanArray::from(&[Some(false), Some(true), None]));
 /// ```
 pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
+    if lhs.null_count() == 0 && rhs.null_count() == 0 {
+        let left_buffer = lhs.values();
+        let right_buffer = rhs.values();
+
+        match (left_buffer.unset_bits(), right_buffer.unset_bits()) {
+            // all values are `true` on both sides
+            (0, 0) => {
+                assert_lengths(lhs, rhs);
+                return lhs.clone();
+            }
+            // all values are `false` on left side
+            (l, _) if l == lhs.len() => {
+                assert_lengths(lhs, rhs);
+                return lhs.clone();
+            }
+            // all values are `false` on right side
+            (_, r) if r == rhs.len() => {
+                assert_lengths(lhs, rhs);
+                return rhs.clone();
+            }
+            // ignore the rest
+            _ => {}
+        }
+    }
+
     binary_boolean_kernel(lhs, rhs, |lhs, rhs| lhs & rhs)
 }
 
@@ -58,6 +90,31 @@ pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
 /// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), Some(true), None]));
 /// ```
 pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
+    if lhs.null_count() == 0 && rhs.null_count() == 0 {
+        let left_buffer = lhs.values();
+        let right_buffer = rhs.values();
+
+        match (left_buffer.unset_bits(), right_buffer.unset_bits()) {
+            // all values are `true` on left side
+            (0, _) => {
+                assert_lengths(lhs, rhs);
+                return lhs.clone();
+            }
+            // all values are `true` on right side
+            (_, 0) => {
+                assert_lengths(lhs, rhs);
+                return rhs.clone();
+            }
+            // all values on lhs and rhs are `false`
+            (l, r) if l == lhs.len() && r == rhs.len() => {
+                assert_lengths(lhs, rhs);
+                return rhs.clone();
+            }
+            // ignore the rest
+            _ => {}
+        }
+    }
+
     binary_boolean_kernel(lhs, rhs, |lhs, rhs| lhs | rhs)
 }
 
diff --git a/src/compute/comparison/mod.rs b/src/compute/comparison/mod.rs
index 4ad7920705..ba6a84e49b 100644
--- a/src/compute/comparison/mod.rs
+++ b/src/compute/comparison/mod.rs
@@ -57,7 +57,7 @@ mod simd;
 pub use simd::{Simd8, Simd8Lanes, Simd8PartialEq, Simd8PartialOrd};
 
 use super::take::take_boolean;
-use crate::bitmap::Bitmap;
+use crate::bitmap::{binary, Bitmap};
 use crate::compute;
 pub(crate) use primitive::{
     compare_values_op as primitive_compare_values_op,
@@ -521,10 +521,43 @@ fn finish_eq_validities(
             &BooleanArray::new(DataType::Boolean, rhs, None),
         ),
         (Some(lhs), Some(rhs)) => {
+            let lhs_validity_unset_bits = lhs.unset_bits();
+            let rhs_validity_unset_bits = rhs.unset_bits();
+
+            // this branch is a bit more complicated as both arrays can have masked out values
+            // these masked out values might differ and lead to a `eq == false` that has to
+            // be corrected as both should be `null == null = true`
+
             let lhs = BooleanArray::new(DataType::Boolean, lhs, None);
             let rhs = BooleanArray::new(DataType::Boolean, rhs, None);
             let eq_validities = compute::comparison::boolean::eq(&lhs, &rhs);
-            compute::boolean::and(&output_without_validities, &eq_validities)
+
+            // validity_bits are equal AND values are equal
+            let equal = compute::boolean::and(&output_without_validities, &eq_validities);
+
+            match (lhs_validity_unset_bits, rhs_validity_unset_bits) {
+                // there is at least one side with all values valid
+                // so we don't have to correct.
+                (0, _) | (_, 0) => equal,
+                _ => {
+                    // we use the binary kernel here to save allocations
+                    // and apply `!(lhs | rhs)` in one step
+                    let both_sides_invalid =
+                        compute::boolean::binary_boolean_kernel(&lhs, &rhs, |lhs, rhs| {
+                            binary(lhs, rhs, |lhs, rhs| !(lhs | rhs))
+                        });
+                    // this still might include incorrect masked out values
+                    // under the validity bits, so we must correct for that
+
+                    // if not all true, e.g. at least one is set.
+                    // then we propagate that null as `true` in equality
+                    if both_sides_invalid.values().unset_bits() != both_sides_invalid.len() {
+                        compute::boolean::or(&equal, &both_sides_invalid)
+                    } else {
+                        equal
+                    }
+                }
+            }
         }
     }
 }
@@ -547,10 +580,46 @@ fn finish_neq_validities(
             compute::boolean::or(&output_without_validities, &rhs_negated)
         }
         (Some(lhs), Some(rhs)) => {
+            let lhs_validity_unset_bits = lhs.unset_bits();
+            let rhs_validity_unset_bits = rhs.unset_bits();
+
+            // this branch is a bit more complicated as both arrays can have masked out values
+            // these masked out values might differ and lead to a `neq == true` that has to
+            // be corrected as both should be `null != null = false`
             let lhs = BooleanArray::new(DataType::Boolean, lhs, None);
             let rhs = BooleanArray::new(DataType::Boolean, rhs, None);
             let neq_validities = compute::comparison::boolean::neq(&lhs, &rhs);
-            compute::boolean::or(&output_without_validities, &neq_validities)
+
+            // validity_bits are not equal OR values not equal
+            let or = compute::boolean::or(&output_without_validities, &neq_validities);
+
+            match (lhs_validity_unset_bits, rhs_validity_unset_bits) {
+                // there is at least one side with all values valid
+                // so we don't have to correct.
+                (0, _) | (_, 0) => or,
+                _ => {
+                    // we use the binary kernel here to save allocations
+                    // and apply `!(lhs | rhs)` in one step
+                    let both_sides_invalid =
+                        compute::boolean::binary_boolean_kernel(&lhs, &rhs, |lhs, rhs| {
+                            binary(lhs, rhs, |lhs, rhs| !(lhs | rhs))
+                        });
+                    // this still might include incorrect masked out values
+                    // under the validity bits, so we must correct for that
+
+                    // if not all true, e.g. at least one is set.
+                    // then we propagate that null as `false` as the nulls are equal
+                    if both_sides_invalid.values().unset_bits() != both_sides_invalid.len() {
+                        // we use the `binary` kernel directly to save allocations
+                        // and apply `lhs & !rhs)` in one shot.
+                        compute::boolean::binary_boolean_kernel(&lhs, &rhs, |lhs, rhs| {
+                            binary(lhs, rhs, |lhs, rhs| (lhs & !rhs))
+                        })
+                    } else {
+                        or
+                    }
+                }
+            }
         }
     }
 }


Test patch:
diff --git a/tests/it/compute/comparison.rs b/tests/it/compute/comparison.rs
index 09f3fd6145..b077e99610 100644
--- a/tests/it/compute/comparison.rs
+++ b/tests/it/compute/comparison.rs
@@ -1,7 +1,7 @@
 use arrow2::array::*;
 use arrow2::bitmap::Bitmap;
-use arrow2::compute::comparison::{self, boolean::*, primitive};
-use arrow2::datatypes::{DataType::*, IntegerType, IntervalUnit, TimeUnit};
+use arrow2::compute::comparison::{self, boolean::*, primitive, utf8};
+use arrow2::datatypes::{DataType, DataType::*, IntegerType, IntervalUnit, TimeUnit};
 use arrow2::scalar::new_scalar;
 
 #[test]
@@ -381,3 +381,23 @@ fn primitive_gt_eq() {
         BooleanArray::from([Some(true), Some(false), Some(true), None, None])
     )
 }
+
+#[test]
+#[cfg(any(feature = "compute_cast", feature = "compute_boolean_kleene"))]
+fn utf8_and_validity() {
+    use arrow2::compute::cast::CastOptions;
+    let a1 = Utf8Array::<i32>::from([Some("0"), Some("1"), None, Some("2")]);
+    let a2 = Int32Array::from([Some(0), Some(1), None, Some(2)]);
+
+    // due to the cast the values underneath the validity bits differ
+    let a2 = arrow2::compute::cast::cast(&a2, &DataType::Utf8, CastOptions::default()).unwrap();
+    let a2 = a2.as_any().downcast_ref::<Utf8Array<i32>>().unwrap();
+
+    let expected = BooleanArray::from_slice([true, true, true, true]);
+    assert_eq!(utf8::eq_and_validity(&a1, &a1), expected);
+    assert_eq!(utf8::eq_and_validity(&a1, a2), expected);
+
+    let expected = BooleanArray::from_slice([false, false, false, false]);
+    assert_eq!(utf8::neq_and_validity(&a1, &a1), expected);
+    assert_eq!(utf8::neq_and_validity(&a1, a2), expected);
+}