Speedup scalar boolean operations (#546)

jorgecarleitao · Oct 22, 2021 · b424e90 · b424e90
1 parent b62184d
commit b424e90
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 8 deletions.
diff --git a/benches/comparison_kernels.rs b/benches/comparison_kernels.rs
@@ -40,7 +40,7 @@ fn add_benchmark(c: &mut Criterion) {
             b.iter(|| bench_op(&arr_a, &arr_b, Operator::Eq))
         });
         c.bench_function(&format!("bool scalar 2^{}", log2_size), |b| {
-            b.iter(|| bench_op_scalar(&arr_a, &BooleanScalar::from(Some(true)), Operator::Eq))
+            b.iter(|| bench_op_scalar(&arr_a, &BooleanScalar::from(Some(false)), Operator::Eq))
         });
 
         let arr_a = create_string_array::<i32>(size, 4, 0.1, 42);

diff --git a/src/compute/comparison/boolean.rs b/src/compute/comparison/boolean.rs
@@ -82,7 +82,11 @@ pub fn eq(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
 
 /// Perform `left == right` operation on an array and a scalar value.
 pub fn eq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
-    compare_op_scalar(lhs, rhs, |a, b| !(a ^ b))
+    if rhs {
+        lhs.clone()
+    } else {
+        compare_op_scalar(lhs, rhs, |a, _| !a)
+    }
 }
 
 /// Perform `left != right` operation on two arrays.
@@ -92,7 +96,7 @@ pub fn neq(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
 
 /// Perform `left != right` operation on an array and a scalar value.
 pub fn neq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
-    compare_op_scalar(lhs, rhs, |a, b| a ^ b)
+    eq_scalar(lhs, !rhs)
 }
 
 /// Perform `left < right` operation on two arrays.
@@ -102,7 +106,15 @@ pub fn lt(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
 
 /// Perform `left < right` operation on an array and a scalar value.
 pub fn lt_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
-    compare_op_scalar(lhs, rhs, |a, b| !a & b)
+    if rhs {
+        compare_op_scalar(lhs, rhs, |a, _| !a)
+    } else {
+        BooleanArray::from_data(
+            DataType::Boolean,
+            Bitmap::new_zeroed(lhs.len()),
+            lhs.validity().cloned(),
+        )
+    }
 }
 
 /// Perform `left <= right` operation on two arrays.
@@ -113,7 +125,11 @@ pub fn lt_eq(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
 /// Perform `left <= right` operation on an array and a scalar value.
 /// Null values are less than non-null values.
 pub fn lt_eq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
-    compare_op_scalar(lhs, rhs, |a, b| !a | b)
+    if rhs {
+        compare_op_scalar(lhs, rhs, |_, _| 0b11111111)
+    } else {
+        compare_op_scalar(lhs, rhs, |a, _| !a)
+    }
 }
 
 /// Perform `left > right` operation on two arrays. Non-null values are greater than null
@@ -125,7 +141,15 @@ pub fn gt(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
 /// Perform `left > right` operation on an array and a scalar value.
 /// Non-null values are greater than null values.
 pub fn gt_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
-    compare_op_scalar(lhs, rhs, |a, b| a & !b)
+    if rhs {
+        BooleanArray::from_data(
+            DataType::Boolean,
+            Bitmap::new_zeroed(lhs.len()),
+            lhs.validity().cloned(),
+        )
+    } else {
+        lhs.clone()
+    }
 }
 
 /// Perform `left >= right` operation on two arrays. Non-null values are greater than null
@@ -137,7 +161,11 @@ pub fn gt_eq(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
 /// Perform `left >= right` operation on an array and a scalar value.
 /// Non-null values are greater than null values.
 pub fn gt_eq_scalar(lhs: &BooleanArray, rhs: bool) -> BooleanArray {
-    compare_op_scalar(lhs, rhs, |a, b| a | !b)
+    if rhs {
+        lhs.clone()
+    } else {
+        compare_op_scalar(lhs, rhs, |_, _| 0b11111111)
+    }
 }
 
 /// Compare two [`BooleanArray`]s using the given [`Operator`].
@@ -263,10 +291,45 @@ mod tests {
     }
 
     #[test]
-    fn test_lt_scalar() {
+    fn test_lt_scalar_true() {
         cmp_bool_scalar!(lt_scalar, &[false, true], true, &[true, false]);
     }
 
+    #[test]
+    fn test_lt_scalar_false() {
+        cmp_bool_scalar!(lt_scalar, &[false, true], false, &[false, false]);
+    }
+
+    #[test]
+    fn test_lt_eq_scalar_true() {
+        cmp_bool_scalar!(lt_eq_scalar, &[false, true], true, &[true, true]);
+    }
+
+    #[test]
+    fn test_lt_eq_scalar_false() {
+        cmp_bool_scalar!(lt_eq_scalar, &[false, true], false, &[true, false]);
+    }
+
+    #[test]
+    fn test_gt_scalar_true() {
+        cmp_bool_scalar!(gt_scalar, &[false, true], true, &[false, false]);
+    }
+
+    #[test]
+    fn test_gt_scalar_false() {
+        cmp_bool_scalar!(gt_scalar, &[false, true], false, &[false, true]);
+    }
+
+    #[test]
+    fn test_gt_eq_scalar_true() {
+        cmp_bool_scalar!(gt_eq_scalar, &[false, true], true, &[false, true]);
+    }
+
+    #[test]
+    fn test_gt_eq_scalar_false() {
+        cmp_bool_scalar!(gt_eq_scalar, &[false, true], false, &[true, true]);
+    }
+
     #[test]
     fn eq_nulls() {
         cmp_bool_options!(