Use arrow kernels for bitwise operations (#6098)

* Use arrow kernels for bitwise operations

* Apply suggestions from code review

Co-authored-by: Liang-Chi Hsieh <[email protected]>

* Add link to upstream arrow ticket

---------

Co-authored-by: Liang-Chi Hsieh <[email protected]>

Author: alamb

datafusion/physical-expr/src/expressions/binary.rs

@@ -63,9 +63,9 @@ use datafusion_common::scalar::{
 };
 use datafusion_expr::type_coercion::{is_decimal, is_timestamp, is_utf8_or_large_utf8};
 use kernels::{
-    bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_shift_left,
-    bitwise_shift_left_scalar, bitwise_shift_right, bitwise_shift_right_scalar,
-    bitwise_xor, bitwise_xor_scalar,
+    bitwise_and_dyn, bitwise_and_dyn_scalar, bitwise_or_dyn, bitwise_or_dyn_scalar,
+    bitwise_shift_left_dyn, bitwise_shift_left_dyn_scalar, bitwise_shift_right_dyn,
+    bitwise_shift_right_dyn_scalar, bitwise_xor_dyn, bitwise_xor_dyn_scalar,
 };
 use kernels_arrow::{
     add_decimal_dyn_scalar, add_dyn_decimal, add_dyn_temporal, add_dyn_temporal_scalar,
@@ -1135,11 +1135,11 @@ impl BinaryExpr {
                 true,
                 true
             ),
-            Operator::BitwiseAnd => bitwise_and_scalar(array, scalar),
-            Operator::BitwiseOr => bitwise_or_scalar(array, scalar),
-            Operator::BitwiseXor => bitwise_xor_scalar(array, scalar),
-            Operator::BitwiseShiftRight => bitwise_shift_right_scalar(array, scalar),
-            Operator::BitwiseShiftLeft => bitwise_shift_left_scalar(array, scalar),
+            Operator::BitwiseAnd => bitwise_and_dyn_scalar(array, scalar),
+            Operator::BitwiseOr => bitwise_or_dyn_scalar(array, scalar),
+            Operator::BitwiseXor => bitwise_xor_dyn_scalar(array, scalar),
+            Operator::BitwiseShiftRight => bitwise_shift_right_dyn_scalar(array, scalar),
+            Operator::BitwiseShiftLeft => bitwise_shift_left_dyn_scalar(array, scalar),
             // if scalar operation is not supported - fallback to array implementation
             _ => None,
         };
@@ -1257,11 +1257,11 @@ impl BinaryExpr {
             Operator::RegexNotIMatch => {
                 binary_string_array_flag_op!(left, right, regexp_is_match, true, true)
             }
-            Operator::BitwiseAnd => bitwise_and(left, right),
-            Operator::BitwiseOr => bitwise_or(left, right),
-            Operator::BitwiseXor => bitwise_xor(left, right),
-            Operator::BitwiseShiftRight => bitwise_shift_right(left, right),
-            Operator::BitwiseShiftLeft => bitwise_shift_left(left, right),
+            Operator::BitwiseAnd => bitwise_and_dyn(left, right),
+            Operator::BitwiseOr => bitwise_or_dyn(left, right),
+            Operator::BitwiseXor => bitwise_xor_dyn(left, right),
+            Operator::BitwiseShiftRight => bitwise_shift_right_dyn(left, right),
+            Operator::BitwiseShiftLeft => bitwise_shift_left_dyn(left, right),
             Operator::StringConcat => {
                 binary_string_array_op!(left, right, concat_elements)
             }
@@ -5147,31 +5147,31 @@ mod tests {
         let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
         let right =
             Arc::new(Int32Array::from(vec![Some(1), Some(3), Some(7)])) as ArrayRef;
-        let mut result = bitwise_and(left.clone(), right.clone())?;
+        let mut result = bitwise_and_dyn(left.clone(), right.clone())?;
         let expected = Int32Array::from(vec![Some(0), None, Some(3)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_or(left.clone(), right.clone())?;
+        result = bitwise_or_dyn(left.clone(), right.clone())?;
         let expected = Int32Array::from(vec![Some(13), None, Some(15)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_xor(left.clone(), right.clone())?;
+        result = bitwise_xor_dyn(left.clone(), right.clone())?;
         let expected = Int32Array::from(vec![Some(13), None, Some(12)]);
         assert_eq!(result.as_ref(), &expected);
 
         let left =
             Arc::new(UInt32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
         let right =
             Arc::new(UInt32Array::from(vec![Some(1), Some(3), Some(7)])) as ArrayRef;
-        let mut result = bitwise_and(left.clone(), right.clone())?;
+        let mut result = bitwise_and_dyn(left.clone(), right.clone())?;
         let expected = UInt32Array::from(vec![Some(0), None, Some(3)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_or(left.clone(), right.clone())?;
+        result = bitwise_or_dyn(left.clone(), right.clone())?;
         let expected = UInt32Array::from(vec![Some(13), None, Some(15)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_xor(left.clone(), right.clone())?;
+        result = bitwise_xor_dyn(left.clone(), right.clone())?;
         let expected = UInt32Array::from(vec![Some(13), None, Some(12)]);
         assert_eq!(result.as_ref(), &expected);
 
@@ -5183,24 +5183,24 @@ mod tests {
         let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(10)])) as ArrayRef;
         let modules =
             Arc::new(Int32Array::from(vec![Some(2), Some(4), Some(8)])) as ArrayRef;
-        let mut result = bitwise_shift_left(input.clone(), modules.clone())?;
+        let mut result = bitwise_shift_left_dyn(input.clone(), modules.clone())?;
 
         let expected = Int32Array::from(vec![Some(8), None, Some(2560)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_shift_right(result.clone(), modules.clone())?;
+        result = bitwise_shift_right_dyn(result.clone(), modules.clone())?;
         assert_eq!(result.as_ref(), &input);
 
         let input =
             Arc::new(UInt32Array::from(vec![Some(2), None, Some(10)])) as ArrayRef;
         let modules =
             Arc::new(UInt32Array::from(vec![Some(2), Some(4), Some(8)])) as ArrayRef;
-        let mut result = bitwise_shift_left(input.clone(), modules.clone())?;
+        let mut result = bitwise_shift_left_dyn(input.clone(), modules.clone())?;
 
         let expected = UInt32Array::from(vec![Some(8), None, Some(2560)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_shift_right(result.clone(), modules.clone())?;
+        result = bitwise_shift_right_dyn(result.clone(), modules.clone())?;
         assert_eq!(result.as_ref(), &input);
         Ok(())
     }
@@ -5209,14 +5209,14 @@ mod tests {
     fn bitwise_shift_array_overflow_test() -> Result<()> {
         let input = Arc::new(Int32Array::from(vec![Some(2)])) as ArrayRef;
         let modules = Arc::new(Int32Array::from(vec![Some(100)])) as ArrayRef;
-        let result = bitwise_shift_left(input.clone(), modules.clone())?;
+        let result = bitwise_shift_left_dyn(input.clone(), modules.clone())?;
 
         let expected = Int32Array::from(vec![Some(32)]);
         assert_eq!(result.as_ref(), &expected);
 
         let input = Arc::new(UInt32Array::from(vec![Some(2)])) as ArrayRef;
         let modules = Arc::new(UInt32Array::from(vec![Some(100)])) as ArrayRef;
-        let result = bitwise_shift_left(input.clone(), modules.clone())?;
+        let result = bitwise_shift_left_dyn(input.clone(), modules.clone())?;
 
         let expected = UInt32Array::from(vec![Some(32)]);
         assert_eq!(result.as_ref(), &expected);
@@ -5227,30 +5227,30 @@ mod tests {
     fn bitwise_scalar_test() -> Result<()> {
         let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
         let right = ScalarValue::from(3i32);
-        let mut result = bitwise_and_scalar(&left, right.clone()).unwrap()?;
+        let mut result = bitwise_and_dyn_scalar(&left, right.clone()).unwrap()?;
         let expected = Int32Array::from(vec![Some(0), None, Some(3)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_or_scalar(&left, right.clone()).unwrap()?;
+        result = bitwise_or_dyn_scalar(&left, right.clone()).unwrap()?;
         let expected = Int32Array::from(vec![Some(15), None, Some(11)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_xor_scalar(&left, right).unwrap()?;
+        result = bitwise_xor_dyn_scalar(&left, right).unwrap()?;
         let expected = Int32Array::from(vec![Some(15), None, Some(8)]);
         assert_eq!(result.as_ref(), &expected);
 
         let left =
             Arc::new(UInt32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
         let right = ScalarValue::from(3u32);
-        let mut result = bitwise_and_scalar(&left, right.clone()).unwrap()?;
+        let mut result = bitwise_and_dyn_scalar(&left, right.clone()).unwrap()?;
         let expected = UInt32Array::from(vec![Some(0), None, Some(3)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_or_scalar(&left, right.clone()).unwrap()?;
+        result = bitwise_or_dyn_scalar(&left, right.clone()).unwrap()?;
         let expected = UInt32Array::from(vec![Some(15), None, Some(11)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_xor_scalar(&left, right).unwrap()?;
+        result = bitwise_xor_dyn_scalar(&left, right).unwrap()?;
         let expected = UInt32Array::from(vec![Some(15), None, Some(8)]);
         assert_eq!(result.as_ref(), &expected);
         Ok(())
@@ -5260,22 +5260,24 @@ mod tests {
     fn bitwise_shift_scalar_test() -> Result<()> {
         let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef;
         let module = ScalarValue::from(10i32);
-        let mut result = bitwise_shift_left_scalar(&input, module.clone()).unwrap()?;
+        let mut result =
+            bitwise_shift_left_dyn_scalar(&input, module.clone()).unwrap()?;
 
         let expected = Int32Array::from(vec![Some(2048), None, Some(4096)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_shift_right_scalar(&result, module).unwrap()?;
+        result = bitwise_shift_right_dyn_scalar(&result, module).unwrap()?;
         assert_eq!(result.as_ref(), &input);
 
         let input = Arc::new(UInt32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef;
         let module = ScalarValue::from(10u32);
-        let mut result = bitwise_shift_left_scalar(&input, module.clone()).unwrap()?;
+        let mut result =
+            bitwise_shift_left_dyn_scalar(&input, module.clone()).unwrap()?;
 
         let expected = UInt32Array::from(vec![Some(2048), None, Some(4096)]);
         assert_eq!(result.as_ref(), &expected);
 
-        result = bitwise_shift_right_scalar(&result, module).unwrap()?;
+        result = bitwise_shift_right_dyn_scalar(&result, module).unwrap()?;
         assert_eq!(result.as_ref(), &input);
         Ok(())
     }