From ba6dbb8e27bbdc0453f4e8694bd24f09b8d7e7b1 Mon Sep 17 00:00:00 2001
From: RT_Enzyme <58059931+RTEnzyme@users.noreply.github.com>
Date: Mon, 15 Jan 2024 22:27:43 +0800
Subject: [PATCH] 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 c7885952f8ba..c829a3c29fff 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 269a36d66c2b..ea8e12abbe2c 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, &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
@@ -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![
diff --git a/arrow-buffer/src/buffer/ops.rs b/arrow-buffer/src/buffer/ops.rs
index ca00e41bea21..c69e5c6deb10 100644
--- a/arrow-buffer/src/buffer/ops.rs
+++ b/arrow-buffer/src/buffer/ops.rs
@@ -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 {