Skip to content

Commit

Permalink
Add support for more fused boolean operations (#5298)
Browse files Browse the repository at this point in the history
* feat: don't check schema and batch

* Add support for more fused boolean operations

* Add support for more fused boolean operations

* remove git error message in code.

* format code

* fix the Clippy error.
  • Loading branch information
RTEnzyme authored Jan 15, 2024
1 parent 8345991 commit ba6dbb8
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 1 deletion.
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

0 comments on commit ba6dbb8

Please sign in to comment.