support bitwise op: or,xor,not (#2716)

Author: liukun4515

arrow/src/compute/kernels/bitwise.rs

@@ -19,7 +19,7 @@ use crate::array::PrimitiveArray;
 use crate::compute::{binary, unary};
 use crate::datatypes::ArrowNumericType;
 use crate::error::{ArrowError, Result};
-use std::ops::BitAnd;
+use std::ops::{BitAnd, BitOr, BitXor, Not};
 
 // The helper function for bitwise operation with two array
 fn bitwise_op<T, F>(
@@ -52,7 +52,43 @@ where
     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
+/// Perform `left | right` operation on two arrays. If either left or right value is null
+/// then the result is also null.
+pub fn bitwise_or<T>(
+    left: &PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>>
+where
+    T: ArrowNumericType,
+    T::Native: BitOr<Output = T::Native>,
+{
+    bitwise_op(left, right, |a, b| a | b)
+}
+
+/// Perform `left ^ right` operation on two arrays. If either left or right value is null
+/// then the result is also null.
+pub fn bitwise_xor<T>(
+    left: &PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>>
+where
+    T: ArrowNumericType,
+    T::Native: BitXor<Output = T::Native>,
+{
+    bitwise_op(left, right, |a, b| a ^ b)
+}
+
+/// Perform `!array` operation on array. If array value is null
+/// then the result is also null.
+pub fn bitwise_not<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
+where
+    T: ArrowNumericType,
+    T::Native: Not<Output = T::Native>,
+{
+    Ok(unary(array, |value| !value))
+}
+
+/// 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>(
     array: &PrimitiveArray<T>,
@@ -65,10 +101,39 @@ where
     Ok(unary(array, |value| value & scalar))
 }
 
+/// Perform bitwise `or` 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_or_scalar<T>(
+    array: &PrimitiveArray<T>,
+    scalar: T::Native,
+) -> Result<PrimitiveArray<T>>
+where
+    T: ArrowNumericType,
+    T::Native: BitOr<Output = T::Native>,
+{
+    Ok(unary(array, |value| value | scalar))
+}
+
+/// Perform bitwise `xor` 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_xor_scalar<T>(
+    array: &PrimitiveArray<T>,
+    scalar: T::Native,
+) -> Result<PrimitiveArray<T>>
+where
+    T: ArrowNumericType,
+    T::Native: BitXor<Output = T::Native>,
+{
+    Ok(unary(array, |value| value ^ scalar))
+}
+
 #[cfg(test)]
 mod tests {
     use crate::array::{Int32Array, UInt64Array};
-    use crate::compute::kernels::bitwise::{bitwise_and, bitwise_and_scalar};
+    use crate::compute::kernels::bitwise::{
+        bitwise_and, bitwise_and_scalar, bitwise_not, bitwise_or, bitwise_or_scalar,
+        bitwise_xor, bitwise_xor_scalar,
+    };
     use crate::error::Result;
 
     #[test]
@@ -82,7 +147,7 @@ mod tests {
 
         // signed value
         let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
-        let right = Int32Array::from(vec![Some(5), Some(10), Some(8), Some(12)]);
+        let right = Int32Array::from(vec![Some(5), Some(-10), Some(8), Some(12)]);
         let expected = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
         let result = bitwise_and(&left, &right)?;
         assert_eq!(expected, result);
@@ -100,10 +165,102 @@ mod tests {
 
         // signed value
         let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
-        let scalar = 20;
+        let scalar = -20;
         let expected = Int32Array::from(vec![Some(0), Some(0), None, Some(4)]);
         let result = bitwise_and_scalar(&left, scalar)?;
         assert_eq!(expected, result);
         Ok(())
     }
+
+    #[test]
+    fn test_bitwise_or_array() -> Result<()> {
+        // unsigned value
+        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
+        let expected = UInt64Array::from(vec![Some(7), Some(7), None, Some(13)]);
+        let result = bitwise_or(&left, &right)?;
+        assert_eq!(expected, result);
+
+        // signed value
+        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
+        let expected = Int32Array::from(vec![Some(-7), Some(-5), None, Some(13)]);
+        let result = bitwise_or(&left, &right)?;
+        assert_eq!(expected, result);
+        Ok(())
+    }
+
+    #[test]
+    fn test_bitwise_not_array() -> Result<()> {
+        // unsigned value
+        let array = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let expected = UInt64Array::from(vec![
+            Some(18446744073709551614),
+            Some(18446744073709551613),
+            None,
+            Some(18446744073709551611),
+        ]);
+        let result = bitwise_not(&array)?;
+        assert_eq!(expected, result);
+        // signed value
+        let array = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let expected = Int32Array::from(vec![Some(-2), Some(-3), None, Some(-5)]);
+        let result = bitwise_not(&array)?;
+        assert_eq!(expected, result);
+        Ok(())
+    }
+
+    #[test]
+    fn test_bitwise_or_array_scalar() -> Result<()> {
+        // unsigned value
+        let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
+        let scalar = 7;
+        let expected = UInt64Array::from(vec![Some(15), Some(7), None, Some(7)]);
+        let result = bitwise_or_scalar(&left, scalar)?;
+        assert_eq!(expected, result);
+
+        // signed value
+        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let scalar = 20;
+        let expected = Int32Array::from(vec![Some(21), Some(22), None, Some(20)]);
+        let result = bitwise_or_scalar(&left, scalar)?;
+        assert_eq!(expected, result);
+        Ok(())
+    }
+
+    #[test]
+    fn test_bitwise_xor_array() -> Result<()> {
+        // unsigned value
+        let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
+        let expected = UInt64Array::from(vec![Some(6), Some(7), None, Some(9)]);
+        let result = bitwise_xor(&left, &right)?;
+        assert_eq!(expected, result);
+
+        // signed value
+        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let right = Int32Array::from(vec![Some(-7), Some(5), Some(8), Some(-13)]);
+        let expected = Int32Array::from(vec![Some(-8), Some(7), None, Some(-9)]);
+        let result = bitwise_xor(&left, &right)?;
+        assert_eq!(expected, result);
+        Ok(())
+    }
+
+    #[test]
+    fn test_bitwise_xor_array_scalar() -> Result<()> {
+        // unsigned value
+        let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
+        let scalar = 7;
+        let expected = UInt64Array::from(vec![Some(8), Some(5), None, Some(3)]);
+        let result = bitwise_xor_scalar(&left, scalar)?;
+        assert_eq!(expected, result);
+
+        // signed value
+        let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
+        let scalar = -20;
+        let expected = Int32Array::from(vec![Some(-19), Some(-18), None, Some(-24)]);
+        let result = bitwise_xor_scalar(&left, scalar)?;
+        assert_eq!(expected, result);
+        Ok(())
+    }
 }