feat: add optimization rules for bitwise operations

datafusion/common/src/scalar.rs

@@ -1019,7 +1019,7 @@ impl ScalarValue {
         Self::List(scalars, Box::new(Field::new("item", child_type, true)))
     }
 
-    // Create a zero value in the given type.
+    /// Create a zero value in the given type.
     pub fn new_zero(datatype: &DataType) -> Result<ScalarValue> {
         assert!(datatype.is_primitive());
         Ok(match datatype {
@@ -1042,6 +1042,24 @@ impl ScalarValue {
         })
     }
 
+    /// Create a negative one value in the given type.
+    pub fn new_negative_one(datatype: &DataType) -> Result<ScalarValue> {
+        assert!(datatype.is_primitive());
+        Ok(match datatype {
+            DataType::Int8 | DataType::UInt8 => ScalarValue::Int8(Some(-1)),
+            DataType::Int16 | DataType::UInt16 => ScalarValue::Int16(Some(-1)),
+            DataType::Int32 | DataType::UInt32 => ScalarValue::Int32(Some(-1)),
+            DataType::Int64 | DataType::UInt64 => ScalarValue::Int64(Some(-1)),
+            DataType::Float32 => ScalarValue::Float32(Some(-1.0)),
+            DataType::Float64 => ScalarValue::Float64(Some(-1.0)),
+            _ => {
+                return Err(DataFusionError::NotImplemented(format!(
+                    "Can't create a negative one scalar from data_type \"{datatype:?}\""
+                )));
+            }
+        })
+    }
+
     /// Getter for the `DataType` of the value
     pub fn get_datatype(&self) -> DataType {
         match self {

datafusion/core/tests/simplification.rs

@@ -49,6 +49,13 @@ impl SimplifyInfo for MyInfo {
     fn execution_props(&self) -> &ExecutionProps {
         &self.execution_props
     }
+
+    fn get_data_type(&self, expr: &Expr) -> Result<DataType> {
+        match expr.get_type(&self.schema) {
+            Ok(expr_data_type) => Ok(expr_data_type),
+            Err(e) => Err(e),
+        }
+    }
 }
 
 impl From<DFSchema> for MyInfo {

datafusion/expr/src/expr.rs

@@ -639,6 +639,31 @@ impl Expr {
         binary_expr(self, Operator::Or, other)
     }
 
+    /// Return `self & other`
+    pub fn bitwise_and(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::BitwiseAnd, other)
+    }
+
+    /// Return `self | other`
+    pub fn bitwise_or(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::BitwiseOr, other)
+    }
+
+    /// Return `self ^ other`
+    pub fn bitwise_xor(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::BitwiseXor, other)
+    }
+
+    /// Return `self >> other`
+    pub fn bitwise_shift_right(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::BitwiseShiftRight, other)
+    }
+
+    /// Return `self << other`
+    pub fn bitwise_shift_left(self, other: Expr) -> Expr {
+        binary_expr(self, Operator::BitwiseShiftLeft, other)
+    }
+
     /// Return `!self`
     #[allow(clippy::should_implement_trait)]
     pub fn not(self) -> Expr {

datafusion/expr/src/expr_fn.rs

@@ -112,6 +112,51 @@ pub fn count(expr: Expr) -> Expr {
     ))
 }
 
+/// Return a new expression with bitwise AND
+pub fn bitwise_and(left: Expr, right: Expr) -> Expr {
+    Expr::BinaryExpr(BinaryExpr::new(
+        Box::new(left),
+        Operator::BitwiseAnd,
+        Box::new(right),
+    ))
+}
+
+/// Return a new expression with bitwise OR
+pub fn bitwise_or(left: Expr, right: Expr) -> Expr {
+    Expr::BinaryExpr(BinaryExpr::new(
+        Box::new(left),
+        Operator::BitwiseOr,
+        Box::new(right),
+    ))
+}
+
+/// Return a new expression with bitwise XOR
+pub fn bitwise_xor(left: Expr, right: Expr) -> Expr {
+    Expr::BinaryExpr(BinaryExpr::new(
+        Box::new(left),
+        Operator::BitwiseXor,
+        Box::new(right),
+    ))
+}
+
+/// Return a new expression with bitwise SHIFT RIGHT
+pub fn bitwise_shift_right(left: Expr, right: Expr) -> Expr {
+    Expr::BinaryExpr(BinaryExpr::new(
+        Box::new(left),
+        Operator::BitwiseShiftRight,
+        Box::new(right),
+    ))
+}
+
+/// Return a new expression with bitwise SHIFT LEFT
+pub fn bitwise_shift_left(left: Expr, right: Expr) -> Expr {
+    Expr::BinaryExpr(BinaryExpr::new(
+        Box::new(left),
+        Operator::BitwiseShiftLeft,
+        Box::new(right),
+    ))
+}
+
 /// Create an expression to represent the count(distinct) aggregate function
 pub fn count_distinct(expr: Expr) -> Expr {
     Expr::AggregateFunction(AggregateFunction::new(

datafusion/optimizer/src/simplify_expressions/context.rs

@@ -38,6 +38,9 @@ pub trait SimplifyInfo {
 
     /// Returns details needed for partial expression evaluation
     fn execution_props(&self) -> &ExecutionProps;
+
+    /// Returns data type of this expr needed for determining optimized int type of a value
+    fn get_data_type(&self, expr: &Expr) -> Result<DataType>;
 }
 
 /// Provides simplification information based on DFSchema and
@@ -123,6 +126,21 @@ impl<'a> SimplifyInfo for SimplifyContext<'a> {
             })
     }
 
+    /// Returns data type of this expr needed for determining optimized int type of a value
+    fn get_data_type(&self, expr: &Expr) -> Result<DataType> {
+        if self.schemas.len() == 1 {
+            match expr.get_type(&self.schemas[0]) {
+                Ok(expr_data_type) => Ok(expr_data_type),
+                Err(e) => Err(e),
+            }
+        } else {
+            Err(DataFusionError::Internal(
+                "The expr has more than one schema, could not determine data type"
+                    .to_string(),
+            ))
+        }
+    }
+
     fn execution_props(&self) -> &ExecutionProps {
         self.props
     }