diff --git a/datafusion-expr/src/operator.rs b/datafusion-expr/src/operator.rs index 0d3f17776c7c..e4a9871e674d 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-physical-expr/src/coercion_rule/binary_rule.rs b/datafusion-physical-expr/src/coercion_rule/binary_rule.rs index 67a052b86de2..22070794bda0 100644 --- a/datafusion-physical-expr/src/coercion_rule/binary_rule.rs +++ b/datafusion-physical-expr/src/coercion_rule/binary_rule.rs @@ -32,7 +32,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-physical-expr/src/expressions/binary.rs b/datafusion-physical-expr/src/expressions/binary.rs index 09990773b9a9..6b40c8f5af83 100644 --- a/datafusion-physical-expr/src/expressions/binary.rs +++ b/datafusion-physical-expr/src/expressions/binary.rs @@ -462,6 +462,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, @@ -488,6 +510,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 { @@ -1021,7 +1066,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 @@ -1198,6 +1243,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, }; @@ -1287,6 +1333,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 8b59ccdca644..c365ac981d24 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