diff --git a/src/compute/boolean.rs b/src/compute/boolean.rs index bd6b9be8e6d..4ef5b46b4b0 100644 --- 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 index 5468909df10..af77d98ed5b 100644 --- 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); +}