[arrow-rs] branch active_release updated: Add boolean comparison to scalar kernels for less then, greater than (#977) (#1005)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch active_release
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/active_release by this push:
     new 6c8b293  Add boolean comparison to scalar kernels for less then, greater than (#977) (#1005)
6c8b293 is described below

commit 6c8b2936d7b07e1e2f5d1d48eea425a385382dfb
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Mon Dec 6 15:39:42 2021 -0500

    Add boolean comparison to scalar kernels for less then, greater than (#977) (#1005)
    
    Co-authored-by: Carlos <wx...@gmail.com>
---
 arrow/src/compute/kernels/comparison.rs | 81 +++++++++++++++++++++++++++++++--
 1 file changed, 77 insertions(+), 4 deletions(-)

diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs
index 3b65f33..125be91 100644
--- a/arrow/src/compute/kernels/comparison.rs
+++ b/arrow/src/compute/kernels/comparison.rs
@@ -781,6 +781,26 @@ pub fn eq_bool_scalar(left: &BooleanArray, right: bool) -> Result<BooleanArray>
     Ok(BooleanArray::from(data))
 }
 
+/// Perform `left < right` operation on [`BooleanArray`] and a scalar
+pub fn lt_bool_scalar(left: &BooleanArray, right: bool) -> Result<BooleanArray> {
+    compare_op_scalar!(left, right, |a: bool, b: bool| !a & b)
+}
+
+/// Perform `left <= right` operation on [`BooleanArray`] and a scalar
+pub fn lt_eq_bool_scalar(left: &BooleanArray, right: bool) -> Result<BooleanArray> {
+    compare_op_scalar!(left, right, |a, b| a <= b)
+}
+
+/// Perform `left > right` operation on [`BooleanArray`] and a scalar
+pub fn gt_bool_scalar(left: &BooleanArray, right: bool) -> Result<BooleanArray> {
+    compare_op_scalar!(left, right, |a: bool, b: bool| a & !b)
+}
+
+/// Perform `left >= right` operation on [`BooleanArray`] and a scalar
+pub fn gt_eq_bool_scalar(left: &BooleanArray, right: bool) -> Result<BooleanArray> {
+    compare_op_scalar!(left, right, |a, b| a >= b)
+}
+
 /// Perform `left != right` operation on [`BooleanArray`] and a scalar
 pub fn neq_bool_scalar(left: &BooleanArray, right: bool) -> Result<BooleanArray> {
     eq_bool_scalar(left, !right)
@@ -897,10 +917,7 @@ where
     let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
 
     // this is currently the case for all our datatypes and allows us to always append full bytes
-    assert!(
-        lanes % 8 == 0,
-        "Number of vector lanes must be multiple of 8"
-    );
+    assert_eq!(lanes % 8, 0, "Number of vector lanes must be multiple of 8");
     let mut left_chunks = left.values().chunks_exact(lanes);
     let mut right_chunks = right.values().chunks_exact(lanes);
 
@@ -1688,6 +1705,62 @@ mod tests {
     }
 
     #[test]
+    fn test_boolean_array_lt_scalar() {
+        let a: BooleanArray = vec![Some(true), Some(false), None].into();
+
+        let res1: Vec<Option<bool>> = lt_bool_scalar(&a, false).unwrap().iter().collect();
+
+        assert_eq!(res1, vec![Some(false), Some(false), None]);
+
+        let res2: Vec<Option<bool>> = lt_bool_scalar(&a, true).unwrap().iter().collect();
+
+        assert_eq!(res2, vec![Some(false), Some(true), None]);
+    }
+
+    #[test]
+    fn test_boolean_array_lt_eq_scalar() {
+        let a: BooleanArray = vec![Some(true), Some(false), None].into();
+
+        let res1: Vec<Option<bool>> =
+            lt_eq_bool_scalar(&a, false).unwrap().iter().collect();
+
+        assert_eq!(res1, vec![Some(false), Some(true), None]);
+
+        let res2: Vec<Option<bool>> =
+            lt_eq_bool_scalar(&a, true).unwrap().iter().collect();
+
+        assert_eq!(res2, vec![Some(true), Some(true), None]);
+    }
+
+    #[test]
+    fn test_boolean_array_gt_scalar() {
+        let a: BooleanArray = vec![Some(true), Some(false), None].into();
+
+        let res1: Vec<Option<bool>> = gt_bool_scalar(&a, false).unwrap().iter().collect();
+
+        assert_eq!(res1, vec![Some(true), Some(false), None]);
+
+        let res2: Vec<Option<bool>> = gt_bool_scalar(&a, true).unwrap().iter().collect();
+
+        assert_eq!(res2, vec![Some(false), Some(false), None]);
+    }
+
+    #[test]
+    fn test_boolean_array_gt_eq_scalar() {
+        let a: BooleanArray = vec![Some(true), Some(false), None].into();
+
+        let res1: Vec<Option<bool>> =
+            gt_eq_bool_scalar(&a, false).unwrap().iter().collect();
+
+        assert_eq!(res1, vec![Some(true), Some(true), None]);
+
+        let res2: Vec<Option<bool>> =
+            gt_eq_bool_scalar(&a, true).unwrap().iter().collect();
+
+        assert_eq!(res2, vec![Some(true), Some(false), None]);
+    }
+
+    #[test]
     fn test_primitive_array_lt() {
         cmp_i64!(
             lt,