feat: Support Binary bitwise shift operators (<< and >>)

datafusion/core/tests/sql/expr.rs

@@ -644,6 +644,16 @@ async fn test_struct_literals() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn binary_bitwise_shift() -> Result<()> {
+    test_expression!("2 << 10", "2048");
+    test_expression!("2048 >> 10", "2");
+    test_expression!("2048 << NULL", "NULL");
+    test_expression!("2048 >> NULL", "NULL");
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_interval_expressions() -> Result<()> {
     // day nano intervals

datafusion/expr/src/binary_rule.rs

@@ -55,7 +55,10 @@ pub fn binary_operator_data_type(
         | Operator::IsDistinctFrom
         | Operator::IsNotDistinctFrom => Ok(DataType::Boolean),
         // bitwise operations return the common coerced type
-        Operator::BitwiseAnd | Operator::BitwiseOr => Ok(result_type),
+        Operator::BitwiseAnd
+        | Operator::BitwiseOr
+        | Operator::BitwiseShiftLeft
+        | Operator::BitwiseShiftRight => Ok(result_type),
         // math operations return the same value as the common coerced type
         Operator::Plus
         | Operator::Minus
@@ -76,9 +79,10 @@ pub fn coerce_types(
 ) -> Result<DataType> {
     // This result MUST be compatible with `binary_coerce`
     let result = match op {
-        Operator::BitwiseAnd | Operator::BitwiseOr => {
-            bitwise_coercion(lhs_type, rhs_type)
-        }
+        Operator::BitwiseAnd
+        | Operator::BitwiseOr
+        | Operator::BitwiseShiftRight
+        | Operator::BitwiseShiftLeft => bitwise_coercion(lhs_type, rhs_type),
         Operator::And | Operator::Or => match (lhs_type, rhs_type) {
             // logical binary boolean operators can only be evaluated in bools
             (DataType::Boolean, DataType::Boolean) => Some(DataType::Boolean),
@@ -135,12 +139,14 @@ pub fn coerce_types(
 fn bitwise_coercion(left_type: &DataType, right_type: &DataType) -> Option<DataType> {
     use arrow::datatypes::DataType::*;
 
-    if !is_numeric(left_type) || !is_numeric(right_type) {
+    if !both_numeric_or_null_and_numeric(left_type, right_type) {
         return None;
     }
+
     if left_type == right_type && !is_dictionary(left_type) {
         return Some(left_type.clone());
     }
+
     // TODO support other data type
     match (left_type, right_type) {
         (Int64, _) | (_, Int64) => Some(Int64),

datafusion/expr/src/operator.rs

@@ -71,6 +71,10 @@ pub enum Operator {
     BitwiseAnd,
     /// Bitwise or, like `|`
     BitwiseOr,
+    /// Bitwise right, like `>>`
+    BitwiseShiftRight,
+    /// Bitwise right, like `<<`
+    BitwiseShiftLeft,
     /// String concat
     StringConcat,
 }
@@ -101,6 +105,8 @@ impl fmt::Display for Operator {
             Operator::IsNotDistinctFrom => "IS NOT DISTINCT FROM",
             Operator::BitwiseAnd => "&",
             Operator::BitwiseOr => "|",
+            Operator::BitwiseShiftRight => ">>",
+            Operator::BitwiseShiftLeft => "<<",
             Operator::StringConcat => "||",
         };
         write!(f, "{}", display)

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

@@ -50,7 +50,10 @@ use arrow::compute::kernels::comparison::{
 
 use adapter::{eq_dyn, gt_dyn, gt_eq_dyn, lt_dyn, lt_eq_dyn, neq_dyn};
 use arrow::compute::kernels::concat_elements::concat_elements_utf8;
-use kernels::{bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar};
+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,
+};
 use kernels_arrow::{
     add_decimal, add_decimal_scalar, divide_decimal, divide_decimal_scalar,
     eq_decimal_scalar, gt_decimal_scalar, gt_eq_decimal_scalar, is_distinct_from,
@@ -740,6 +743,12 @@ impl BinaryExpr {
             ),
             Operator::BitwiseAnd => bitwise_and_scalar(array, scalar.clone()),
             Operator::BitwiseOr => bitwise_or_scalar(array, scalar.clone()),
+            Operator::BitwiseShiftRight => {
+                bitwise_shift_right_scalar(array, scalar.clone())
+            }
+            Operator::BitwiseShiftLeft => {
+                bitwise_shift_left_scalar(array, scalar.clone())
+            }
             // if scalar operation is not supported - fallback to array implementation
             _ => None,
         };
@@ -850,6 +859,8 @@ impl BinaryExpr {
             }
             Operator::BitwiseAnd => bitwise_and(left, right),
             Operator::BitwiseOr => bitwise_or(left, right),
+            Operator::BitwiseShiftRight => bitwise_shift_right(left, right),
+            Operator::BitwiseShiftLeft => bitwise_shift_left(left, right),
             Operator::StringConcat => {
                 binary_string_array_op!(left, right, concat_elements)
             }
@@ -2481,6 +2492,34 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn bitwise_shift_array_test() -> Result<()> {
+        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 expected = Int32Array::from(vec![Some(8), None, Some(2560)]);
+        assert_eq!(result.as_ref(), &expected);
+
+        result = bitwise_shift_right(result.clone(), modules.clone())?;
+        assert_eq!(result.as_ref(), &input);
+
+        Ok(())
+    }
+
+    #[test]
+    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 expected = Int32Array::from(vec![Some(32)]);
+        assert_eq!(result.as_ref(), &expected);
+
+        Ok(())
+    }
+
     #[test]
     fn bitwise_scalar_test() -> Result<()> {
         let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
@@ -2494,4 +2533,19 @@ mod tests {
         assert_eq!(result.as_ref(), &expected);
         Ok(())
     }
+
+    #[test]
+    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 expected = Int32Array::from(vec![Some(2048), None, Some(4096)]);
+        assert_eq!(result.as_ref(), &expected);
+
+        result = bitwise_shift_right_scalar(&result, module).unwrap()?;
+        assert_eq!(result.as_ref(), &input);
+
+        Ok(())
+    }
 }