Skip to content

Commit 1d3860d

Browse files
izveigorizveigor
authored
refactor: bitwise kernel right and left shifts (#6585)
1 parent 919ac96 commit 1d3860d

File tree

1 file changed

+7
-365
lines changed
  • datafusion/physical-expr/src/expressions/binary

1 file changed

+7
-365
lines changed

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

Lines changed: 7 additions & 365 deletions
Original file line numberDiff line numberDiff line change
@@ -18,48 +18,16 @@
1818
//! This module contains computation kernels that are specific to
1919
//! datafusion and not (yet) targeted to port upstream to arrow
2020
use arrow::array::*;
21-
use arrow::compute::binary;
2221
use arrow::compute::kernels::bitwise::{
23-
bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_xor,
24-
bitwise_xor_scalar,
22+
bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_shift_left,
23+
bitwise_shift_left_scalar, bitwise_shift_right, bitwise_shift_right_scalar,
24+
bitwise_xor, bitwise_xor_scalar,
2525
};
2626
use arrow::datatypes::DataType;
2727
use datafusion_common::{DataFusionError, Result, ScalarValue};
28-
use datafusion_expr::Operator;
2928

3029
use std::sync::Arc;
3130

32-
/// The binary_bitwise_array_op macro only evaluates for integer types
33-
/// like int64, int32.
34-
/// It is used to do bitwise operation.
35-
macro_rules! binary_bitwise_array_op {
36-
($LEFT:expr, $RIGHT:expr, $METHOD:expr, $ARRAY_TYPE:ident) => {{
37-
let left = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
38-
let right = $RIGHT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
39-
let result: $ARRAY_TYPE = binary(left, right, $METHOD)?;
40-
Ok(Arc::new(result))
41-
}};
42-
}
43-
44-
/// The binary_bitwise_array_op macro only evaluates for integer types
45-
/// like int64, int32.
46-
/// It is used to do bitwise operation on an array with a scalar.
47-
macro_rules! binary_bitwise_array_scalar {
48-
($LEFT:expr, $RIGHT:expr, $METHOD:expr, $ARRAY_TYPE:ident, $TYPE:ty) => {{
49-
let len = $LEFT.len();
50-
let array = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap();
51-
let scalar = $RIGHT;
52-
if scalar.is_null() {
53-
Ok(new_null_array(array.data_type(), len))
54-
} else {
55-
let right: $TYPE = scalar.try_into().unwrap();
56-
let method = $METHOD(right);
57-
let result: $ARRAY_TYPE = array.unary(method);
58-
Ok(Arc::new(result) as ArrayRef)
59-
}
60-
}};
61-
}
62-
6331
/// Downcasts $LEFT and $RIGHT to $ARRAY_TYPE and then calls $KERNEL($LEFT, $RIGHT)
6432
macro_rules! call_bitwise_kernel {
6533
($LEFT:expr, $RIGHT:expr, $KERNEL:expr, $ARRAY_TYPE:ident) => {{
@@ -113,162 +81,8 @@ macro_rules! create_dyn_kernel {
11381
create_dyn_kernel!(bitwise_or_dyn, bitwise_or);
11482
create_dyn_kernel!(bitwise_xor_dyn, bitwise_xor);
11583
create_dyn_kernel!(bitwise_and_dyn, bitwise_and);
116-
117-
// TODO: use create_dyn_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented
118-
pub(crate) fn bitwise_shift_right_dyn(
119-
left: ArrayRef,
120-
right: ArrayRef,
121-
) -> Result<ArrayRef> {
122-
match &left.data_type() {
123-
DataType::Int8 => {
124-
binary_bitwise_array_op!(
125-
left,
126-
right,
127-
|a: i8, b: i8| a.wrapping_shr(b as u32),
128-
Int8Array
129-
)
130-
}
131-
DataType::Int16 => {
132-
binary_bitwise_array_op!(
133-
left,
134-
right,
135-
|a: i16, b: i16| a.wrapping_shr(b as u32),
136-
Int16Array
137-
)
138-
}
139-
DataType::Int32 => {
140-
binary_bitwise_array_op!(
141-
left,
142-
right,
143-
|a: i32, b: i32| a.wrapping_shr(b as u32),
144-
Int32Array
145-
)
146-
}
147-
DataType::Int64 => {
148-
binary_bitwise_array_op!(
149-
left,
150-
right,
151-
|a: i64, b: i64| a.wrapping_shr(b as u32),
152-
Int64Array
153-
)
154-
}
155-
DataType::UInt8 => {
156-
binary_bitwise_array_op!(
157-
left,
158-
right,
159-
|a: u8, b: u8| a.wrapping_shr(b as u32),
160-
UInt8Array
161-
)
162-
}
163-
DataType::UInt16 => {
164-
binary_bitwise_array_op!(
165-
left,
166-
right,
167-
|a: u16, b: u16| a.wrapping_shr(b as u32),
168-
UInt16Array
169-
)
170-
}
171-
DataType::UInt32 => {
172-
binary_bitwise_array_op!(
173-
left,
174-
right,
175-
|a: u32, b: u32| a.wrapping_shr(b),
176-
UInt32Array
177-
)
178-
}
179-
DataType::UInt64 => {
180-
binary_bitwise_array_op!(
181-
left,
182-
right,
183-
|a: u64, b: u64| a.wrapping_shr(b.try_into().unwrap()),
184-
UInt64Array
185-
)
186-
}
187-
other => Err(DataFusionError::Internal(format!(
188-
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
189-
other,
190-
Operator::BitwiseShiftRight
191-
))),
192-
}
193-
}
194-
195-
// TODO: use create_dyn_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented
196-
pub(crate) fn bitwise_shift_left_dyn(
197-
left: ArrayRef,
198-
right: ArrayRef,
199-
) -> Result<ArrayRef> {
200-
match &left.data_type() {
201-
DataType::Int8 => {
202-
binary_bitwise_array_op!(
203-
left,
204-
right,
205-
|a: i8, b: i8| a.wrapping_shl(b as u32),
206-
Int8Array
207-
)
208-
}
209-
DataType::Int16 => {
210-
binary_bitwise_array_op!(
211-
left,
212-
right,
213-
|a: i16, b: i16| a.wrapping_shl(b as u32),
214-
Int16Array
215-
)
216-
}
217-
DataType::Int32 => {
218-
binary_bitwise_array_op!(
219-
left,
220-
right,
221-
|a: i32, b: i32| a.wrapping_shl(b as u32),
222-
Int32Array
223-
)
224-
}
225-
DataType::Int64 => {
226-
binary_bitwise_array_op!(
227-
left,
228-
right,
229-
|a: i64, b: i64| a.wrapping_shl(b as u32),
230-
Int64Array
231-
)
232-
}
233-
DataType::UInt8 => {
234-
binary_bitwise_array_op!(
235-
left,
236-
right,
237-
|a: u8, b: u8| a.wrapping_shl(b as u32),
238-
UInt8Array
239-
)
240-
}
241-
DataType::UInt16 => {
242-
binary_bitwise_array_op!(
243-
left,
244-
right,
245-
|a: u16, b: u16| a.wrapping_shl(b as u32),
246-
UInt16Array
247-
)
248-
}
249-
DataType::UInt32 => {
250-
binary_bitwise_array_op!(
251-
left,
252-
right,
253-
|a: u32, b: u32| a.wrapping_shl(b),
254-
UInt32Array
255-
)
256-
}
257-
DataType::UInt64 => {
258-
binary_bitwise_array_op!(
259-
left,
260-
right,
261-
|a: u64, b: u64| a.wrapping_shr(b.try_into().unwrap()),
262-
UInt64Array
263-
)
264-
}
265-
other => Err(DataFusionError::Internal(format!(
266-
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
267-
other,
268-
Operator::BitwiseShiftLeft
269-
))),
270-
}
271-
}
84+
create_dyn_kernel!(bitwise_shift_right_dyn, bitwise_shift_right);
85+
create_dyn_kernel!(bitwise_shift_left_dyn, bitwise_shift_left);
27286

27387
/// Downcasts $LEFT as $ARRAY_TYPE and $RIGHT as TYPE and calls $KERNEL($LEFT, $RIGHT)
27488
macro_rules! call_bitwise_scalar_kernel {
@@ -314,177 +128,5 @@ macro_rules! create_dyn_scalar_kernel {
314128
create_dyn_scalar_kernel!(bitwise_and_dyn_scalar, bitwise_and_scalar);
315129
create_dyn_scalar_kernel!(bitwise_or_dyn_scalar, bitwise_or_scalar);
316130
create_dyn_scalar_kernel!(bitwise_xor_dyn_scalar, bitwise_xor_scalar);
317-
318-
// TODO: use create_dyn_scalar_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented
319-
pub(crate) fn bitwise_shift_right_dyn_scalar(
320-
array: &dyn Array,
321-
scalar: ScalarValue,
322-
) -> Option<Result<ArrayRef>> {
323-
let result = match array.data_type() {
324-
DataType::Int8 => {
325-
binary_bitwise_array_scalar!(
326-
array,
327-
scalar,
328-
|a: i8| move |b: i8| b.wrapping_shr(a as u32),
329-
Int8Array,
330-
i8
331-
)
332-
}
333-
DataType::Int16 => {
334-
binary_bitwise_array_scalar!(
335-
array,
336-
scalar,
337-
|a: i16| move |b: i16| b.wrapping_shr(a as u32),
338-
Int16Array,
339-
i16
340-
)
341-
}
342-
DataType::Int32 => {
343-
binary_bitwise_array_scalar!(
344-
array,
345-
scalar,
346-
|a: i32| move |b: i32| b.wrapping_shr(a as u32),
347-
Int32Array,
348-
i32
349-
)
350-
}
351-
DataType::Int64 => {
352-
binary_bitwise_array_scalar!(
353-
array,
354-
scalar,
355-
|a: i64| move |b: i64| b.wrapping_shr(a as u32),
356-
Int64Array,
357-
i64
358-
)
359-
}
360-
DataType::UInt8 => {
361-
binary_bitwise_array_scalar!(
362-
array,
363-
scalar,
364-
|a: u8| move |b: u8| b.wrapping_shr(a as u32),
365-
UInt8Array,
366-
u8
367-
)
368-
}
369-
DataType::UInt16 => {
370-
binary_bitwise_array_scalar!(
371-
array,
372-
scalar,
373-
|a: u16| move |b: u16| b.wrapping_shr(a as u32),
374-
UInt16Array,
375-
u16
376-
)
377-
}
378-
DataType::UInt32 => {
379-
binary_bitwise_array_scalar!(
380-
array,
381-
scalar,
382-
|a: u32| move |b: u32| b.wrapping_shr(a),
383-
UInt32Array,
384-
u32
385-
)
386-
}
387-
DataType::UInt64 => {
388-
binary_bitwise_array_scalar!(
389-
array,
390-
scalar,
391-
|a: u32| move |b: u64| b.wrapping_shr(a),
392-
UInt64Array,
393-
u32
394-
)
395-
}
396-
other => Err(DataFusionError::Internal(format!(
397-
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
398-
other,
399-
Operator::BitwiseShiftRight
400-
))),
401-
};
402-
Some(result)
403-
}
404-
405-
// TODO: use create_dyn_scalar_kernel! when https://github.com/apache/arrow-rs/issues/2741 is implemented
406-
pub(crate) fn bitwise_shift_left_dyn_scalar(
407-
array: &dyn Array,
408-
scalar: ScalarValue,
409-
) -> Option<Result<ArrayRef>> {
410-
let result = match array.data_type() {
411-
DataType::Int8 => {
412-
binary_bitwise_array_scalar!(
413-
array,
414-
scalar,
415-
|a: i8| move |b: i8| b.wrapping_shl(a as u32),
416-
Int8Array,
417-
i8
418-
)
419-
}
420-
DataType::Int16 => {
421-
binary_bitwise_array_scalar!(
422-
array,
423-
scalar,
424-
|a: i16| move |b: i16| b.wrapping_shl(a as u32),
425-
Int16Array,
426-
i16
427-
)
428-
}
429-
DataType::Int32 => {
430-
binary_bitwise_array_scalar!(
431-
array,
432-
scalar,
433-
|a: i32| move |b: i32| b.wrapping_shl(a as u32),
434-
Int32Array,
435-
i32
436-
)
437-
}
438-
DataType::Int64 => {
439-
binary_bitwise_array_scalar!(
440-
array,
441-
scalar,
442-
|a: i64| move |b: i64| b.wrapping_shl(a as u32),
443-
Int64Array,
444-
i64
445-
)
446-
}
447-
DataType::UInt8 => {
448-
binary_bitwise_array_scalar!(
449-
array,
450-
scalar,
451-
|a: u8| move |b: u8| b.wrapping_shl(a as u32),
452-
UInt8Array,
453-
u8
454-
)
455-
}
456-
DataType::UInt16 => {
457-
binary_bitwise_array_scalar!(
458-
array,
459-
scalar,
460-
|a: u16| move |b: u16| b.wrapping_shl(a as u32),
461-
UInt16Array,
462-
u16
463-
)
464-
}
465-
DataType::UInt32 => {
466-
binary_bitwise_array_scalar!(
467-
array,
468-
scalar,
469-
|a: u32| move |b: u32| b.wrapping_shl(a),
470-
UInt32Array,
471-
u32
472-
)
473-
}
474-
DataType::UInt64 => {
475-
binary_bitwise_array_scalar!(
476-
array,
477-
scalar,
478-
|a: u32| move |b: u64| b.wrapping_shr(a),
479-
UInt64Array,
480-
u32
481-
)
482-
}
483-
other => Err(DataFusionError::Internal(format!(
484-
"Data type {:?} not supported for binary operation '{}' on dyn arrays",
485-
other,
486-
Operator::BitwiseShiftLeft
487-
))),
488-
};
489-
Some(result)
490-
}
131+
create_dyn_scalar_kernel!(bitwise_shift_right_dyn_scalar, bitwise_shift_right_scalar);
132+
create_dyn_scalar_kernel!(bitwise_shift_left_dyn_scalar, bitwise_shift_left_scalar);

0 commit comments

Comments
 (0)