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

[alamb]

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

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


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

commit 434e3f71a8ff6d1c6f8e8d7db63e697643bbf3cc
Author: Carlos <wxd963996...@gmail.com>
AuthorDate: Sat Nov 27 19:21:29 2021 +0800

    Add boolean comparison to scalar kernels for less then, greater than (#977)
---
 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,