Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Simpler boolean compute API
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorgecarleitao
jorgecarleitao committed Jul 12, 2022
1 parent 0042f4f commit a2d9d1a
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 96 deletions.
69 changes: 28 additions & 41 deletions src/compute/boolean.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,20 @@
use crate::array::{Array, BooleanArray};
use crate::bitmap::{Bitmap, MutableBitmap};
use crate::datatypes::DataType;
use crate::error::{Error, Result};
use crate::scalar::BooleanScalar;

use super::utils::combine_validities;

/// Helper function to implement binary kernels
fn binary_boolean_kernel<F>(lhs: &BooleanArray, rhs: &BooleanArray, op: F) -> Result<BooleanArray>
fn binary_boolean_kernel<F>(lhs: &BooleanArray, rhs: &BooleanArray, op: F) -> BooleanArray
where
F: Fn(&Bitmap, &Bitmap) -> Bitmap,
{
if lhs.len() != rhs.len() {
return Err(Error::InvalidArgumentError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
assert_eq!(
lhs.len(),
rhs.len(),
"lhs and rhs must have the same length"
);

let validity = combine_validities(lhs.validity(), rhs.validity());

Expand All @@ -25,48 +24,40 @@ where

let values = op(left_buffer, right_buffer);

Ok(BooleanArray::new(DataType::Boolean, values, validity))
BooleanArray::new(DataType::Boolean, values, validity)
}

/// Performs `AND` operation on two arrays. If either left or right value is null then the
/// result is also null.
/// # Error
/// This function errors when the arrays have different lengths.
/// # Example
/// Performs `&&` operation on two [`BooleanArray`], combining the validities.
/// # Panics
/// This function panics iff the arrays have different lengths.
/// # Examples
/// ```rust
/// use arrow2::array::BooleanArray;
/// use arrow2::error::Result;
/// use arrow2::compute::boolean::and;
/// # fn main() -> Result<()> {
///
/// let a = BooleanArray::from(&[Some(false), Some(true), None]);
/// let b = BooleanArray::from(&[Some(true), Some(true), Some(false)]);
/// let and_ab = and(&a, &b)?;
/// let and_ab = and(&a, &b);
/// assert_eq!(and_ab, BooleanArray::from(&[Some(false), Some(true), None]));
/// # Ok(())
/// # }
/// ```
pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
binary_boolean_kernel(lhs, rhs, |lhs, rhs| lhs & rhs)
}

/// Performs `OR` operation on two arrays. If either left or right value is null then the
/// result is also null.
/// # Error
/// This function errors when the arrays have different lengths.
/// # Example
/// Performs `||` operation on two [`BooleanArray`], combining the validities.
/// # Panics
/// This function panics iff the arrays have different lengths.
/// # Examples
/// ```rust
/// use arrow2::array::BooleanArray;
/// use arrow2::error::Result;
/// use arrow2::compute::boolean::or;
/// # fn main() -> Result<()> {
///
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
/// let or_ab = or(&a, &b)?;
/// let or_ab = or(&a, &b);
/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), Some(true), None]));
/// # Ok(())
/// # }
/// ```
pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
binary_boolean_kernel(lhs, rhs, |lhs, rhs| lhs | rhs)
}

Expand All @@ -76,21 +67,18 @@ pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
/// ```rust
/// use arrow2::array::BooleanArray;
/// use arrow2::compute::boolean::not;
/// # fn main() {
///
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let not_a = not(&a);
/// assert_eq!(not_a, BooleanArray::from(vec![Some(true), Some(false), None]));
/// # }
/// ```
pub fn not(array: &BooleanArray) -> BooleanArray {
let values = !array.values();
let validity = array.validity().cloned();
BooleanArray::new(DataType::Boolean, values, validity)
}

/// Returns a non-null [BooleanArray] with whether each value of the array is null.
/// # Error
/// This function never errors.
/// Returns a non-null [`BooleanArray`] with whether each value of the array is null.
/// # Example
/// ```rust
/// use arrow2::array::BooleanArray;
Expand All @@ -112,16 +100,15 @@ pub fn is_null(input: &dyn Array) -> BooleanArray {
BooleanArray::new(DataType::Boolean, values, None)
}

/// Returns a non-null [BooleanArray] with whether each value of the array is not null.
/// Returns a non-null [`BooleanArray`] with whether each value of the array is not null.
/// # Example
/// ```rust
/// use arrow2::array::BooleanArray;
/// use arrow2::compute::boolean::is_not_null;
/// # fn main() {
///
/// let a = BooleanArray::from(&vec![Some(false), Some(true), None]);
/// let a_is_not_null = is_not_null(&a);
/// assert_eq!(a_is_not_null, BooleanArray::from_slice(&vec![true, true, false]));
/// # }
/// ```
pub fn is_not_null(input: &dyn Array) -> BooleanArray {
let values = match input.validity() {
Expand All @@ -142,12 +129,12 @@ pub fn is_not_null(input: &dyn Array) -> BooleanArray {
/// 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() {
Expand Down Expand Up @@ -186,7 +173,7 @@ pub fn or_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
}
}

/// Check if any of the values in the array is `true`
/// Returns whether any of the values in the array is `true`
pub fn any(array: &BooleanArray) -> bool {
if array.is_empty() {
false
Expand Down
89 changes: 47 additions & 42 deletions src/compute/boolean_kleene.rs
View file
Original file line number Diff line number Diff line change
@@ -1,35 +1,31 @@
//! Boolean operators of [Kleene logic](https://en.wikipedia.org/wiki/Three-valued_logic#Kleene_and_Priest_logics).
use crate::datatypes::DataType;
use crate::error::{Error, Result};
use crate::scalar::BooleanScalar;
use crate::{
array::BooleanArray,
array::{Array, BooleanArray},
bitmap::{binary, quaternary, ternary, unary, Bitmap, MutableBitmap},
};

/// 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.
/// # Panics
/// This function panics iff the arrays have a different length
/// # Example
///
/// ```rust
/// # use arrow2::error::Result;
/// use arrow2::array::BooleanArray;
/// use arrow2::compute::boolean_kleene::or;
/// # fn main() -> Result<()> {
///
/// let a = BooleanArray::from(&[Some(true), Some(false), None]);
/// let b = BooleanArray::from(&[None, None, None]);
/// let or_ab = or(&a, &b)?;
/// let or_ab = or(&a, &b);
/// assert_eq!(or_ab, BooleanArray::from(&[Some(true), None, None]));
/// # Ok(())
/// # }
/// ```
pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
if lhs.len() != rhs.len() {
return Err(Error::InvalidArgumentError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
assert_eq!(
lhs.len(),
rhs.len(),
"lhs and rhs must have the same length"
);

let lhs_values = lhs.values();
let rhs_values = rhs.values();
Expand Down Expand Up @@ -90,36 +86,29 @@ pub fn or(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
}
(None, None) => None,
};
Ok(BooleanArray::new(
DataType::Boolean,
lhs_values | rhs_values,
validity,
))
BooleanArray::new(DataType::Boolean, lhs_values | rhs_values, validity)
}

/// 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.
/// # Panics
/// This function panics iff the arrays have a different length
/// # Example
///
/// ```rust
/// # use arrow2::error::Result;
/// use arrow2::array::BooleanArray;
/// use arrow2::compute::boolean_kleene::and;
/// # fn main() -> Result<()> {
///
/// let a = BooleanArray::from(&[Some(true), Some(false), None]);
/// let b = BooleanArray::from(&[None, None, None]);
/// let and_ab = and(&a, &b)?;
/// let and_ab = and(&a, &b);
/// assert_eq!(and_ab, BooleanArray::from(&[None, Some(false), None]));
/// # Ok(())
/// # }
/// ```
pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
if lhs.len() != rhs.len() {
return Err(Error::InvalidArgumentError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> BooleanArray {
assert_eq!(
lhs.len(),
rhs.len(),
"lhs and rhs must have the same length"
);

let lhs_values = lhs.values();
let rhs_values = rhs.values();
Expand Down Expand Up @@ -179,11 +168,7 @@ pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
}
(None, None) => None,
};
Ok(BooleanArray::new(
DataType::Boolean,
lhs_values & rhs_values,
validity,
))
BooleanArray::new(DataType::Boolean, lhs_values & rhs_values, 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)
Expand All @@ -193,12 +178,11 @@ pub fn and(lhs: &BooleanArray, rhs: &BooleanArray) -> Result<BooleanArray> {
/// 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() {
Expand Down Expand Up @@ -226,12 +210,11 @@ pub fn or_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray {
/// 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() {
Expand All @@ -250,3 +233,25 @@ pub fn and_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray
}
}
}

/// Returns whether any of the values in the array is `true`
pub fn any(array: &BooleanArray) -> bool {
if array.is_empty() {
false
} else if array.validity().is_some() {
array.into_iter().any(|v| v == Some(true))
} else {
let vals = array.values();
vals.unset_bits() != vals.len()
}
}

/// Returns whether all values in the array are `true`
pub fn all(array: &BooleanArray) -> bool {
if array.is_empty() || array.null_count() > 0 {
false
} else {
let vals = array.values();
vals.unset_bits() == 0
}
}
15 changes: 7 additions & 8 deletions src/compute/comparison/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -515,21 +515,20 @@ fn finish_eq_validities(
(Some(lhs), None) => compute::boolean::and(
&BooleanArray::new(DataType::Boolean, lhs, None),
&output_without_validities,
)
.unwrap(),
),
(None, Some(rhs)) => compute::boolean::and(
&output_without_validities,
&BooleanArray::new(DataType::Boolean, rhs, None),
)
.unwrap(),
),
(Some(lhs), Some(rhs)) => {
let lhs = BooleanArray::new(DataType::Boolean, lhs, None);
let rhs = BooleanArray::new(DataType::Boolean, rhs, None);
let eq_validities = compute::comparison::boolean::eq(&lhs, &rhs);
compute::boolean::and(&output_without_validities, &eq_validities).unwrap()
compute::boolean::and(&output_without_validities, &eq_validities)
}
}
}

fn finish_neq_validities(
output_without_validities: BooleanArray,
validity_lhs: Option<Bitmap>,
Expand All @@ -540,18 +539,18 @@ fn finish_neq_validities(
(Some(lhs), None) => {
let lhs_negated =
compute::boolean::not(&BooleanArray::new(DataType::Boolean, lhs, None));
compute::boolean::or(&lhs_negated, &output_without_validities).unwrap()
compute::boolean::or(&lhs_negated, &output_without_validities)
}
(None, Some(rhs)) => {
let rhs_negated =
compute::boolean::not(&BooleanArray::new(DataType::Boolean, rhs, None));
compute::boolean::or(&output_without_validities, &rhs_negated).unwrap()
compute::boolean::or(&output_without_validities, &rhs_negated)
}
(Some(lhs), Some(rhs)) => {
let lhs = BooleanArray::new(DataType::Boolean, lhs, None);
let rhs = BooleanArray::new(DataType::Boolean, rhs, None);
let neq_validities = compute::comparison::boolean::neq(&lhs, &rhs);
compute::boolean::or(&output_without_validities, &neq_validities).unwrap()
compute::boolean::or(&output_without_validities, &neq_validities)
}
}
}

0 comments on commit a2d9d1a

Please sign in to comment.