From 6c3cfe8a7bce9340a1c5d2a3345966e3a3865e9a Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 23 Apr 2023 07:12:51 -0400 Subject: [PATCH 1/3] Use arrow kernels for bitwise operations --- .../physical-expr/src/expressions/binary.rs | 72 +-- .../src/expressions/binary/kernels.rs | 455 ++++-------------- 2 files changed, 139 insertions(+), 388 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index b98a2dc420a6..e873fa0ab57f 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -63,9 +63,9 @@ use datafusion_common::scalar::{ }; use datafusion_expr::type_coercion::{is_decimal, is_timestamp, is_utf8_or_large_utf8}; use kernels::{ - bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_shift_left, - bitwise_shift_left_scalar, bitwise_shift_right, bitwise_shift_right_scalar, - bitwise_xor, bitwise_xor_scalar, + bitwise_and_dyn, bitwise_and_dyn_scalar, bitwise_or_dyn, bitwise_or_dyn_scalar, + bitwise_shift_left_dyn, bitwise_shift_left_dyn_scalar, bitwise_shift_right_dyn, + bitwise_shift_right_dyn_scalar, bitwise_xor_dyn, bitwise_xor_dyn_scalar, }; use kernels_arrow::{ add_decimal_dyn_scalar, add_dyn_decimal, divide_decimal_dyn_scalar, @@ -1134,11 +1134,11 @@ impl BinaryExpr { true, true ), - Operator::BitwiseAnd => bitwise_and_scalar(array, scalar), - Operator::BitwiseOr => bitwise_or_scalar(array, scalar), - Operator::BitwiseXor => bitwise_xor_scalar(array, scalar), - Operator::BitwiseShiftRight => bitwise_shift_right_scalar(array, scalar), - Operator::BitwiseShiftLeft => bitwise_shift_left_scalar(array, scalar), + Operator::BitwiseAnd => bitwise_and_dyn_scalar(array, scalar), + Operator::BitwiseOr => bitwise_or_dyn_scalar(array, scalar), + Operator::BitwiseXor => bitwise_xor_dyn_scalar(array, scalar), + Operator::BitwiseShiftRight => bitwise_shift_right_dyn_scalar(array, scalar), + Operator::BitwiseShiftLeft => bitwise_shift_left_dyn_scalar(array, scalar), // if scalar operation is not supported - fallback to array implementation _ => None, }; @@ -1256,11 +1256,11 @@ impl BinaryExpr { Operator::RegexNotIMatch => { binary_string_array_flag_op!(left, right, regexp_is_match, true, true) } - Operator::BitwiseAnd => bitwise_and(left, right), - Operator::BitwiseOr => bitwise_or(left, right), - Operator::BitwiseXor => bitwise_xor(left, right), - Operator::BitwiseShiftRight => bitwise_shift_right(left, right), - Operator::BitwiseShiftLeft => bitwise_shift_left(left, right), + Operator::BitwiseAnd => bitwise_and_dyn(left, right), + Operator::BitwiseOr => bitwise_or_dyn(left, right), + Operator::BitwiseXor => bitwise_xor_dyn(left, right), + Operator::BitwiseShiftRight => bitwise_shift_right_dyn(left, right), + Operator::BitwiseShiftLeft => bitwise_shift_left_dyn(left, right), Operator::StringConcat => { binary_string_array_op!(left, right, concat_elements) } @@ -5117,15 +5117,15 @@ 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 mut result = bitwise_and(left.clone(), right.clone())?; + let mut result = bitwise_and_dyn(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())?; + result = bitwise_or_dyn(left.clone(), right.clone())?; let expected = Int32Array::from(vec![Some(13), None, Some(15)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_xor(left.clone(), right.clone())?; + result = bitwise_xor_dyn(left.clone(), right.clone())?; let expected = Int32Array::from(vec![Some(13), None, Some(12)]); assert_eq!(result.as_ref(), &expected); @@ -5133,15 +5133,15 @@ mod tests { Arc::new(UInt32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef; let right = Arc::new(UInt32Array::from(vec![Some(1), Some(3), Some(7)])) as ArrayRef; - let mut result = bitwise_and(left.clone(), right.clone())?; + let mut result = bitwise_and_dyn(left.clone(), right.clone())?; let expected = UInt32Array::from(vec![Some(0), None, Some(3)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_or(left.clone(), right.clone())?; + result = bitwise_or_dyn(left.clone(), right.clone())?; let expected = UInt32Array::from(vec![Some(13), None, Some(15)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_xor(left.clone(), right.clone())?; + result = bitwise_xor_dyn(left.clone(), right.clone())?; let expected = UInt32Array::from(vec![Some(13), None, Some(12)]); assert_eq!(result.as_ref(), &expected); @@ -5153,24 +5153,24 @@ mod tests { let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(10)])) as ArrayRef; let modules = Arc::new(Int32Array::from(vec![Some(2), Some(4), Some(8)])) as ArrayRef; - let mut result = bitwise_shift_left(input.clone(), modules.clone())?; + let mut result = bitwise_shift_left_dyn(input.clone(), modules.clone())?; let expected = Int32Array::from(vec![Some(8), None, Some(2560)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_shift_right(result.clone(), modules.clone())?; + result = bitwise_shift_right_dyn(result.clone(), modules.clone())?; assert_eq!(result.as_ref(), &input); let input = Arc::new(UInt32Array::from(vec![Some(2), None, Some(10)])) as ArrayRef; let modules = Arc::new(UInt32Array::from(vec![Some(2), Some(4), Some(8)])) as ArrayRef; - let mut result = bitwise_shift_left(input.clone(), modules.clone())?; + let mut result = bitwise_shift_left_dyn(input.clone(), modules.clone())?; let expected = UInt32Array::from(vec![Some(8), None, Some(2560)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_shift_right(result.clone(), modules.clone())?; + result = bitwise_shift_right_dyn(result.clone(), modules.clone())?; assert_eq!(result.as_ref(), &input); Ok(()) } @@ -5179,14 +5179,14 @@ mod tests { fn bitwise_shift_array_overflow_test() -> Result<()> { let input = Arc::new(Int32Array::from(vec![Some(2)])) as ArrayRef; let modules = Arc::new(Int32Array::from(vec![Some(100)])) as ArrayRef; - let result = bitwise_shift_left(input.clone(), modules.clone())?; + let result = bitwise_shift_left_dyn(input.clone(), modules.clone())?; let expected = Int32Array::from(vec![Some(32)]); assert_eq!(result.as_ref(), &expected); let input = Arc::new(UInt32Array::from(vec![Some(2)])) as ArrayRef; let modules = Arc::new(UInt32Array::from(vec![Some(100)])) as ArrayRef; - let result = bitwise_shift_left(input.clone(), modules.clone())?; + let result = bitwise_shift_left_dyn(input.clone(), modules.clone())?; let expected = UInt32Array::from(vec![Some(32)]); assert_eq!(result.as_ref(), &expected); @@ -5197,30 +5197,30 @@ 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 mut result = bitwise_and_scalar(&left, right.clone()).unwrap()?; + let mut result = bitwise_and_dyn_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.clone()).unwrap()?; + result = bitwise_or_dyn_scalar(&left, right.clone()).unwrap()?; let expected = Int32Array::from(vec![Some(15), None, Some(11)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_xor_scalar(&left, right).unwrap()?; + result = bitwise_xor_dyn_scalar(&left, right).unwrap()?; let expected = Int32Array::from(vec![Some(15), None, Some(8)]); assert_eq!(result.as_ref(), &expected); let left = Arc::new(UInt32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef; let right = ScalarValue::from(3u32); - let mut result = bitwise_and_scalar(&left, right.clone()).unwrap()?; + let mut result = bitwise_and_dyn_scalar(&left, right.clone()).unwrap()?; let expected = UInt32Array::from(vec![Some(0), None, Some(3)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_or_scalar(&left, right.clone()).unwrap()?; + result = bitwise_or_dyn_scalar(&left, right.clone()).unwrap()?; let expected = UInt32Array::from(vec![Some(15), None, Some(11)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_xor_scalar(&left, right).unwrap()?; + result = bitwise_xor_dyn_scalar(&left, right).unwrap()?; let expected = UInt32Array::from(vec![Some(15), None, Some(8)]); assert_eq!(result.as_ref(), &expected); Ok(()) @@ -5230,22 +5230,24 @@ mod tests { fn bitwise_shift_scalar_test() -> Result<()> { let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef; let module = ScalarValue::from(10i32); - let mut result = bitwise_shift_left_scalar(&input, module.clone()).unwrap()?; + let mut result = + bitwise_shift_left_dyn_scalar(&input, module.clone()).unwrap()?; let expected = Int32Array::from(vec![Some(2048), None, Some(4096)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_shift_right_scalar(&result, module).unwrap()?; + result = bitwise_shift_right_dyn_scalar(&result, module).unwrap()?; assert_eq!(result.as_ref(), &input); let input = Arc::new(UInt32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef; let module = ScalarValue::from(10u32); - let mut result = bitwise_shift_left_scalar(&input, module.clone()).unwrap()?; + let mut result = + bitwise_shift_left_dyn_scalar(&input, module.clone()).unwrap()?; let expected = UInt32Array::from(vec![Some(2048), None, Some(4096)]); assert_eq!(result.as_ref(), &expected); - result = bitwise_shift_right_scalar(&result, module).unwrap()?; + result = bitwise_shift_right_dyn_scalar(&result, module).unwrap()?; assert_eq!(result.as_ref(), &input); Ok(()) } diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index 6284f9bed473..4f3504ab8134 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -19,6 +19,10 @@ //! datafusion and not (yet) targeted to port upstream to arrow use arrow::array::*; use arrow::compute::binary; +use arrow::compute::kernels::bitwise::{ + bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_xor, + bitwise_xor_scalar, +}; use arrow::datatypes::DataType; use datafusion_common::{DataFusionError, Result, ScalarValue}; use datafusion_expr::Operator; @@ -56,41 +60,64 @@ macro_rules! binary_bitwise_array_scalar { }}; } -pub(crate) fn bitwise_and(left: ArrayRef, right: ArrayRef) -> Result { - match &left.data_type() { - DataType::Int8 => { - binary_bitwise_array_op!(left, right, |a: i8, b: i8| a & b, Int8Array) - } - DataType::Int16 => { - binary_bitwise_array_op!(left, right, |a: i16, b: i16| a & b, Int16Array) - } - DataType::Int32 => { - binary_bitwise_array_op!(left, right, |a: i32, b: i32| a & b, Int32Array) - } - DataType::Int64 => { - binary_bitwise_array_op!(left, right, |a: i64, b: i64| a & b, Int64Array) - } - DataType::UInt8 => { - binary_bitwise_array_op!(left, right, |a: u8, b: u8| a & b, UInt8Array) - } - DataType::UInt16 => { - binary_bitwise_array_op!(left, right, |a: u16, b: u16| a & b, UInt16Array) - } - DataType::UInt32 => { - binary_bitwise_array_op!(left, right, |a: u32, b: u32| a & b, UInt32Array) - } - DataType::UInt64 => { - binary_bitwise_array_op!(left, right, |a: u64, b: u64| a & b, UInt64Array) +/// Downcasts $LEFT and $RIGHT to $ARRAY_TYPE and then calls $KERNEL($LEFT, $RIGHT) +macro_rules! call_bitwise_kernel { + ($LEFT:expr, $RIGHT:expr, $KERNEL:expr, $ARRAY_TYPE:ident) => {{ + let left = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); + let right = $RIGHT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); + let result: $ARRAY_TYPE = $KERNEL(left, right)?; + Ok(Arc::new(result)) + }}; +} + +/// Creates a $FUNC(left: ArrayRef, right: ArrayRef) that +/// downcasts left / right to the appropriate integral type and calls the kernel +macro_rules! create_dyn_kernel { + ($FUNC:ident, $KERNEL:ident) => { + pub(crate) fn $FUNC(left: ArrayRef, right: ArrayRef) -> Result { + match &left.data_type() { + DataType::Int8 => { + call_bitwise_kernel!(left, right, $KERNEL, Int8Array) + } + DataType::Int16 => { + call_bitwise_kernel!(left, right, $KERNEL, Int16Array) + } + DataType::Int32 => { + call_bitwise_kernel!(left, right, $KERNEL, Int32Array) + } + DataType::Int64 => { + call_bitwise_kernel!(left, right, $KERNEL, Int64Array) + } + DataType::UInt8 => { + call_bitwise_kernel!(left, right, $KERNEL, UInt8Array) + } + DataType::UInt16 => { + call_bitwise_kernel!(left, right, $KERNEL, UInt16Array) + } + DataType::UInt32 => { + call_bitwise_kernel!(left, right, $KERNEL, UInt32Array) + } + DataType::UInt64 => { + call_bitwise_kernel!(left, right, $KERNEL, UInt64Array) + } + other => Err(DataFusionError::Internal(format!( + "Data type {:?} not supported for binary operation '{}' on dyn arrays", + other, + stringify!($KERNEL), + ))), + } } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseAnd - ))), - } + }; } -pub(crate) fn bitwise_shift_right(left: ArrayRef, right: ArrayRef) -> Result { +create_dyn_kernel!(bitwise_or_dyn, bitwise_or); +create_dyn_kernel!(bitwise_xor_dyn, bitwise_xor); +create_dyn_kernel!(bitwise_and_dyn, bitwise_and); + +pub(crate) fn bitwise_shift_right_dyn( + left: ArrayRef, + right: ArrayRef, +) -> Result { match &left.data_type() { DataType::Int8 => { binary_bitwise_array_op!( @@ -164,7 +191,10 @@ pub(crate) fn bitwise_shift_right(left: ArrayRef, right: ArrayRef) -> Result Result { +pub(crate) fn bitwise_shift_left_dyn( + left: ArrayRef, + right: ArrayRef, +) -> Result { match &left.data_type() { DataType::Int8 => { binary_bitwise_array_op!( @@ -238,333 +268,52 @@ pub(crate) fn bitwise_shift_left(left: ArrayRef, right: ArrayRef) -> Result Result { - match &left.data_type() { - DataType::Int8 => { - binary_bitwise_array_op!(left, right, |a: i8, b: i8| a | b, Int8Array) - } - DataType::Int16 => { - binary_bitwise_array_op!(left, right, |a: i16, b: i16| a | b, Int16Array) - } - DataType::Int32 => { - binary_bitwise_array_op!(left, right, |a: i32, b: i32| a | b, Int32Array) - } - DataType::Int64 => { - binary_bitwise_array_op!(left, right, |a: i64, b: i64| a | b, Int64Array) - } - DataType::UInt8 => { - binary_bitwise_array_op!(left, right, |a: u8, b: u8| a | b, UInt8Array) - } - DataType::UInt16 => { - binary_bitwise_array_op!(left, right, |a: u16, b: u16| a | b, UInt16Array) - } - DataType::UInt32 => { - binary_bitwise_array_op!(left, right, |a: u32, b: u32| a | b, UInt32Array) - } - DataType::UInt64 => { - binary_bitwise_array_op!(left, right, |a: u64, b: u64| a | b, UInt64Array) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseOr - ))), - } -} - -pub(crate) fn bitwise_xor(left: ArrayRef, right: ArrayRef) -> Result { - match &left.data_type() { - DataType::Int8 => { - binary_bitwise_array_op!(left, right, |a: i8, b: i8| a ^ b, Int8Array) - } - DataType::Int16 => { - binary_bitwise_array_op!(left, right, |a: i16, b: i16| a ^ b, Int16Array) - } - DataType::Int32 => { - binary_bitwise_array_op!(left, right, |a: i32, b: i32| a ^ b, Int32Array) - } - DataType::Int64 => { - binary_bitwise_array_op!(left, right, |a: i64, b: i64| a ^ b, Int64Array) - } - DataType::UInt8 => { - binary_bitwise_array_op!(left, right, |a: u8, b: u8| a ^ b, UInt8Array) - } - DataType::UInt16 => { - binary_bitwise_array_op!(left, right, |a: u16, b: u16| a ^ b, UInt16Array) - } - DataType::UInt32 => { - binary_bitwise_array_op!(left, right, |a: u32, b: u32| a ^ b, UInt32Array) - } - DataType::UInt64 => { - binary_bitwise_array_op!(left, right, |a: u64, b: u64| a ^ b, UInt64Array) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseXor - ))), - } -} - -pub(crate) fn bitwise_and_scalar( - array: &dyn Array, - scalar: ScalarValue, -) -> Option> { - let result = match array.data_type() { - DataType::Int8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i8| move |b: i8| a & b, - Int8Array, - i8 - ) - } - DataType::Int16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i16| move |b: i16| a & b, - Int16Array, - i16 - ) - } - DataType::Int32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i32| move |b: i32| a & b, - Int32Array, - i32 - ) - } - DataType::Int64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i64| move |b: i64| a & b, - Int64Array, - i64 - ) - } - DataType::UInt8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u8| move |b: u8| a & b, - UInt8Array, - u8 - ) - } - DataType::UInt16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u16| move |b: u16| a & b, - UInt16Array, - u16 - ) - } - DataType::UInt32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u32| a & b, - UInt32Array, - u32 - ) - } - DataType::UInt64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u64| move |b: u64| a & b, - UInt64Array, - u64 - ) +/// Downcasts $LEFT as $ARRAY_TYPE and $RIGHT as TYPE and calls $KERNEL($LEFT, RIGHT) +macro_rules! call_bitwise_scalar_kernel { + ($LEFT:expr, $RIGHT:expr, $KERNEL:ident, $ARRAY_TYPE:ident, $TYPE:ty) => {{ + let len = $LEFT.len(); + let array = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); + let scalar = $RIGHT; + if scalar.is_null() { + Ok(new_null_array(array.data_type(), len)) + } else { + let scalar: $TYPE = scalar.try_into().unwrap(); + let result: $ARRAY_TYPE = $KERNEL(array, scalar).unwrap(); + 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) + }}; } -pub(crate) fn bitwise_or_scalar( - array: &dyn Array, - scalar: ScalarValue, -) -> Option> { - let result = match array.data_type() { - DataType::Int8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i8| move |b: i8| a | b, - Int8Array, - i8 - ) - } - DataType::Int16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i16| move |b: i16| a | b, - Int16Array, - i16 - ) - } - DataType::Int32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i32| move |b: i32| a | b, - Int32Array, - i32 - ) - } - DataType::Int64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i64| move |b: i64| a | b, - Int64Array, - i64 - ) - } - DataType::UInt8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u8| move |b: u8| a | b, - UInt8Array, - u8 - ) - } - DataType::UInt16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u16| move |b: u16| a | b, - UInt16Array, - u16 - ) +/// Creates a $FUNC(left: ArrayRef, right: ArrayRef) that +/// downcasts left / right to the appropriate integral type and calls the kernel +macro_rules! create_dyn_scalar_kernel { + ($FUNC:ident, $KERNEL:ident) => { + pub(crate) fn $FUNC(array: &dyn Array, scalar: ScalarValue) -> Option> { + let result = match array.data_type() { + DataType::Int8 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int8Array, i8), + DataType::Int16 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int16Array, i16), + DataType::Int32 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int32Array, i32), + DataType::Int64 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int64Array, i64), + DataType::UInt8 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt8Array, u8), + DataType::UInt16 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt16Array, u16), + DataType::UInt32 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt32Array, u32), + DataType::UInt64 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt64Array, u64), + other => Err(DataFusionError::Internal(format!( + "Data type {:?} not supported for binary operation '{}' on dyn arrays", + other, + stringify!($KERNEL), + ))), + }; + Some(result) } - DataType::UInt32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u32| a | b, - UInt32Array, - u32 - ) - } - DataType::UInt64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u64| move |b: u64| a | b, - UInt64Array, - u64 - ) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseOr - ))), }; - Some(result) } -pub(crate) fn bitwise_xor_scalar( - array: &dyn Array, - scalar: ScalarValue, -) -> Option> { - let result = match array.data_type() { - DataType::Int8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i8| move |b: i8| a ^ b, - Int8Array, - i8 - ) - } - DataType::Int16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i16| move |b: i16| a ^ b, - Int16Array, - i16 - ) - } - DataType::Int32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i32| move |b: i32| a ^ b, - Int32Array, - i32 - ) - } - DataType::Int64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i64| move |b: i64| a ^ b, - Int64Array, - i64 - ) - } - DataType::UInt8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u8| move |b: u8| a ^ b, - UInt8Array, - u8 - ) - } - DataType::UInt16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u16| move |b: u16| a ^ b, - UInt16Array, - u16 - ) - } - DataType::UInt32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u32| a ^ b, - UInt32Array, - u32 - ) - } - DataType::UInt64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u64| move |b: u64| a ^ b, - UInt64Array, - u64 - ) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseXor - ))), - }; - Some(result) -} +create_dyn_scalar_kernel!(bitwise_and_dyn_scalar, bitwise_and_scalar); +create_dyn_scalar_kernel!(bitwise_or_dyn_scalar, bitwise_or_scalar); +create_dyn_scalar_kernel!(bitwise_xor_dyn_scalar, bitwise_xor_scalar); -pub(crate) fn bitwise_shift_right_scalar( +pub(crate) fn bitwise_shift_right_dyn_scalar( array: &dyn Array, scalar: ScalarValue, ) -> Option> { @@ -650,7 +399,7 @@ pub(crate) fn bitwise_shift_right_scalar( Some(result) } -pub(crate) fn bitwise_shift_left_scalar( +pub(crate) fn bitwise_shift_left_dyn_scalar( array: &dyn Array, scalar: ScalarValue, ) -> Option> { From dd03fb9482676c12e8d06ced8b7fce1bc4cbe414 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 24 Apr 2023 10:18:04 -0400 Subject: [PATCH 2/3] Apply suggestions from code review Co-authored-by: Liang-Chi Hsieh --- datafusion/physical-expr/src/expressions/binary/kernels.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index 4f3504ab8134..98b21f55af8a 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -268,7 +268,7 @@ pub(crate) fn bitwise_shift_left_dyn( } } -/// Downcasts $LEFT as $ARRAY_TYPE and $RIGHT as TYPE and calls $KERNEL($LEFT, RIGHT) +/// Downcasts $LEFT as $ARRAY_TYPE and $RIGHT as TYPE and calls $KERNEL($LEFT, $RIGHT) macro_rules! call_bitwise_scalar_kernel { ($LEFT:expr, $RIGHT:expr, $KERNEL:ident, $ARRAY_TYPE:ident, $TYPE:ty) => {{ let len = $LEFT.len(); @@ -284,7 +284,7 @@ macro_rules! call_bitwise_scalar_kernel { }}; } -/// Creates a $FUNC(left: ArrayRef, right: ArrayRef) that +/// Creates a $FUNC(left: ArrayRef, right: ScalarValue) that /// downcasts left / right to the appropriate integral type and calls the kernel macro_rules! create_dyn_scalar_kernel { ($FUNC:ident, $KERNEL:ident) => { From 4de167316cad9db05042a387b9fad270e68f040c Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Mon, 24 Apr 2023 10:36:31 -0400 Subject: [PATCH 3/3] Add link to upstream arrow ticket --- datafusion/physical-expr/src/expressions/binary/kernels.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index 98b21f55af8a..5d30c7f423aa 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -114,6 +114,7 @@ create_dyn_kernel!(bitwise_or_dyn, bitwise_or); create_dyn_kernel!(bitwise_xor_dyn, bitwise_xor); create_dyn_kernel!(bitwise_and_dyn, bitwise_and); +// TODO: use create_dyn_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented pub(crate) fn bitwise_shift_right_dyn( left: ArrayRef, right: ArrayRef, @@ -191,6 +192,7 @@ pub(crate) fn bitwise_shift_right_dyn( } } +// TODO: use create_dyn_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented pub(crate) fn bitwise_shift_left_dyn( left: ArrayRef, right: ArrayRef, @@ -313,6 +315,7 @@ create_dyn_scalar_kernel!(bitwise_and_dyn_scalar, bitwise_and_scalar); create_dyn_scalar_kernel!(bitwise_or_dyn_scalar, bitwise_or_scalar); create_dyn_scalar_kernel!(bitwise_xor_dyn_scalar, bitwise_xor_scalar); +// TODO: use create_dyn_scalar_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented pub(crate) fn bitwise_shift_right_dyn_scalar( array: &dyn Array, scalar: ScalarValue, @@ -399,6 +402,7 @@ pub(crate) fn bitwise_shift_right_dyn_scalar( Some(result) } +// TODO: use create_dyn_scalar_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented pub(crate) fn bitwise_shift_left_dyn_scalar( array: &dyn Array, scalar: ScalarValue,