Add and_scalar and or_scalar for boolean (#707)

jorgecarleitao · Dec 25, 2021 · 64043a0 · 64043a0
1 parent 87a08f4
commit 64043a0
Show file tree

Hide file tree

Showing 2 changed files with 134 additions and 0 deletions.
diff --git a/src/compute/boolean.rs b/src/compute/boolean.rs
@@ -3,6 +3,7 @@ use crate::array::{Array, BooleanArray};
 use crate::bitmap::{Bitmap, MutableBitmap};
 use crate::datatypes::DataType;
 use crate::error::{ArrowError, Result};
+use crate::scalar::BooleanScalar;
 
 use super::utils::combine_validities;
 
@@ -133,3 +134,54 @@ pub fn is_not_null(input: &dyn Array) -> BooleanArray {
     };
     BooleanArray::from_data(DataType::Boolean, values, None)
 }
+
+/// Performs `AND` operation on an array and a scalar value. If either left or right value
+/// is null then the result is also null.
+/// # Example
+/// ```rust
+/// use arrow2::array::BooleanArray;
+/// use arrow2::compute::boolean::and_scalar;
+/// use arrow2::scalar::BooleanScalar;
+/// # fn main() {
+/// let array = BooleanArray::from_slice(&[false, false, true, true]);
+/// let scalar = BooleanScalar::new(Some(true));
+/// let result = and_scalar(&array, &scalar);
+/// assert_eq!(result, BooleanArray::from_slice(&[false, false, true, true]));
+/// # }
+/// ```
+pub fn and_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
+    match scalar.value() {
+        Some(true) => array.clone(),
+        Some(false) => {
+            let values = Bitmap::new_zeroed(array.len());
+            BooleanArray::from_data(DataType::Boolean, values, array.validity().cloned())
+        }
+        None => BooleanArray::new_null(DataType::Boolean, array.len()),
+    }
+}
+
+/// Performs `OR` operation on an array and a scalar value. If either left or right value
+/// is null then the result is also null.
+/// # Example
+/// ```rust
+/// use arrow2::array::BooleanArray;
+/// use arrow2::compute::boolean::or_scalar;
+/// use arrow2::scalar::BooleanScalar;
+/// # fn main() {
+/// let array = BooleanArray::from_slice(&[false, false, true, true]);
+/// let scalar = BooleanScalar::new(Some(true));
+/// let result = or_scalar(&array, &scalar);
+/// assert_eq!(result, BooleanArray::from_slice(&[true, true, true, true]));
+/// # }
+/// ```
+pub fn or_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
+    match scalar.value() {
+        Some(true) => {
+            let mut values = MutableBitmap::new();
+            values.extend_constant(array.len(), true);
+            BooleanArray::from_data(DataType::Boolean, values.into(), array.validity().cloned())
+        }
+        Some(false) => array.clone(),
+        None => BooleanArray::new_null(DataType::Boolean, array.len()),
+    }
+}
diff --git a/tests/it/compute/boolean.rs b/tests/it/compute/boolean.rs
@@ -1,5 +1,6 @@
 use arrow2::array::*;
 use arrow2::compute::boolean::*;
+use arrow2::scalar::BooleanScalar;
 
 #[test]
 fn array_and() {
@@ -336,3 +337,84 @@ fn test_nullable_array_with_offset_is_not_null() {
 
     assert_eq!(expected, res);
 }
+
+#[test]
+fn array_and_scalar() {
+    let array = BooleanArray::from_slice(&[false, false, true, true]);
+
+    let scalar = BooleanScalar::new(Some(true));
+    let real = and_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from_slice(&[false, false, true, true]);
+    assert_eq!(real, expected);
+
+    let scalar = BooleanScalar::new(Some(false));
+    let real = and_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from_slice(&[false, false, false, false]);
+
+    assert_eq!(real, expected);
+}
+
+#[test]
+fn array_and_scalar_validity() {
+    let array = BooleanArray::from(&[None, Some(false), Some(true)]);
+
+    let scalar = BooleanScalar::new(Some(true));
+    let real = and_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from(&[None, Some(false), Some(true)]);
+    assert_eq!(real, expected);
+
+    let scalar = BooleanScalar::new(None);
+    let real = and_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from(&[None; 3]);
+    assert_eq!(real, expected);
+
+    let array = BooleanArray::from_slice(&[true, false, true]);
+    let real = and_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from(&[None; 3]);
+    assert_eq!(real, expected);
+}
+
+#[test]
+fn array_or_scalar() {
+    let array = BooleanArray::from_slice(&[false, false, true, true]);
+
+    let scalar = BooleanScalar::new(Some(true));
+    let real = or_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from_slice(&[true, true, true, true]);
+    assert_eq!(real, expected);
+
+    let scalar = BooleanScalar::new(Some(false));
+    let real = or_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from_slice(&[false, false, true, true]);
+    assert_eq!(real, expected);
+}
+
+#[test]
+fn array_or_scalar_validity() {
+    let array = BooleanArray::from(&[None, Some(false), Some(true)]);
+
+    let scalar = BooleanScalar::new(Some(true));
+    let real = or_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from(&[None, Some(true), Some(true)]);
+    assert_eq!(real, expected);
+
+    let scalar = BooleanScalar::new(None);
+    let real = or_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from(&[None; 3]);
+    assert_eq!(real, expected);
+
+    let array = BooleanArray::from_slice(&[true, false, true]);
+    let real = and_scalar(&array, &scalar);
+
+    let expected = BooleanArray::from(&[None; 3]);
+    assert_eq!(real, expected);
+}