diff --git a/arrow/src/compute/kernels/boolean.rs b/arrow/src/compute/kernels/boolean.rs index f9e17839e237..fd3539455b5a 100644 --- a/arrow/src/compute/kernels/boolean.rs +++ b/arrow/src/compute/kernels/boolean.rs @@ -175,7 +175,7 @@ where } /// Helper function to implement binary kernels -fn binary_boolean_kernel( +pub(crate) fn binary_boolean_kernel( left: &BooleanArray, right: &BooleanArray, op: F, diff --git a/arrow/src/compute/kernels/comparison.rs b/arrow/src/compute/kernels/comparison.rs index f246b24372aa..22a6f85da137 100644 --- a/arrow/src/compute/kernels/comparison.rs +++ b/arrow/src/compute/kernels/comparison.rs @@ -26,7 +26,8 @@ use regex::Regex; use std::collections::HashMap; use crate::array::*; -use crate::buffer::{Buffer, MutableBuffer}; +use crate::buffer::{bitwise_bin_op_helper, buffer_unary_not, Buffer, MutableBuffer}; +use crate::compute::binary_boolean_kernel; use crate::compute::util::combine_option_bitmap; use crate::datatypes::{ArrowNumericType, DataType}; use crate::error::{ArrowError, Result}; @@ -623,6 +624,84 @@ pub fn eq_utf8_scalar( compare_op_scalar!(left, right, |a, b| a == b) } +/// Perform `left == right` operation on [`BooleanArray`] +fn eq_bool(left: &BooleanArray, right: &BooleanArray) -> Result { + binary_boolean_kernel( + left, + right, + |left: &Buffer, + left_offset_in_bits: usize, + right: &Buffer, + right_offset_in_bits: usize, + len_in_bits: usize| { + bitwise_bin_op_helper( + left, + left_offset_in_bits, + right, + right_offset_in_bits, + len_in_bits, + |a, b| !(a ^ b), + ) + }, + ) +} + +/// Perform `left != right` operation on [`BooleanArray`] +fn neq_bool(left: &BooleanArray, right: &BooleanArray) -> Result { + binary_boolean_kernel( + left, + right, + |left: &Buffer, + left_offset_in_bits: usize, + right: &Buffer, + right_offset_in_bits: usize, + len_in_bits: usize| { + bitwise_bin_op_helper( + left, + left_offset_in_bits, + right, + right_offset_in_bits, + len_in_bits, + |a, b| (a ^ b), + ) + }, + ) +} + +/// Perform `left == right` operation on [`BooleanArray`] and a scalar +fn eq_bool_scalar(left: &BooleanArray, right: bool) -> Result { + let len = left.len(); + let left_offset = left.offset(); + + let values = if right { + left.values().bit_slice(left_offset, len) + } else { + buffer_unary_not(left.values(), left.offset(), left.len()) + }; + + let data = unsafe { + ArrayData::new_unchecked( + DataType::Boolean, + len, + None, + left.data_ref() + .null_bitmap() + .as_ref() + .map(|b| b.bits.bit_slice(left_offset, len)), + 0, + vec![values], + vec![], + ) + }; + + Ok(BooleanArray::from(data)) +} + +/// Perform `left != right` operation on [`BooleanArray`] and a scalar +fn neq_bool_scalar(left: &BooleanArray, right: bool) -> Result { + eq_bool_scalar(left, !right) +} + /// Perform `left != right` operation on [`StringArray`] / [`LargeStringArray`]. pub fn neq_utf8( left: &GenericStringArray, @@ -1252,6 +1331,67 @@ mod tests { ); } + #[test] + fn test_boolean_array_eq() { + let a: BooleanArray = + vec![Some(true), Some(false), Some(false), Some(true), Some(true), None] + .into(); + let b: BooleanArray = + vec![Some(true), Some(true), Some(false), Some(false), None, Some(false)] + .into(); + + let res: Vec> = eq_bool(&a, &b).unwrap().iter().collect(); + + assert_eq!( + res, + vec![Some(true), Some(false), Some(true), Some(false), None, None] + ) + } + + #[test] + fn test_boolean_array_neq() { + let a: BooleanArray = + vec![Some(true), Some(false), Some(false), Some(true), Some(true), None] + .into(); + let b: BooleanArray = + vec![Some(true), Some(true), Some(false), Some(false), None, Some(false)] + .into(); + + let res: Vec> = neq_bool(&a, &b).unwrap().iter().collect(); + + assert_eq!( + res, + vec![Some(false), Some(true), Some(false), Some(true), None, None] + ) + } + + #[test] + fn test_boolean_array_eq_scalar() { + let a: BooleanArray = vec![Some(true), Some(false), None].into(); + + let res1: Vec> = eq_bool_scalar(&a, false).unwrap().iter().collect(); + + assert_eq!(res1, vec![Some(false), Some(true), None]); + + let res2: Vec> = eq_bool_scalar(&a, true).unwrap().iter().collect(); + + assert_eq!(res2, vec![Some(true), Some(false), None]); + } + + #[test] + fn test_boolean_array_neq_scalar() { + let a: BooleanArray = vec![Some(true), Some(false), None].into(); + + let res1: Vec> = + neq_bool_scalar(&a, false).unwrap().iter().collect(); + + assert_eq!(res1, vec![Some(true), Some(false), None]); + + let res2: Vec> = neq_bool_scalar(&a, true).unwrap().iter().collect(); + + assert_eq!(res2, vec![Some(false), Some(true), None]); + } + #[test] fn test_primitive_array_lt() { cmp_i64!(