Coverage Report

Created: 2024-10-13 08:39

/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
}