From ef7937dfe56033c2cc491482c67587b52cd91554 Mon Sep 17 00:00:00 2001 From: Steven Date: Sun, 2 Jan 2022 16:27:57 +0800 Subject: [PATCH] Added `and_scalar` and `or_scalar` for boolean kleene. (#723) --- src/compute/boolean_kleene.rs | 72 +++++++++++++++++++++++-- tests/it/compute/boolean_kleene.rs | 85 ++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+), 3 deletions(-) diff --git a/src/compute/boolean_kleene.rs b/src/compute/boolean_kleene.rs index 350cf6b8bd0..ec5d7e0db95 100644 --- a/src/compute/boolean_kleene.rs +++ b/src/compute/boolean_kleene.rs @@ -1,12 +1,13 @@ //! Boolean operators of [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics). use crate::datatypes::DataType; use crate::error::{ArrowError, Result}; +use crate::scalar::BooleanScalar; use crate::{ array::BooleanArray, - bitmap::{quaternary, ternary}, + bitmap::{binary, quaternary, ternary, unary, Bitmap, MutableBitmap}, }; -/// Logical 'or' with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics) +/// Logical 'or' operation on two arrays with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics) /// # Errors /// This function errors if the operands have different lengths. /// # Example @@ -96,7 +97,7 @@ pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> Result { )) } -/// Logical 'and' with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics) +/// Logical 'and' operation on two arrays with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics) /// # Errors /// This function errors if the operands have different lengths. /// # Example @@ -184,3 +185,68 @@ pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> Result { validity, )) } + +/// Logical 'or' operation on an array and a scalar value with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics) +/// # Example +/// +/// ```rust +/// use arrow2::array::BooleanArray; +/// use arrow2::scalar::BooleanScalar; +/// use arrow2::compute::boolean_kleene::or_scalar; +/// # fn main() { +/// let array = BooleanArray::from(&[Some(true), Some(false), None]); +/// let scalar = BooleanScalar::new(Some(false)); +/// let result = or_scalar(&array, &scalar); +/// assert_eq!(result, BooleanArray::from(&[Some(true), Some(false), None])); +/// # } +/// ``` +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(), None) + } + Some(false) => array.clone(), + None => { + let values = array.values(); + let validity = match array.validity() { + Some(validity) => binary(values, validity, |value, validity| validity & value), + None => unary(values, |value| value), + }; + BooleanArray::from_data(DataType::Boolean, values.clone(), Some(validity)) + } + } +} + +/// Logical 'and' operation on an array and a scalar value with [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics) +/// # Example +/// +/// ```rust +/// use arrow2::array::BooleanArray; +/// use arrow2::scalar::BooleanScalar; +/// use arrow2::compute::boolean_kleene::and_scalar; +/// # fn main() { +/// let array = BooleanArray::from(&[Some(true), Some(false), None]); +/// let scalar = BooleanScalar::new(None); +/// let result = and_scalar(&array, &scalar); +/// assert_eq!(result, BooleanArray::from(&[None, Some(false), None])); +/// # } +/// ``` +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, None) + } + None => { + let values = array.values(); + let validity = match array.validity() { + Some(validity) => binary(values, validity, |value, validity| validity & !value), + None => unary(values, |value| !value), + }; + BooleanArray::from_data(DataType::Boolean, array.values().clone(), Some(validity)) + } + } +} diff --git a/tests/it/compute/boolean_kleene.rs b/tests/it/compute/boolean_kleene.rs index 09e2fc501a3..3bf0b8ef37e 100644 --- a/tests/it/compute/boolean_kleene.rs +++ b/tests/it/compute/boolean_kleene.rs @@ -1,5 +1,6 @@ use arrow2::array::BooleanArray; use arrow2::compute::boolean_kleene::*; +use arrow2::scalar::BooleanScalar; #[test] fn and_generic() { @@ -129,3 +130,87 @@ fn or_left_nulls() { assert_eq!(c, expected); } + +#[test] +fn array_and_true() { + let array = BooleanArray::from(&[Some(true), Some(false), None, Some(true), Some(false), None]); + + let scalar = BooleanScalar::new(Some(true)); + let result = and_scalar(&array, &scalar); + + // Should be same as argument array if scalar is true. + assert_eq!(result, array); +} + +#[test] +fn array_and_false() { + let array = BooleanArray::from(&[Some(true), Some(false), None, Some(true), Some(false), None]); + + let scalar = BooleanScalar::new(Some(false)); + let result = and_scalar(&array, &scalar); + + let expected = BooleanArray::from(&[ + Some(false), + Some(false), + Some(false), + Some(false), + Some(false), + Some(false), + ]); + + assert_eq!(result, expected); +} + +#[test] +fn array_and_none() { + let array = BooleanArray::from(&[Some(true), Some(false), None, Some(true), Some(false), None]); + + let scalar = BooleanScalar::new(None); + let result = and_scalar(&array, &scalar); + + let expected = BooleanArray::from(&[None, Some(false), None, None, Some(false), None]); + + assert_eq!(result, expected); +} + +#[test] +fn array_or_true() { + let array = BooleanArray::from(&[Some(true), Some(false), None, Some(true), Some(false), None]); + + let scalar = BooleanScalar::new(Some(true)); + let result = or_scalar(&array, &scalar); + + let expected = BooleanArray::from(&[ + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + Some(true), + ]); + + assert_eq!(result, expected); +} + +#[test] +fn array_or_false() { + let array = BooleanArray::from(&[Some(true), Some(false), None, Some(true), Some(false), None]); + + let scalar = BooleanScalar::new(Some(false)); + let result = or_scalar(&array, &scalar); + + // Should be same as argument array if scalar is false. + assert_eq!(result, array); +} + +#[test] +fn array_or_none() { + let array = BooleanArray::from(&[Some(true), Some(false), None, Some(true), Some(false), None]); + + let scalar = BooleanScalar::new(None); + let result = or_scalar(&array, &scalar); + + let expected = BooleanArray::from(&[Some(true), None, None, Some(true), None, None]); + + assert_eq!(result, expected); +}