-
Notifications
You must be signed in to change notification settings - Fork 784
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Implement boolean equality kernels #844
Changes from 2 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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,107 @@ 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( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. How about using There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. That only should be used together with a |
||
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( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not sure if there is a more elegant way of doing this. |
||
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> { | ||
let len = left.len(); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I wonder if this could simply be a call to There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good idea! |
||
let left_offset = left.offset(); | ||
|
||
let values = if right { | ||
buffer_unary_not(left.values(), left.offset(), left.len()) | ||
} else { | ||
left.values().bit_slice(left_offset, 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 [`StringArray`] / [`LargeStringArray`]. | ||
pub fn neq_utf8<OffsetSize: StringOffsetSizeTrait>( | ||
left: &GenericStringArray<OffsetSize>, | ||
|
@@ -1252,6 +1354,67 @@ mod tests { | |
); | ||
} | ||
|
||
#[test] | ||
fn test_boolean_array_eq() { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Normally since there is special case handling for processing these at 64 bits at a time, I would suggest tests that have at least 65 entries to test the boundary conditions. However, in this case as it is using a pre-existing helper method |
||
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!( | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This works per pair of
u64
instead of per bool.