Skip to content
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

Add support for more fused boolean operations #5298

Merged
merged 8 commits into from
Jan 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions arrow-arith/src/bitwise.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,20 @@ where
Ok(unary(array, |value| !value))
}

/// Perform `left & !right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_and_not<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitAnd<Output = T::Native>,
T::Native: Not<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a & !b)
}

/// Perform bitwise `and` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_and_scalar<T>(
Expand Down Expand Up @@ -298,6 +312,31 @@ mod tests {
assert_eq!(expected, result);
}

#[test]
fn test_bitwise_and_not_array() {
// unsigned value
let left = UInt64Array::from(vec![Some(8), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
let expected = UInt64Array::from(vec![Some(8), Some(2), None, Some(0)]);
let result = bitwise_and_not(&left, &right).unwrap();
assert_eq!(expected, result);
assert_eq!(
bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
result
);

// signed value
let left = Int32Array::from(vec![Some(2), Some(1), None, Some(3)]);
let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
let expected = Int32Array::from(vec![Some(2), Some(0), None, Some(2)]);
let result = bitwise_and_not(&left, &right).unwrap();
assert_eq!(expected, result);
assert_eq!(
bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
result
);
}

#[test]
fn test_bitwise_or_array_scalar() {
// unsigned value
Expand Down
35 changes: 34 additions & 1 deletion arrow-arith/src/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@

use arrow_array::*;
use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper};
use arrow_buffer::{BooleanBuffer, NullBuffer};
use arrow_buffer::{buffer_bin_and_not, BooleanBuffer, NullBuffer};
use arrow_schema::ArrowError;

/// Logical 'and' boolean values with Kleene logic
Expand Down Expand Up @@ -272,6 +272,27 @@ pub fn or(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, Arr
binary_boolean_kernel(left, right, |a, b| a | b)
}

/// Performs `AND_NOT` 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
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::{and, not, and_not};
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
/// let andn_ab = and_not(&a, &b).unwrap();
/// assert_eq!(andn_ab, BooleanArray::from(vec![Some(false), Some(false), None]));
/// // It's equal to and(left, not(right))
/// assert_eq!(andn_ab, and(&a, &not(&b).unwrap()).unwrap());
pub fn and_not(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| {
let buffer = buffer_bin_and_not(a.inner(), b.offset(), b.inner(), a.offset(), a.len());
BooleanBuffer::new(buffer, left.offset(), left.len())
})
}

/// Performs unary `NOT` operation on an arrays. If value is null then the result is also
/// null.
/// # Error
Expand Down Expand Up @@ -356,6 +377,18 @@ mod tests {
assert_eq!(c, expected);
}

#[test]
fn test_bool_array_and_not() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = and_not(&a, &b).unwrap();

let expected = BooleanArray::from(vec![false, false, true, false]);

assert_eq!(c, expected);
assert_eq!(c, and(&a, &not(&b).unwrap()).unwrap());
}

#[test]
fn test_bool_array_or_nulls() {
let a = BooleanArray::from(vec![
Expand Down
19 changes: 19 additions & 0 deletions arrow-buffer/src/buffer/ops.rs
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,25 @@ pub fn buffer_bin_xor(
)
}

/// Apply a bitwise and_not to two inputs and return the result as a Buffer.
/// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits.
pub fn buffer_bin_and_not(
left: &Buffer,
left_offset_in_bits: usize,
right: &Buffer,
right_offset_in_bits: usize,
len_in_bits: usize,
) -> Buffer {
bitwise_bin_op_helper(
left,
left_offset_in_bits,
right,
right_offset_in_bits,
len_in_bits,
|a, b| a & !b,
)
}

/// Apply a bitwise not to one input and return the result as a Buffer.
/// The input is treated as a bitmap, meaning that offset and length are specified in number of bits.
pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer {
Expand Down
Loading