From 09ea7671279cd0bab3b0573b9073a08bf32c7537 Mon Sep 17 00:00:00 2001 From: izveigor Date: Wed, 7 Jun 2023 16:20:30 +0300 Subject: [PATCH] refactor: bitwise kernel right and left shifts --- .../src/expressions/binary/kernels.rs | 372 +----------------- 1 file changed, 7 insertions(+), 365 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index 5d30c7f423aa..185b9416d3d0 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -18,48 +18,16 @@ //! This module contains computation kernels that are specific to //! 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, + 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, }; use arrow::datatypes::DataType; use datafusion_common::{DataFusionError, Result, ScalarValue}; -use datafusion_expr::Operator; use std::sync::Arc; -/// The binary_bitwise_array_op macro only evaluates for integer types -/// like int64, int32. -/// It is used to do bitwise operation. -macro_rules! binary_bitwise_array_op { - ($LEFT:expr, $RIGHT:expr, $METHOD: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 = binary(left, right, $METHOD)?; - Ok(Arc::new(result)) - }}; -} - -/// The binary_bitwise_array_op macro only evaluates for integer types -/// like int64, int32. -/// It is used to do bitwise operation on an array with a scalar. -macro_rules! binary_bitwise_array_scalar { - ($LEFT:expr, $RIGHT:expr, $METHOD:expr, $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 right: $TYPE = scalar.try_into().unwrap(); - let method = $METHOD(right); - let result: $ARRAY_TYPE = array.unary(method); - Ok(Arc::new(result) as ArrayRef) - } - }}; -} - /// 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) => {{ @@ -113,162 +81,8 @@ macro_rules! create_dyn_kernel { 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, -) -> Result { - match &left.data_type() { - DataType::Int8 => { - binary_bitwise_array_op!( - left, - right, - |a: i8, b: i8| a.wrapping_shr(b as u32), - Int8Array - ) - } - DataType::Int16 => { - binary_bitwise_array_op!( - left, - right, - |a: i16, b: i16| a.wrapping_shr(b as u32), - Int16Array - ) - } - DataType::Int32 => { - binary_bitwise_array_op!( - left, - right, - |a: i32, b: i32| a.wrapping_shr(b as u32), - Int32Array - ) - } - DataType::Int64 => { - binary_bitwise_array_op!( - left, - right, - |a: i64, b: i64| a.wrapping_shr(b as u32), - Int64Array - ) - } - DataType::UInt8 => { - binary_bitwise_array_op!( - left, - right, - |a: u8, b: u8| a.wrapping_shr(b as u32), - UInt8Array - ) - } - DataType::UInt16 => { - binary_bitwise_array_op!( - left, - right, - |a: u16, b: u16| a.wrapping_shr(b as u32), - UInt16Array - ) - } - DataType::UInt32 => { - binary_bitwise_array_op!( - left, - right, - |a: u32, b: u32| a.wrapping_shr(b), - UInt32Array - ) - } - DataType::UInt64 => { - binary_bitwise_array_op!( - left, - right, - |a: u64, b: u64| a.wrapping_shr(b.try_into().unwrap()), - UInt64Array - ) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseShiftRight - ))), - } -} - -// 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, -) -> Result { - match &left.data_type() { - DataType::Int8 => { - binary_bitwise_array_op!( - left, - right, - |a: i8, b: i8| a.wrapping_shl(b as u32), - Int8Array - ) - } - DataType::Int16 => { - binary_bitwise_array_op!( - left, - right, - |a: i16, b: i16| a.wrapping_shl(b as u32), - Int16Array - ) - } - DataType::Int32 => { - binary_bitwise_array_op!( - left, - right, - |a: i32, b: i32| a.wrapping_shl(b as u32), - Int32Array - ) - } - DataType::Int64 => { - binary_bitwise_array_op!( - left, - right, - |a: i64, b: i64| a.wrapping_shl(b as u32), - Int64Array - ) - } - DataType::UInt8 => { - binary_bitwise_array_op!( - left, - right, - |a: u8, b: u8| a.wrapping_shl(b as u32), - UInt8Array - ) - } - DataType::UInt16 => { - binary_bitwise_array_op!( - left, - right, - |a: u16, b: u16| a.wrapping_shl(b as u32), - UInt16Array - ) - } - DataType::UInt32 => { - binary_bitwise_array_op!( - left, - right, - |a: u32, b: u32| a.wrapping_shl(b), - UInt32Array - ) - } - DataType::UInt64 => { - binary_bitwise_array_op!( - left, - right, - |a: u64, b: u64| a.wrapping_shr(b.try_into().unwrap()), - UInt64Array - ) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseShiftLeft - ))), - } -} +create_dyn_kernel!(bitwise_shift_right_dyn, bitwise_shift_right); +create_dyn_kernel!(bitwise_shift_left_dyn, bitwise_shift_left); /// Downcasts $LEFT as $ARRAY_TYPE and $RIGHT as TYPE and calls $KERNEL($LEFT, $RIGHT) macro_rules! call_bitwise_scalar_kernel { @@ -314,177 +128,5 @@ macro_rules! create_dyn_scalar_kernel { 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, -) -> Option> { - let result = match array.data_type() { - DataType::Int8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i8| move |b: i8| b.wrapping_shr(a as u32), - Int8Array, - i8 - ) - } - DataType::Int16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i16| move |b: i16| b.wrapping_shr(a as u32), - Int16Array, - i16 - ) - } - DataType::Int32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i32| move |b: i32| b.wrapping_shr(a as u32), - Int32Array, - i32 - ) - } - DataType::Int64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i64| move |b: i64| b.wrapping_shr(a as u32), - Int64Array, - i64 - ) - } - DataType::UInt8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u8| move |b: u8| b.wrapping_shr(a as u32), - UInt8Array, - u8 - ) - } - DataType::UInt16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u16| move |b: u16| b.wrapping_shr(a as u32), - UInt16Array, - u16 - ) - } - DataType::UInt32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u32| b.wrapping_shr(a), - UInt32Array, - u32 - ) - } - DataType::UInt64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u64| b.wrapping_shr(a), - UInt64Array, - u32 - ) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseShiftRight - ))), - }; - 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, -) -> Option> { - let result = match array.data_type() { - DataType::Int8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i8| move |b: i8| b.wrapping_shl(a as u32), - Int8Array, - i8 - ) - } - DataType::Int16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i16| move |b: i16| b.wrapping_shl(a as u32), - Int16Array, - i16 - ) - } - DataType::Int32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i32| move |b: i32| b.wrapping_shl(a as u32), - Int32Array, - i32 - ) - } - DataType::Int64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: i64| move |b: i64| b.wrapping_shl(a as u32), - Int64Array, - i64 - ) - } - DataType::UInt8 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u8| move |b: u8| b.wrapping_shl(a as u32), - UInt8Array, - u8 - ) - } - DataType::UInt16 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u16| move |b: u16| b.wrapping_shl(a as u32), - UInt16Array, - u16 - ) - } - DataType::UInt32 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u32| b.wrapping_shl(a), - UInt32Array, - u32 - ) - } - DataType::UInt64 => { - binary_bitwise_array_scalar!( - array, - scalar, - |a: u32| move |b: u64| b.wrapping_shr(a), - UInt64Array, - u32 - ) - } - other => Err(DataFusionError::Internal(format!( - "Data type {:?} not supported for binary operation '{}' on dyn arrays", - other, - Operator::BitwiseShiftLeft - ))), - }; - Some(result) -} +create_dyn_scalar_kernel!(bitwise_shift_right_dyn_scalar, bitwise_shift_right_scalar); +create_dyn_scalar_kernel!(bitwise_shift_left_dyn_scalar, bitwise_shift_left_scalar);