This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git
The following commit(s) were added to refs/heads/master by this push:
new 3261e2b support bitwise or/'|' operation (#1876)
3261e2b is described below
commit 3261e2bcdfa13e460d7a11a64dd72f272c7d872b
Author: Kun Liu <[email protected]>
AuthorDate: Fri Mar 4 20:54:50 2022 +0800
support bitwise or/'|' operation (#1876)
---
datafusion-expr/src/operator.rs | 3 +
.../src/coercion_rule/binary_rule.rs | 4 +-
datafusion-physical-expr/src/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 0d3f177..e4a9871 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 67a052b..2207079 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<DataType> {
// 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 0999077..6b40c8f 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<ArrayRef> {
}
}
+fn bitwise_or(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
+ 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<Result<ArrayRef>> {
+ 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 8b59ccd..c365ac9 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
