diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index ea551000df9a..6769032bff6d 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -49,10 +49,10 @@ use arrow::compute::kernels::comparison::{ }; use adapter::{eq_dyn, gt_dyn, gt_eq_dyn, lt_dyn, lt_eq_dyn, neq_dyn}; +use arrow::compute::kernels::concat_elements::concat_elements_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, - string_concat, }; use kernels_arrow::{ add_decimal, add_decimal_scalar, divide_decimal, divide_decimal_scalar, @@ -2508,6 +2508,18 @@ mod tests { Ok(()) } + #[test] + 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 expected = Int32Array::from(vec![Some(32)]); + assert_eq!(result.as_ref(), &expected); + + Ok(()) + } + #[test] fn bitwise_scalar_test() -> Result<()> { let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef; diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index f5d7b0dd93f3..3ca4a447c875 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -21,13 +21,14 @@ use arrow::array::*; 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, $OP:tt, $ARRAY_TYPE:ident, $TYPE:ty) => {{ + ($LEFT:expr, $RIGHT:expr, $METHOD:expr, $ARRAY_TYPE:ident) => {{ let len = $LEFT.len(); let left = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); let right = $RIGHT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); @@ -37,7 +38,7 @@ macro_rules! binary_bitwise_array_op { if left.is_null(i) || right.is_null(i) { None } else { - Some(left.value(i) $OP right.value(i)) + Some($METHOD(left.value(i), right.value(i))) } }) .collect::<$ARRAY_TYPE>(); @@ -49,7 +50,7 @@ macro_rules! binary_bitwise_array_op { /// 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, $OP:tt, $ARRAY_TYPE:ident, $TYPE:ty) => {{ + ($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; @@ -63,7 +64,7 @@ macro_rules! binary_bitwise_array_scalar { if array.is_null(i) { None } else { - Some(array.value(i) $OP right) + Some($METHOD(array.value(i), right)) } }) .collect::<$ARRAY_TYPE>(); @@ -75,16 +76,16 @@ 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, &, Int8Array, i8) + binary_bitwise_array_op!(left, right, |a, b| a & b, Int8Array) } DataType::Int16 => { - binary_bitwise_array_op!(left, right, &, Int16Array, i16) + binary_bitwise_array_op!(left, right, |a, b| a & b, Int16Array) } DataType::Int32 => { - binary_bitwise_array_op!(left, right, &, Int32Array, i32) + binary_bitwise_array_op!(left, right, |a, b| a & b, Int32Array) } DataType::Int64 => { - binary_bitwise_array_op!(left, right, &, Int64Array, i64) + binary_bitwise_array_op!(left, right, |a, b| a & b, Int64Array) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -97,16 +98,36 @@ pub(crate) fn bitwise_and(left: ArrayRef, right: ArrayRef) -> Result { pub(crate) fn bitwise_shift_right(left: ArrayRef, right: ArrayRef) -> Result { match &left.data_type() { DataType::Int8 => { - binary_bitwise_array_op!(left, right, >>, Int8Array, i8) + 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, >>, Int16Array, i16) + 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, >>, Int32Array, i32) + 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, >>, Int64Array, i64) + binary_bitwise_array_op!( + left, + right, + |a: i64, b: i64| a.wrapping_shr(b as u32), + Int64Array + ) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -119,16 +140,36 @@ pub(crate) fn bitwise_shift_right(left: ArrayRef, right: ArrayRef) -> Result Result { match &left.data_type() { DataType::Int8 => { - binary_bitwise_array_op!(left, right, <<, Int8Array, i8) + 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, <<, Int16Array, i16) + 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, <<, Int32Array, i32) + 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, <<, Int64Array, i64) + binary_bitwise_array_op!( + left, + right, + |a: i64, b: i64| a.wrapping_shl(b as u32), + Int64Array + ) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -141,16 +182,16 @@ pub(crate) fn bitwise_shift_left(left: ArrayRef, right: ArrayRef) -> Result Result { match &left.data_type() { DataType::Int8 => { - binary_bitwise_array_op!(left, right, |, Int8Array, i8) + binary_bitwise_array_op!(left, right, |a, b| a | b, Int8Array) } DataType::Int16 => { - binary_bitwise_array_op!(left, right, |, Int16Array, i16) + binary_bitwise_array_op!(left, right, |a, b| a | b, Int16Array) } DataType::Int32 => { - binary_bitwise_array_op!(left, right, |, Int32Array, i32) + binary_bitwise_array_op!(left, right, |a, b| a | b, Int32Array) } DataType::Int64 => { - binary_bitwise_array_op!(left, right, |, Int64Array, i64) + binary_bitwise_array_op!(left, right, |a, b| a | b, Int64Array) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -166,16 +207,16 @@ pub(crate) fn bitwise_and_scalar( ) -> Option> { let result = match array.data_type() { DataType::Int8 => { - binary_bitwise_array_scalar!(array, scalar, &, Int8Array, i8) + binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int8Array, i8) } DataType::Int16 => { - binary_bitwise_array_scalar!(array, scalar, &, Int16Array, i16) + binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int16Array, i16) } DataType::Int32 => { - binary_bitwise_array_scalar!(array, scalar, &, Int32Array, i32) + binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int32Array, i32) } DataType::Int64 => { - binary_bitwise_array_scalar!(array, scalar, &, Int64Array, i64) + binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int64Array, i64) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -192,16 +233,16 @@ pub(crate) fn bitwise_or_scalar( ) -> Option> { let result = match array.data_type() { DataType::Int8 => { - binary_bitwise_array_scalar!(array, scalar, |, Int8Array, i8) + binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int8Array, i8) } DataType::Int16 => { - binary_bitwise_array_scalar!(array, scalar, |, Int16Array, i16) + binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int16Array, i16) } DataType::Int32 => { - binary_bitwise_array_scalar!(array, scalar, |, Int32Array, i32) + binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int32Array, i32) } DataType::Int64 => { - binary_bitwise_array_scalar!(array, scalar, |, Int64Array, i64) + binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int64Array, i64) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -218,16 +259,40 @@ pub(crate) fn bitwise_shift_right_scalar( ) -> Option> { let result = match array.data_type() { DataType::Int8 => { - binary_bitwise_array_scalar!(array, scalar, >>, Int8Array, i8) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i8, b: i8| a.wrapping_shr(b as u32), + Int8Array, + i8 + ) } DataType::Int16 => { - binary_bitwise_array_scalar!(array, scalar, >>, Int16Array, i16) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i16, b: i16| a.wrapping_shr(b as u32), + Int16Array, + i16 + ) } DataType::Int32 => { - binary_bitwise_array_scalar!(array, scalar, >>, Int32Array, i32) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i32, b: i32| a.wrapping_shr(b as u32), + Int32Array, + i32 + ) } DataType::Int64 => { - binary_bitwise_array_scalar!(array, scalar, >>, Int64Array, i64) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i64, b: i64| a.wrapping_shr(b as u32), + Int64Array, + i64 + ) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays", @@ -244,16 +309,40 @@ pub(crate) fn bitwise_shift_left_scalar( ) -> Option> { let result = match array.data_type() { DataType::Int8 => { - binary_bitwise_array_scalar!(array, scalar, <<, Int8Array, i8) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i8, b: i8| a.wrapping_shl(b as u32), + Int8Array, + i8 + ) } DataType::Int16 => { - binary_bitwise_array_scalar!(array, scalar, <<, Int16Array, i16) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i16, b: i16| a.wrapping_shl(b as u32), + Int16Array, + i16 + ) } DataType::Int32 => { - binary_bitwise_array_scalar!(array, scalar, <<, Int32Array, i32) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i32, b: i32| a.wrapping_shl(b as u32), + Int32Array, + i32 + ) } DataType::Int64 => { - binary_bitwise_array_scalar!(array, scalar, <<, Int64Array, i64) + binary_bitwise_array_scalar!( + array, + scalar, + |a: i64, b: i64| a.wrapping_shl(b as u32), + Int64Array, + i64 + ) } other => Err(DataFusionError::Internal(format!( "Data type {:?} not supported for binary operation '{}' on dyn arrays",