This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 3ea0e18 Implement boolean equality kernels (#844)
3ea0e18 is described below
commit 3ea0e18acd81c0a959b3a7cbb8863e6627ee69de
Author: Daniël Heres <[email protected]>
AuthorDate: Sat Oct 23 13:50:40 2021 +0200
Implement boolean equality kernels (#844)
* Implement boolean equality kernels
* Respect offset
* Simplify
---
arrow/src/compute/kernels/boolean.rs | 2 +-
arrow/src/compute/kernels/comparison.rs | 142 +++++++++++++++++++++++++++++++-
2 files changed, 142 insertions(+), 2 deletions(-)
diff --git a/arrow/src/compute/kernels/boolean.rs
b/arrow/src/compute/kernels/boolean.rs
index f9e1783..fd35394 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<F>(
+pub(crate) fn binary_boolean_kernel<F>(
left: &BooleanArray,
right: &BooleanArray,
op: F,
diff --git a/arrow/src/compute/kernels/comparison.rs
b/arrow/src/compute/kernels/comparison.rs
index f246b24..22a6f85 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<OffsetSize: StringOffsetSizeTrait>(
compare_op_scalar!(left, right, |a, b| a == b)
}
+/// Perform `left == right` operation on [`BooleanArray`]
+fn eq_bool(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray> {
+ 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<BooleanArray>
{
+ 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<BooleanArray> {
+ 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<BooleanArray> {
+ eq_bool_scalar(left, !right)
+}
+
/// Perform `left != right` operation on [`StringArray`] /
[`LargeStringArray`].
pub fn neq_utf8<OffsetSize: StringOffsetSizeTrait>(
left: &GenericStringArray<OffsetSize>,
@@ -1253,6 +1332,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<Option<bool>> = 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<Option<bool>> = 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<Option<bool>> = eq_bool_scalar(&a,
false).unwrap().iter().collect();
+
+ assert_eq!(res1, vec![Some(false), Some(true), None]);
+
+ let res2: Vec<Option<bool>> = 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<Option<bool>> =
+ neq_bool_scalar(&a, false).unwrap().iter().collect();
+
+ assert_eq!(res1, vec![Some(true), Some(false), None]);
+
+ let res2: Vec<Option<bool>> = 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!(
lt,
