This is an automated email from the ASF dual-hosted git repository.
tustvold 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 ba6dbb8e27 Add support for more fused boolean operations (#5298)
ba6dbb8e27 is described below
commit ba6dbb8e27bbdc0453f4e8694bd24f09b8d7e7b1
Author: RT_Enzyme <[email protected]>
AuthorDate: Mon Jan 15 22:27:43 2024 +0800
Add support for more fused boolean operations (#5298)
* 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.
---
arrow-arith/src/bitwise.rs | 39 +++++++++++++++++++++++++++++++++++++++
arrow-arith/src/boolean.rs | 35 ++++++++++++++++++++++++++++++++++-
arrow-buffer/src/buffer/ops.rs | 19 +++++++++++++++++++
3 files changed, 92 insertions(+), 1 deletion(-)
diff --git a/arrow-arith/src/bitwise.rs b/arrow-arith/src/bitwise.rs
index c7885952f8..c829a3c29f 100644
--- a/arrow-arith/src/bitwise.rs
+++ b/arrow-arith/src/bitwise.rs
@@ -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>(
@@ -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
diff --git a/arrow-arith/src/boolean.rs b/arrow-arith/src/boolean.rs
index 269a36d66c..ea8e12abbe 100644
--- a/arrow-arith/src/boolean.rs
+++ b/arrow-arith/src/boolean.rs
@@ -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
@@ -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, ¬(&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
@@ -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, ¬(&b).unwrap()).unwrap());
+ }
+
#[test]
fn test_bool_array_or_nulls() {
let a = BooleanArray::from(vec