From 5720d9a2450823b68f2174807af74a44c058b0b8 Mon Sep 17 00:00:00 2001 From: Stijn de Gooijer Date: Sat, 19 Aug 2023 10:05:20 +0200 Subject: [PATCH] Fix `any`/`all` for Kleene logic (#1545) --- src/compute/boolean.rs | 42 +++++++++++++++++-- src/compute/boolean_kleene.rs | 66 +++++++++++++++++++++++++----- tests/it/compute/boolean.rs | 9 ++-- tests/it/compute/boolean_kleene.rs | 4 +- 4 files changed, 101 insertions(+), 20 deletions(-) diff --git a/src/compute/boolean.rs b/src/compute/boolean.rs index b8200ef5a95..e34b90c6378 100644 --- a/src/compute/boolean.rs +++ b/src/compute/boolean.rs @@ -230,11 +230,28 @@ pub fn or_scalar(array: &BooleanArray, scalar: &BooleanScalar) -> BooleanArray { } } -/// Returns whether any of the values in the array is `true` +/// Returns whether any of the values in the array are `true`. +/// +/// Null values are ignored. +/// +/// # Example +/// +/// ``` +/// use arrow2::array::BooleanArray; +/// use arrow2::compute::boolean::any; +/// +/// let a = BooleanArray::from(&[Some(true), Some(false)]); +/// let b = BooleanArray::from(&[Some(false), Some(false)]); +/// let c = BooleanArray::from(&[None, Some(false)]); +/// +/// assert_eq!(any(&a), true); +/// assert_eq!(any(&b), false); +/// assert_eq!(any(&c), false); +/// ``` pub fn any(array: &BooleanArray) -> bool { if array.is_empty() { false - } else if array.validity().is_some() { + } else if array.null_count() > 0 { array.into_iter().any(|v| v == Some(true)) } else { let vals = array.values(); @@ -242,12 +259,29 @@ pub fn any(array: &BooleanArray) -> bool { } } -/// Check if all of the values in the array are `true` +/// Returns whether all values in the array are `true`. +/// +/// Null values are ignored. +/// +/// # Example +/// +/// ``` +/// use arrow2::array::BooleanArray; +/// use arrow2::compute::boolean::all; +/// +/// let a = BooleanArray::from(&[Some(true), Some(true)]); +/// let b = BooleanArray::from(&[Some(false), Some(true)]); +/// let c = BooleanArray::from(&[None, Some(true)]); +/// +/// assert_eq!(all(&a), true); +/// assert_eq!(all(&b), false); +/// assert_eq!(all(&c), true); +/// ``` pub fn all(array: &BooleanArray) -> bool { if array.is_empty() { true } else if array.null_count() > 0 { - false + !array.into_iter().any(|v| v == Some(false)) } else { let vals = array.values(); vals.unset_bits() == 0 diff --git a/src/compute/boolean_kleene.rs b/src/compute/boolean_kleene.rs index fc6b717543f..b19efeaa78d 100644 --- a/src/compute/boolean_kleene.rs +++ b/src/compute/boolean_kleene.rs @@ -234,26 +234,70 @@ 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 { +/// Returns whether any of the values in the array are `true`. +/// +/// The output is unknown (`None`) if the array contains any null values and +/// no `true` values. +/// +/// # Example +/// +/// ``` +/// use arrow2::array::BooleanArray; +/// use arrow2::compute::boolean_kleene::any; +/// +/// let a = BooleanArray::from(&[Some(true), Some(false)]); +/// let b = BooleanArray::from(&[Some(false), Some(false)]); +/// let c = BooleanArray::from(&[None, Some(false)]); +/// +/// assert_eq!(any(&a), Some(true)); +/// assert_eq!(any(&b), Some(false)); +/// assert_eq!(any(&c), None); +/// ``` +pub fn any(array: &BooleanArray) -> Option { if array.is_empty() { - false - } else if array.validity().is_some() { - array.into_iter().any(|v| v == Some(true)) + Some(false) + } else if array.null_count() > 0 { + if array.into_iter().any(|v| v == Some(true)) { + Some(true) + } else { + None + } } else { let vals = array.values(); - vals.unset_bits() != vals.len() + Some(vals.unset_bits() != vals.len()) } } -/// Returns whether all values in the array are `true` -pub fn all(array: &BooleanArray) -> bool { +/// Returns whether all values in the array are `true`. +/// +/// The output is unknown (`None`) if the array contains any null values and +/// no `false` values. +/// +/// # Example +/// +/// ``` +/// use arrow2::array::BooleanArray; +/// use arrow2::compute::boolean_kleene::all; +/// +/// let a = BooleanArray::from(&[Some(true), Some(true)]); +/// let b = BooleanArray::from(&[Some(false), Some(true)]); +/// let c = BooleanArray::from(&[None, Some(true)]); +/// +/// assert_eq!(all(&a), Some(true)); +/// assert_eq!(all(&b), Some(false)); +/// assert_eq!(all(&c), None); +/// ``` +pub fn all(array: &BooleanArray) -> Option { if array.is_empty() { - true + Some(true) } else if array.null_count() > 0 { - false + if array.into_iter().any(|v| v == Some(false)) { + Some(false) + } else { + None + } } else { let vals = array.values(); - vals.unset_bits() == 0 + Some(vals.unset_bits() == 0) } } diff --git a/tests/it/compute/boolean.rs b/tests/it/compute/boolean.rs index 8c505a164fc..ae4c0fde85b 100644 --- a/tests/it/compute/boolean.rs +++ b/tests/it/compute/boolean.rs @@ -429,21 +429,24 @@ fn test_any_all() { assert!(!any(&array)); assert!(!all(&array)); let array = BooleanArray::from(&[None, Some(true), Some(true)]); - assert!(!all(&array)); assert!(any(&array)); + assert!(all(&array)); let array = BooleanArray::from_iter(std::iter::repeat(false).take(10).map(Some)); assert!(!any(&array)); assert!(!all(&array)); let array = BooleanArray::from_iter(std::iter::repeat(true).take(10).map(Some)); - assert!(all(&array)); assert!(any(&array)); + assert!(all(&array)); let array = BooleanArray::from_iter([true, false, true, true].map(Some)); - assert!(!all(&array)); assert!(any(&array)); + assert!(!all(&array)); let array = BooleanArray::from(&[Some(true)]); assert!(any(&array)); assert!(all(&array)); let array = BooleanArray::from(&[Some(false)]); assert!(!any(&array)); assert!(!all(&array)); + let array = BooleanArray::from(&[]); + assert!(!any(&array)); + assert!(all(&array)); } diff --git a/tests/it/compute/boolean_kleene.rs b/tests/it/compute/boolean_kleene.rs index 8dac6e63c46..902e5b425ac 100644 --- a/tests/it/compute/boolean_kleene.rs +++ b/tests/it/compute/boolean_kleene.rs @@ -218,6 +218,6 @@ fn array_or_none() { #[test] fn array_empty() { let array = BooleanArray::from(&[]); - assert!(!any(&array)); - assert!(all(&array)); + assert_eq!(any(&array), Some(false)); + assert_eq!(all(&array), Some(true)); }