suppport bitwise and as an example

Author: liukun4515

datafusion/src/logical_plan/operators.rs

@@ -64,6 +64,8 @@ pub enum Operator {
     RegexNotMatch,
     /// Case insensitive regex not match
     RegexNotIMatch,
+    /// Bitwise and, like `&`
+    BitwiseAnd,
 }
 
 impl fmt::Display for Operator {
@@ -90,6 +92,7 @@ impl fmt::Display for Operator {
             Operator::RegexNotIMatch => "!~*",
             Operator::IsDistinctFrom => "IS DISTINCT FROM",
             Operator::IsNotDistinctFrom => "IS NOT DISTINCT FROM",
+            Operator::BitwiseAnd => "&",
         };
         write!(f, "{}", display)
     }

datafusion/src/physical_plan/coercion_rule/binary_rule.rs

@@ -31,6 +31,7 @@ pub(crate) fn coerce_types(
 ) -> Result<DataType> {
     // This result MUST be compatible with `binary_coerce`
     let result = match op {
+        Operator::BitwiseAnd => 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),
@@ -72,6 +73,25 @@ pub(crate) 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) {
+        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),
+        (Int32, _) | (_, Int32) => Some(Int32),
+        (Int16, _) | (_, Int16) => Some(Int16),
+        (Int8, _) | (_, Int8) => Some(Int8),
+        _ => None,
+    }
+}
+
 fn comparison_eq_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataType> {
     // can't compare dictionaries directly due to
     // https://github.com/apache/arrow-rs/issues/1201

datafusion/src/physical_plan/expressions/binary.rs

@@ -334,6 +334,141 @@ fn modulus_decimal(left: &DecimalArray, right: &DecimalArray) -> Result<DecimalA
     Ok(decimal_builder.finish())
 }
 
+fn bitwise_and(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
+    let len = left.len();
+    match &left.data_type() {
+        DataType::Int8 => {
+            let left = left.as_any().downcast_ref::<Int8Array>().unwrap();
+            let right = right.as_any().downcast_ref::<Int8Array>().unwrap();
+            let result = (0..len)
+                .into_iter()
+                .map(|i| {
+                    if left.is_null(i) || right.is_null(i) {
+                        None
+                    } else {
+                        Some(left.value(i) & right.value(i))
+                    }
+                })
+                .collect::<Int8Array>();
+            Ok(Arc::new(result))
+        }
+        DataType::Int16 => {
+            let left = left.as_any().downcast_ref::<Int16Array>().unwrap();
+            let right = right.as_any().downcast_ref::<Int16Array>().unwrap();
+            let result = (0..len)
+                .into_iter()
+                .map(|i| {
+                    if left.is_null(i) || right.is_null(i) {
+                        None
+                    } else {
+                        Some(left.value(i) & right.value(i))
+                    }
+                })
+                .collect::<Int16Array>();
+            Ok(Arc::new(result))
+        }
+        DataType::Int32 => {
+            let left = left.as_any().downcast_ref::<Int32Array>().unwrap();
+            let right = right.as_any().downcast_ref::<Int32Array>().unwrap();
+            let result = (0..len)
+                .into_iter()
+                .map(|i| {
+                    if left.is_null(i) || right.is_null(i) {
+                        None
+                    } else {
+                        Some(left.value(i) & right.value(i))
+                    }
+                })
+                .collect::<Int32Array>();
+            Ok(Arc::new(result))
+        }
+        DataType::Int64 => {
+            let left = left.as_any().downcast_ref::<Int64Array>().unwrap();
+            let right = right.as_any().downcast_ref::<Int64Array>().unwrap();
+            let result = (0..len)
+                .into_iter()
+                .map(|i| {
+                    if left.is_null(i) || right.is_null(i) {
+                        None
+                    } else {
+                        Some(left.value(i) & right.value(i))
+                    }
+                })
+                .collect::<Int64Array>();
+            Ok(Arc::new(result))
+        }
+        other => Err(DataFusionError::Internal(format!(
+            "Data type {:?} not supported for binary operation '{}' on dyn arrays",
+            other,
+            Operator::BitwiseAnd
+        ))),
+    }
+}
+
+fn bitwise_and_scalar(array: &ArrayRef, scalar: ScalarValue) -> Option<Result<ArrayRef>> {
+    let len = array.len();
+    let result = match array.data_type() {
+        DataType::Int8 => {
+            let array = array.as_any().downcast_ref::<Int8Array>().unwrap();
+            if scalar.is_null() {
+                Ok(new_null_array(array.data_type(), len))
+            } else {
+                let right: i8 = scalar.try_into().unwrap();
+                let result = (0..len)
+                    .into_iter()
+                    .map(|i| Some(array.value(i) & right))
+                    .collect::<Int8Array>();
+                Ok(Arc::new(result) as ArrayRef)
+            }
+        }
+        DataType::Int16 => {
+            let array = array.as_any().downcast_ref::<Int16Array>().unwrap();
+            if scalar.is_null() {
+                Ok(new_null_array(array.data_type(), len))
+            } else {
+                let right: i16 = scalar.try_into().unwrap();
+                let result = (0..len)
+                    .into_iter()
+                    .map(|i| Some(array.value(i) & right))
+                    .collect::<Int16Array>();
+                Ok(Arc::new(result) as ArrayRef)
+            }
+        }
+        DataType::Int32 => {
+            let array = array.as_any().downcast_ref::<Int32Array>().unwrap();
+            if scalar.is_null() {
+                Ok(new_null_array(array.data_type(), len))
+            } else {
+                let right: i32 = scalar.try_into().unwrap();
+                let result = (0..len)
+                    .into_iter()
+                    .map(|i| Some(array.value(i) & right))
+                    .collect::<Int32Array>();
+                Ok(Arc::new(result) as ArrayRef)
+            }
+        }
+        DataType::Int64 => {
+            let array = array.as_any().downcast_ref::<Int64Array>().unwrap();
+            if scalar.is_null() {
+                Ok(new_null_array(array.data_type(), len))
+            } else {
+                let right: i64 = scalar.try_into().unwrap();
+                let result = (0..len)
+                    .into_iter()
+                    .map(|i| Some(array.value(i) & right))
+                    .collect::<Int64Array>();
+                Ok(Arc::new(result) as ArrayRef)
+            }
+        }
+        other => Err(DataFusionError::Internal(format!(
+            "Data type {:?} not supported for binary operation '{}' on dyn arrays",
+            other,
+            Operator::BitwiseAnd
+        ))),
+    };
+    Some(result)
+}
+
 /// Binary expression
 #[derive(Debug)]
 pub struct BinaryExpr {
@@ -564,6 +699,14 @@ macro_rules! binary_primitive_array_op {
     }};
 }
 
+/// The binary_bitwise_array_op macro only evaluates for integer types
+/// like int64, int32
+// macro_rules! binary_bitwise_array_op {
+//     ($LEFT:expr, $RIGHT: expr, $OP:ident) => {{
+//
+//     }};
+// }
+
 /// Invoke a compute kernel on an array and a scalar
 /// The binary_primitive_array_op_scalar macro only evaluates for primitive
 /// types like integers and floats.
@@ -811,6 +954,8 @@ pub fn binary_operator_data_type(
         | Operator::RegexNotIMatch
         | Operator::IsDistinctFrom
         | Operator::IsNotDistinctFrom => Ok(DataType::Boolean),
+        // bitwise operations return the common coerced type
+        Operator::BitwiseAnd => Ok(result_type),
         // math operations return the same value as the common coerced type
         Operator::Plus
         | Operator::Minus
@@ -939,6 +1084,7 @@ impl BinaryExpr {
                 true,
                 true
             ),
+            Operator::BitwiseAnd => bitwise_and_scalar(array, scalar.clone()),
             // if scalar operation is not supported - fallback to array implementation
             _ => None,
         };
@@ -1027,6 +1173,7 @@ impl BinaryExpr {
             Operator::RegexNotIMatch => {
                 binary_string_array_flag_op!(left, right, regexp_is_match, true, true)
             }
+            Operator::BitwiseAnd => bitwise_and(left, right),
         }
     }
 }

datafusion/src/sql/planner.rs

@@ -1250,6 +1250,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             BinaryOperator::PGRegexIMatch => Ok(Operator::RegexIMatch),
             BinaryOperator::PGRegexNotMatch => Ok(Operator::RegexNotMatch),
             BinaryOperator::PGRegexNotIMatch => Ok(Operator::RegexNotIMatch),
+            BinaryOperator::BitwiseAnd => Ok(Operator::BitwiseAnd),
             _ => Err(DataFusionError::NotImplemented(format!(
                 "Unsupported SQL binary operator {:?}",
                 op