/Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/expressions/binary/kernels.rs
Line | Count | Source (jump to first uncovered line) |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | //! This module contains computation kernels that are specific to |
19 | | //! datafusion and not (yet) targeted to port upstream to arrow |
20 | | use arrow::array::*; |
21 | | use arrow::compute::kernels::bitwise::{ |
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, |
25 | | }; |
26 | | use arrow::datatypes::DataType; |
27 | | use datafusion_common::plan_err; |
28 | | use datafusion_common::{Result, ScalarValue}; |
29 | | |
30 | | use arrow_schema::ArrowError; |
31 | | use std::sync::Arc; |
32 | | |
33 | | /// Downcasts $LEFT and $RIGHT to $ARRAY_TYPE and then calls $KERNEL($LEFT, $RIGHT) |
34 | | macro_rules! call_bitwise_kernel { |
35 | | ($LEFT:expr, $RIGHT:expr, $KERNEL:expr, $ARRAY_TYPE:ident) => {{ |
36 | | let left = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); |
37 | | let right = $RIGHT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); |
38 | | let result: $ARRAY_TYPE = $KERNEL(left, right)?; |
39 | | Ok(Arc::new(result)) |
40 | | }}; |
41 | | } |
42 | | |
43 | | /// Creates a $FUNC(left: ArrayRef, right: ArrayRef) that |
44 | | /// downcasts left / right to the appropriate integral type and calls the kernel |
45 | | macro_rules! create_dyn_kernel { |
46 | | ($FUNC:ident, $KERNEL:ident) => { |
47 | 0 | pub(crate) fn $FUNC(left: ArrayRef, right: ArrayRef) -> Result<ArrayRef> { |
48 | 0 | match &left.data_type() { |
49 | | DataType::Int8 => { |
50 | 0 | call_bitwise_kernel!(left, right, $KERNEL, Int8Array) |
51 | | } |
52 | | DataType::Int16 => { |
53 | 0 | call_bitwise_kernel!(left, right, $KERNEL, Int16Array) |
54 | | } |
55 | | DataType::Int32 => { |
56 | 0 | call_bitwise_kernel!(left, right, $KERNEL, Int32Array) |
57 | | } |
58 | | DataType::Int64 => { |
59 | 0 | call_bitwise_kernel!(left, right, $KERNEL, Int64Array) |
60 | | } |
61 | | DataType::UInt8 => { |
62 | 0 | call_bitwise_kernel!(left, right, $KERNEL, UInt8Array) |
63 | | } |
64 | | DataType::UInt16 => { |
65 | 0 | call_bitwise_kernel!(left, right, $KERNEL, UInt16Array) |
66 | | } |
67 | | DataType::UInt32 => { |
68 | 0 | call_bitwise_kernel!(left, right, $KERNEL, UInt32Array) |
69 | | } |
70 | | DataType::UInt64 => { |
71 | 0 | call_bitwise_kernel!(left, right, $KERNEL, UInt64Array) |
72 | | } |
73 | 0 | other => plan_err!( |
74 | 0 | "Data type {:?} not supported for binary operation '{}' on dyn arrays", |
75 | 0 | other, |
76 | 0 | stringify!($KERNEL) |
77 | 0 | ), |
78 | | } |
79 | 0 | } |
80 | | }; |
81 | | } |
82 | | |
83 | | create_dyn_kernel!(bitwise_or_dyn, bitwise_or); |
84 | | create_dyn_kernel!(bitwise_xor_dyn, bitwise_xor); |
85 | | create_dyn_kernel!(bitwise_and_dyn, bitwise_and); |
86 | | create_dyn_kernel!(bitwise_shift_right_dyn, bitwise_shift_right); |
87 | | create_dyn_kernel!(bitwise_shift_left_dyn, bitwise_shift_left); |
88 | | |
89 | | /// Downcasts $LEFT as $ARRAY_TYPE and $RIGHT as TYPE and calls $KERNEL($LEFT, $RIGHT) |
90 | | macro_rules! call_bitwise_scalar_kernel { |
91 | | ($LEFT:expr, $RIGHT:expr, $KERNEL:ident, $ARRAY_TYPE:ident, $TYPE:ty) => {{ |
92 | | let len = $LEFT.len(); |
93 | | let array = $LEFT.as_any().downcast_ref::<$ARRAY_TYPE>().unwrap(); |
94 | | let scalar = $RIGHT; |
95 | | if scalar.is_null() { |
96 | | Ok(new_null_array(array.data_type(), len)) |
97 | | } else { |
98 | | let scalar: $TYPE = scalar.try_into().unwrap(); |
99 | | let result: $ARRAY_TYPE = $KERNEL(array, scalar).unwrap(); |
100 | | Ok(Arc::new(result) as ArrayRef) |
101 | | } |
102 | | }}; |
103 | | } |
104 | | |
105 | | /// Creates a $FUNC(left: ArrayRef, right: ScalarValue) that |
106 | | /// downcasts left / right to the appropriate integral type and calls the kernel |
107 | | macro_rules! create_dyn_scalar_kernel { |
108 | | ($FUNC:ident, $KERNEL:ident) => { |
109 | 0 | pub(crate) fn $FUNC(array: &dyn Array, scalar: ScalarValue) -> Option<Result<ArrayRef>> { |
110 | 0 | let result = match array.data_type() { |
111 | 0 | DataType::Int8 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int8Array, i8), |
112 | 0 | DataType::Int16 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int16Array, i16), |
113 | 0 | DataType::Int32 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int32Array, i32), |
114 | 0 | DataType::Int64 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, Int64Array, i64), |
115 | 0 | DataType::UInt8 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt8Array, u8), |
116 | 0 | DataType::UInt16 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt16Array, u16), |
117 | 0 | DataType::UInt32 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt32Array, u32), |
118 | 0 | DataType::UInt64 => call_bitwise_scalar_kernel!(array, scalar, $KERNEL, UInt64Array, u64), |
119 | 0 | other => plan_err!( |
120 | 0 | "Data type {:?} not supported for binary operation '{}' on dyn arrays", |
121 | 0 | other, |
122 | 0 | stringify!($KERNEL) |
123 | 0 | ), |
124 | | }; |
125 | 0 | Some(result) |
126 | 0 | } |
127 | | }; |
128 | | } |
129 | | |
130 | | create_dyn_scalar_kernel!(bitwise_and_dyn_scalar, bitwise_and_scalar); |
131 | | create_dyn_scalar_kernel!(bitwise_or_dyn_scalar, bitwise_or_scalar); |
132 | | create_dyn_scalar_kernel!(bitwise_xor_dyn_scalar, bitwise_xor_scalar); |
133 | | create_dyn_scalar_kernel!(bitwise_shift_right_dyn_scalar, bitwise_shift_right_scalar); |
134 | | create_dyn_scalar_kernel!(bitwise_shift_left_dyn_scalar, bitwise_shift_left_scalar); |
135 | | |
136 | 0 | pub fn concat_elements_utf8view( |
137 | 0 | left: &StringViewArray, |
138 | 0 | right: &StringViewArray, |
139 | 0 | ) -> std::result::Result<StringViewArray, ArrowError> { |
140 | 0 | let capacity = left |
141 | 0 | .data_buffers() |
142 | 0 | .iter() |
143 | 0 | .zip(right.data_buffers().iter()) |
144 | 0 | .map(|(b1, b2)| b1.len() + b2.len()) |
145 | 0 | .sum(); |
146 | 0 | let mut result = StringViewBuilder::with_capacity(capacity); |
147 | 0 |
|
148 | 0 | // Avoid reallocations by writing to a reused buffer (note we |
149 | 0 | // could be even more efficient r by creating the view directly |
150 | 0 | // here and avoid the buffer but that would be more complex) |
151 | 0 | let mut buffer = String::new(); |
152 | | |
153 | 0 | for (left, right) in left.iter().zip(right.iter()) { |
154 | 0 | if let (Some(left), Some(right)) = (left, right) { |
155 | 0 | use std::fmt::Write; |
156 | 0 | buffer.clear(); |
157 | 0 | write!(&mut buffer, "{left}{right}") |
158 | 0 | .expect("writing into string buffer failed"); |
159 | 0 | result.append_value(&buffer); |
160 | 0 | } else { |
161 | | // at least one of the values is null, so the output is also null |
162 | 0 | result.append_null() |
163 | | } |
164 | | } |
165 | 0 | Ok(result.finish()) |
166 | 0 | } |