Skip to content

Commit 17fb689

Browse files
ovrovr
committed
no panic on overflow
1 parent efe6cd7 commit 17fb689

2 files changed

Lines changed: 139 additions & 38 deletions

File tree

‎datafusion/physical-expr/src/expressions/binary.rs‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,10 +49,10 @@ use arrow::compute::kernels::comparison::{
4949
};
5050

5151
use adapter::{eq_dyn, gt_dyn, gt_eq_dyn, lt_dyn, lt_eq_dyn, neq_dyn};
52+
use arrow::compute::kernels::concat_elements::concat_elements_utf8;
5253
use kernels::{
5354
bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_shift_left,
5455
bitwise_shift_left_scalar, bitwise_shift_right, bitwise_shift_right_scalar,
55-
string_concat,
5656
};
5757
use kernels_arrow::{
5858
add_decimal, add_decimal_scalar, divide_decimal, divide_decimal_scalar,
@@ -2508,6 +2508,18 @@ mod tests {
25082508
Ok(())
25092509
}
25102510

2511+
#[test]
2512+
fn bitwise_shift_array_overflow_test() -> Result<()> {
2513+
let input = Arc::new(Int32Array::from(vec![Some(2)])) as ArrayRef;
2514+
let modules = Arc::new(Int32Array::from(vec![Some(100)])) as ArrayRef;
2515+
let result = bitwise_shift_left(input.clone(), modules.clone())?;
2516+
2517+
let expected = Int32Array::from(vec![Some(32)]);
2518+
assert_eq!(result.as_ref(), &expected);
2519+
2520+
Ok(())
2521+
}
2522+
25112523
#[test]
25122524
fn bitwise_scalar_test() -> Result<()> {
25132525
let left = Arc::new(Int32Array::from(vec![Some(12), None, Some(11)])) as ArrayRef;
@@ -2525,7 +2537,7 @@ mod tests {
25252537
#[test]
25262538
fn bitwise_shift_scalar_test() -> Result<()> {
25272539
let input = Arc::new(Int32Array::from(vec![Some(2), None, Some(4)])) as ArrayRef;
2528-
let module = ScalarValue::from(10i32);
2540+
let module = ScalarValue::from(10i64);
25292541
let mut result = bitwise_shift_left_scalar(&input, module.clone()).unwrap()?;
25302542

25312543
let expected = Int32Array::from(vec![Some(2048), None, Some(4096)]);

‎datafusion/physical-expr/src/expressions/binary/kernels.rs‎

Lines changed: 125 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,14 @@ use arrow::array::*;
2121
use arrow::datatypes::DataType;
2222
use datafusion_common::{DataFusionError, Result, ScalarValue};
2323
use datafusion_expr::Operator;
24+
2425
use std::sync::Arc;
2526

2627
/// The binary_bitwise_array_op macro only evaluates for integer types
2728
/// like int64, int32.
2829
/// It is used to do bitwise operation.
2930
macro_rules! binary_bitwise_array_op {
30-
($LEFT:expr, $RIGHT:expr, $OP:tt, $ARRAY_TYPE:ident, $TYPE:ty) => {{
31+
($LEFT:expr, $RIGHT:expr, $METHOD:expr, $ARRAY_TYPE:ident) => {{
3132
let len = $LEFT.len();
3233
let left = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
3334
let right = $RIGHT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
@@ -37,7 +38,7 @@ macro_rules! binary_bitwise_array_op {
3738
if left.is_null(i) || right.is_null(i) {
3839
None
3940
} else {
40-
Some(left.value(i) $OP right.value(i))
41+
Some($METHOD(left.value(i), right.value(i)))
4142
}
4243
})
4344
.collect::<$ARRAY_TYPE>();
@@ -49,7 +50,7 @@ macro_rules! binary_bitwise_array_op {
4950
/// like int64, int32.
5051
/// It is used to do bitwise operation on an array with a scalar.
5152
macro_rules! binary_bitwise_array_scalar {
52-
($LEFT:expr, $RIGHT:expr, $OP:tt, $ARRAY_TYPE:ident, $TYPE:ty) => {{
53+
($LEFT:expr, $RIGHT:expr, $METHOD:expr, $ARRAY_TYPE:ident, $TYPE:ty) => {{
5354
let len = $LEFT.len();
5455
let array = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
5556
let scalar = $RIGHT;
@@ -63,7 +64,7 @@ macro_rules! binary_bitwise_array_scalar {
6364
if array.is_null(i) {
6465
None
6566
} else {
66-
Some(array.value(i) $OP right)
67+
Some($METHOD(array.value(i), right))
6768
}
6869
})
6970
.collect::<$ARRAY_TYPE>();
@@ -75,16 +76,16 @@ macro_rules! binary_bitwise_array_scalar {
7576
pub(crate) fn bitwise_and(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
7677
match &left.data_type() {
7778
DataType::Int8 => {
78-
binary_bitwise_array_op!(left, right, &, Int8Array, i8)
79+
binary_bitwise_array_op!(left, right, |a, b| a & b, Int8Array)
7980
}
8081
DataType::Int16 => {
81-
binary_bitwise_array_op!(left, right, &, Int16Array, i16)
82+
binary_bitwise_array_op!(left, right, |a, b| a & b, Int16Array)
8283
}
8384
DataType::Int32 => {
84-
binary_bitwise_array_op!(left, right, &, Int32Array, i32)
85+
binary_bitwise_array_op!(left, right, |a, b| a & b, Int32Array)
8586
}
8687
DataType::Int64 => {
87-
binary_bitwise_array_op!(left, right, &, Int64Array, i64)
88+
binary_bitwise_array_op!(left, right, |a, b| a & b, Int64Array)
8889
}
8990
other => Err(DataFusionError::Internal(format!(
9091
"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> {
9798
pub(crate) fn bitwise_shift_right(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
9899
match &left.data_type() {
99100
DataType::Int8 => {
100-
binary_bitwise_array_op!(left, right, >>, Int8Array, i8)
101+
binary_bitwise_array_op!(
102+
left,
103+
right,
104+
|a: i8, b: i8| a.wrapping_shr(b as u32),
105+
Int8Array
106+
)
101107
}
102108
DataType::Int16 => {
103-
binary_bitwise_array_op!(left, right, >>, Int16Array, i16)
109+
binary_bitwise_array_op!(
110+
left,
111+
right,
112+
|a: i16, b: i16| a.wrapping_shr(b as u32),
113+
Int16Array
114+
)
104115
}
105116
DataType::Int32 => {
106-
binary_bitwise_array_op!(left, right, >>, Int32Array, i32)
117+
binary_bitwise_array_op!(
118+
left,
119+
right,
120+
|a: i32, b: i32| a.wrapping_shr(b as u32),
121+
Int32Array
122+
)
107123
}
108124
DataType::Int64 => {
109-
binary_bitwise_array_op!(left, right, >>, Int64Array, i64)
125+
binary_bitwise_array_op!(
126+
left,
127+
right,
128+
|a: i64, b: i64| a.wrapping_shr(b as u32),
129+
Int64Array
130+
)
110131
}
111132
other => Err(DataFusionError::Internal(format!(
112133
"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
119140
pub(crate) fn bitwise_shift_left(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
120141
match &left.data_type() {
121142
DataType::Int8 => {
122-
binary_bitwise_array_op!(left, right, <<, Int8Array, i8)
143+
binary_bitwise_array_op!(
144+
left,
145+
right,
146+
|a: i8, b: i8| a.wrapping_shl(b as u32),
147+
Int8Array
148+
)
123149
}
124150
DataType::Int16 => {
125-
binary_bitwise_array_op!(left, right, <<, Int16Array, i16)
151+
binary_bitwise_array_op!(
152+
left,
153+
right,
154+
|a: i16, b: i16| a.wrapping_shl(b as u32),
155+
Int16Array
156+
)
126157
}
127158
DataType::Int32 => {
128-
binary_bitwise_array_op!(left, right, <<, Int32Array, i32)
159+
binary_bitwise_array_op!(
160+
left,
161+
right,
162+
|a: i32, b: i32| a.wrapping_shl(b as u32),
163+
Int32Array
164+
)
129165
}
130166
DataType::Int64 => {
131-
binary_bitwise_array_op!(left, right, <<, Int64Array, i64)
167+
binary_bitwise_array_op!(
168+
left,
169+
right,
170+
|a: i64, b: i64| a.wrapping_shl(b as u32),
171+
Int64Array
172+
)
132173
}
133174
other => Err(DataFusionError::Internal(format!(
134175
"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
141182
pub(crate) fn bitwise_or(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> {
142183
match &left.data_type() {
143184
DataType::Int8 => {
144-
binary_bitwise_array_op!(left, right, |, Int8Array, i8)
185+
binary_bitwise_array_op!(left, right, |a, b| a | b, Int8Array)
145186
}
146187
DataType::Int16 => {
147-
binary_bitwise_array_op!(left, right, |, Int16Array, i16)
188+
binary_bitwise_array_op!(left, right, |a, b| a | b, Int16Array)
148189
}
149190
DataType::Int32 => {
150-
binary_bitwise_array_op!(left, right, |, Int32Array, i32)
191+
binary_bitwise_array_op!(left, right, |a, b| a | b, Int32Array)
151192
}
152193
DataType::Int64 => {
153-
binary_bitwise_array_op!(left, right, |, Int64Array, i64)
194+
binary_bitwise_array_op!(left, right, |a, b| a | b, Int64Array)
154195
}
155196
other => Err(DataFusionError::Internal(format!(
156197
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
@@ -166,16 +207,16 @@ pub(crate) fn bitwise_and_scalar(
166207
) -> Option<Result<ArrayRef>> {
167208
let result = match array.data_type() {
168209
DataType::Int8 => {
169-
binary_bitwise_array_scalar!(array, scalar, &, Int8Array, i8)
210+
binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int8Array, i8)
170211
}
171212
DataType::Int16 => {
172-
binary_bitwise_array_scalar!(array, scalar, &, Int16Array, i16)
213+
binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int16Array, i16)
173214
}
174215
DataType::Int32 => {
175-
binary_bitwise_array_scalar!(array, scalar, &, Int32Array, i32)
216+
binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int32Array, i32)
176217
}
177218
DataType::Int64 => {
178-
binary_bitwise_array_scalar!(array, scalar, &, Int64Array, i64)
219+
binary_bitwise_array_scalar!(array, scalar, |a, b| a & b, Int64Array, i64)
179220
}
180221
other => Err(DataFusionError::Internal(format!(
181222
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
@@ -192,16 +233,16 @@ pub(crate) fn bitwise_or_scalar(
192233
) -> Option<Result<ArrayRef>> {
193234
let result = match array.data_type() {
194235
DataType::Int8 => {
195-
binary_bitwise_array_scalar!(array, scalar, |, Int8Array, i8)
236+
binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int8Array, i8)
196237
}
197238
DataType::Int16 => {
198-
binary_bitwise_array_scalar!(array, scalar, |, Int16Array, i16)
239+
binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int16Array, i16)
199240
}
200241
DataType::Int32 => {
201-
binary_bitwise_array_scalar!(array, scalar, |, Int32Array, i32)
242+
binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int32Array, i32)
202243
}
203244
DataType::Int64 => {
204-
binary_bitwise_array_scalar!(array, scalar, |, Int64Array, i64)
245+
binary_bitwise_array_scalar!(array, scalar, |a, b| a | b, Int64Array, i64)
205246
}
206247
other => Err(DataFusionError::Internal(format!(
207248
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
@@ -218,16 +259,40 @@ pub(crate) fn bitwise_shift_right_scalar(
218259
) -> Option<Result<ArrayRef>> {
219260
let result = match array.data_type() {
220261
DataType::Int8 => {
221-
binary_bitwise_array_scalar!(array, scalar, >>, Int8Array, i8)
262+
binary_bitwise_array_scalar!(
263+
array,
264+
scalar,
265+
|a: i8, b: i8| a.wrapping_shr(b as u32),
266+
Int8Array,
267+
i8
268+
)
222269
}
223270
DataType::Int16 => {
224-
binary_bitwise_array_scalar!(array, scalar, >>, Int16Array, i16)
271+
binary_bitwise_array_scalar!(
272+
array,
273+
scalar,
274+
|a: i16, b: i16| a.wrapping_shr(b as u32),
275+
Int16Array,
276+
i16
277+
)
225278
}
226279
DataType::Int32 => {
227-
binary_bitwise_array_scalar!(array, scalar, >>, Int32Array, i32)
280+
binary_bitwise_array_scalar!(
281+
array,
282+
scalar,
283+
|a: i32, b: i32| a.wrapping_shr(b as u32),
284+
Int32Array,
285+
i32
286+
)
228287
}
229288
DataType::Int64 => {
230-
binary_bitwise_array_scalar!(array, scalar, >>, Int64Array, i64)
289+
binary_bitwise_array_scalar!(
290+
array,
291+
scalar,
292+
|a: i64, b: i64| a.wrapping_shr(b as u32),
293+
Int64Array,
294+
i64
295+
)
231296
}
232297
other => Err(DataFusionError::Internal(format!(
233298
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
@@ -244,16 +309,40 @@ pub(crate) fn bitwise_shift_left_scalar(
244309
) -> Option<Result<ArrayRef>> {
245310
let result = match array.data_type() {
246311
DataType::Int8 => {
247-
binary_bitwise_array_scalar!(array, scalar, <<, Int8Array, i8)
312+
binary_bitwise_array_scalar!(
313+
array,
314+
scalar,
315+
|a: i8, b: i8| a.wrapping_shl(b as u32),
316+
Int8Array,
317+
i8
318+
)
248319
}
249320
DataType::Int16 => {
250-
binary_bitwise_array_scalar!(array, scalar, <<, Int16Array, i16)
321+
binary_bitwise_array_scalar!(
322+
array,
323+
scalar,
324+
|a: i16, b: i16| a.wrapping_shl(b as u32),
325+
Int16Array,
326+
i16
327+
)
251328
}
252329
DataType::Int32 => {
253-
binary_bitwise_array_scalar!(array, scalar, <<, Int32Array, i32)
330+
binary_bitwise_array_scalar!(
331+
array,
332+
scalar,
333+
|a: i32, b: i32| a.wrapping_shl(b as u32),
334+
Int32Array,
335+
i32
336+
)
254337
}
255338
DataType::Int64 => {
256-
binary_bitwise_array_scalar!(array, scalar, <<, Int64Array, i64)
339+
binary_bitwise_array_scalar!(
340+
array,
341+
scalar,
342+
|a: i64, b: i64| a.wrapping_shl(b as u32),
343+
Int64Array,
344+
i64
345+
)
257346
}
258347
other => Err(DataFusionError::Internal(format!(
259348
"Data type {:?} not supported for binary operation '{}' on dyn arrays",

0 commit comments

Comments
 (0)