From 38c2531ea511dc6719c8a212032f3228f91c5a25 Mon Sep 17 00:00:00 2001 From: liukun4515 Date: Tue, 22 Feb 2022 09:38:57 +0800 Subject: [PATCH] support bitwise or/'|' operation --- datafusion-expr/src/operator.rs | 3 + .../coercion_rule/binary_rule.rs | 4 +- .../src/physical_plan/expressions/binary.rs | 73 ++++++++++++++++++- datafusion/src/sql/planner.rs | 1 + 4 files changed, 77 insertions(+), 4 deletions(-) diff --git a/datafusion-expr/src/operator.rs b/datafusion-expr/src/operator.rs index 0d3f17776c7cf..e4a9871e674d2 100644 --- a/datafusion-expr/src/operator.rs +++ b/datafusion-expr/src/operator.rs @@ -69,6 +69,8 @@ pub enum Operator { RegexNotIMatch, /// Bitwise and, like `&` BitwiseAnd, + /// Bitwise or, like `|` + BitwiseOr, } impl fmt::Display for Operator { @@ -96,6 +98,7 @@ impl fmt::Display for Operator { Operator::IsDistinctFrom => "IS DISTINCT FROM", Operator::IsNotDistinctFrom => "IS NOT DISTINCT FROM", Operator::BitwiseAnd => "&", + Operator::BitwiseOr => "|", }; write!(f, "{}", display) } diff --git a/datafusion/src/physical_plan/coercion_rule/binary_rule.rs b/datafusion/src/physical_plan/coercion_rule/binary_rule.rs index 7d4dd55bf605f..9419fd1187f64 100644 --- a/datafusion/src/physical_plan/coercion_rule/binary_rule.rs +++ b/datafusion/src/physical_plan/coercion_rule/binary_rule.rs @@ -31,7 +31,9 @@ pub(crate) fn coerce_types( ) -> Result { // This result MUST be compatible with `binary_coerce` let result = match op { - Operator::BitwiseAnd => bitwise_coercion(lhs_type, rhs_type), + Operator::BitwiseAnd | Operator::BitwiseOr => { + 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), diff --git a/datafusion/src/physical_plan/expressions/binary.rs b/datafusion/src/physical_plan/expressions/binary.rs index 6f9084ac18f20..0e5fafbcd99ed 100644 --- a/datafusion/src/physical_plan/expressions/binary.rs +++ b/datafusion/src/physical_plan/expressions/binary.rs @@ -461,6 +461,28 @@ fn bitwise_and(left: ArrayRef, right: ArrayRef) -> Result { } } +fn bitwise_or(left: ArrayRef, right: ArrayRef) -> Result { + match &left.data_type() { + DataType::Int8 => { + binary_bitwise_array_op!(left, right, |, Int8Array, i8) + } + DataType::Int16 => { + binary_bitwise_array_op!(left, right, |, Int16Array, i16) + } + DataType::Int32 => { + binary_bitwise_array_op!(left, right, |, Int32Array, i32) + } + DataType::Int64 => { + binary_bitwise_array_op!(left, right, |, Int64Array, i64) + } + other => Err(DataFusionError::Internal(format!( + "Data type {:?} not supported for binary operation '{}' on dyn arrays", + other, + Operator::BitwiseOr + ))), + } +} + fn bitwise_and_scalar( array: &dyn Array, scalar: ScalarValue, @@ -487,6 +509,29 @@ fn bitwise_and_scalar( Some(result) } +fn bitwise_or_scalar(array: &dyn Array, scalar: ScalarValue) -> Option> { + let result = match array.data_type() { + DataType::Int8 => { + binary_bitwise_array_scalar!(array, scalar, |, Int8Array, i8) + } + DataType::Int16 => { + binary_bitwise_array_scalar!(array, scalar, |, Int16Array, i16) + } + DataType::Int32 => { + binary_bitwise_array_scalar!(array, scalar, |, Int32Array, i32) + } + DataType::Int64 => { + binary_bitwise_array_scalar!(array, scalar, |, Int64Array, i64) + } + other => Err(DataFusionError::Internal(format!( + "Data type {:?} not supported for binary operation '{}' on dyn arrays", + other, + Operator::BitwiseOr + ))), + }; + Some(result) +} + /// Binary expression #[derive(Debug)] pub struct BinaryExpr { @@ -1020,7 +1065,7 @@ pub fn binary_operator_data_type( | Operator::IsDistinctFrom | Operator::IsNotDistinctFrom => Ok(DataType::Boolean), // bitwise operations return the common coerced type - Operator::BitwiseAnd => Ok(result_type), + Operator::BitwiseAnd | Operator::BitwiseOr => Ok(result_type), // math operations return the same value as the common coerced type Operator::Plus | Operator::Minus @@ -1197,6 +1242,7 @@ impl BinaryExpr { true ), Operator::BitwiseAnd => bitwise_and_scalar(array, scalar.clone()), + Operator::BitwiseOr => bitwise_or_scalar(array, scalar.clone()), // if scalar operation is not supported - fallback to array implementation _ => None, }; @@ -1286,6 +1332,7 @@ impl BinaryExpr { binary_string_array_flag_op!(left, right, regexp_is_match, true, true) } Operator::BitwiseAnd => bitwise_and(left, right), + Operator::BitwiseOr => bitwise_or(left, right), } } } @@ -1735,6 +1782,18 @@ mod tests { DataType::Int64, vec![0i64, 0i64, 1i64] ); + test_coercion!( + Int16Array, + DataType::Int16, + vec![1i16, 2i16, 3i16], + Int64Array, + DataType::Int64, + vec![10i64, 4i64, 5i64], + Operator::BitwiseOr, + Int64Array, + DataType::Int64, + vec![11i64, 6i64, 7i64] + ); Ok(()) } @@ -3115,9 +3174,13 @@ 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 result = bitwise_and(left, right)?; + let mut result = bitwise_and(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())?; + let expected = Int32Array::from(vec![Some(13), None, Some(15)]); + assert_eq!(result.as_ref(), &expected); Ok(()) } @@ -3125,9 +3188,13 @@ 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 result = bitwise_and_scalar(&left, right).unwrap()?; + let mut result = bitwise_and_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).unwrap()?; + let expected = Int32Array::from(vec![Some(15), None, Some(11)]); + assert_eq!(result.as_ref(), &expected); Ok(()) } } diff --git a/datafusion/src/sql/planner.rs b/datafusion/src/sql/planner.rs index 8b59ccdca6441..c365ac981d245 100644 --- a/datafusion/src/sql/planner.rs +++ b/datafusion/src/sql/planner.rs @@ -1299,6 +1299,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { BinaryOperator::PGRegexNotMatch => Ok(Operator::RegexNotMatch), BinaryOperator::PGRegexNotIMatch => Ok(Operator::RegexNotIMatch), BinaryOperator::BitwiseAnd => Ok(Operator::BitwiseAnd), + BinaryOperator::BitwiseOr => Ok(Operator::BitwiseOr), _ => Err(DataFusionError::NotImplemented(format!( "Unsupported SQL binary operator {:?}", op