no panic on overflow

apache · Aug 11, 2022 · 92fa53c · 92fa53c
1 parent efe6cd7
commit 92fa53c
Show file tree

Hide file tree

Showing 2 changed files with 138 additions and 37 deletions.
diff --git 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
@@ -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<ArrayRef> {
     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<ArrayRef> {
 pub(crate) fn bitwise_shift_right(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
     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<Arr
 pub(crate) fn bitwise_shift_left(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
     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<Arra
 pub(crate) fn bitwise_or(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
     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<Result<ArrayRef>> {
     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<Result<ArrayRef>> {
     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<Result<ArrayRef>> {
     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<Result<ArrayRef>> {
     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",