Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/math_expressions.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
//! Math expressions
19
20
use std::any::type_name;
21
use std::sync::Arc;
22
23
use arrow::array::ArrayRef;
24
use arrow::array::{BooleanArray, Float32Array, Float64Array};
25
use arrow::datatypes::DataType;
26
use arrow_array::Array;
27
28
use datafusion_common::exec_err;
29
use datafusion_common::{DataFusionError, Result};
30
31
macro_rules! downcast_arg {
32
    ($ARG:expr, $NAME:expr, $ARRAY_TYPE:ident) => {{
33
0
        $ARG.as_any().downcast_ref::<$ARRAY_TYPE>().ok_or_else(|| {
34
0
            DataFusionError::Internal(format!(
35
0
                "could not cast {} from {} to {}",
36
0
                $NAME,
37
0
                $ARG.data_type(),
38
0
                type_name::<$ARRAY_TYPE>()
39
0
            ))
40
0
        })?
41
    }};
42
}
43
44
macro_rules! make_function_scalar_inputs_return_type {
45
    ($ARG: expr, $NAME:expr, $ARGS_TYPE:ident, $RETURN_TYPE:ident, $FUNC: block) => {{
46
        let arg = downcast_arg!($ARG, $NAME, $ARGS_TYPE);
47
48
        arg.iter()
49
0
            .map(|a| match a {
50
0
                Some(a) => Some($FUNC(a)),
51
0
                _ => None,
52
0
            })
53
            .collect::<$RETURN_TYPE>()
54
    }};
55
}
56
57
/// Isnan SQL function
58
pub fn isnan(args: &[ArrayRef]) -> Result<ArrayRef> {
59
    match args[0].data_type() {
60
        DataType::Float64 => Ok(Arc::new(make_function_scalar_inputs_return_type!(
61
            &args[0],
62
            "x",
63
            Float64Array,
64
            BooleanArray,
65
            { f64::is_nan }
66
        )) as ArrayRef),
67
68
        DataType::Float32 => Ok(Arc::new(make_function_scalar_inputs_return_type!(
69
            &args[0],
70
            "x",
71
            Float32Array,
72
            BooleanArray,
73
            { f32::is_nan }
74
        )) as ArrayRef),
75
76
        other => exec_err!("Unsupported data type {other:?} for function isnan"),
77
    }
78
}
79
80
#[cfg(test)]
81
mod tests {
82
83
    use datafusion_common::cast::as_boolean_array;
84
85
    use super::*;
86
87
    #[test]
88
    fn test_isnan_f64() {
89
        let args: Vec<ArrayRef> = vec![Arc::new(Float64Array::from(vec![
90
            1.0,
91
            f64::NAN,
92
            3.0,
93
            -f64::NAN,
94
        ]))];
95
96
        let result = isnan(&args).expect("failed to initialize function isnan");
97
        let booleans =
98
            as_boolean_array(&result).expect("failed to initialize function isnan");
99
100
        assert_eq!(booleans.len(), 4);
101
        assert!(!booleans.value(0));
102
        assert!(booleans.value(1));
103
        assert!(!booleans.value(2));
104
        assert!(booleans.value(3));
105
    }
106
107
    #[test]
108
    fn test_isnan_f32() {
109
        let args: Vec<ArrayRef> = vec![Arc::new(Float32Array::from(vec![
110
            1.0,
111
            f32::NAN,
112
            3.0,
113
            f32::NAN,
114
        ]))];
115
116
        let result = isnan(&args).expect("failed to initialize function isnan");
117
        let booleans =
118
            as_boolean_array(&result).expect("failed to initialize function isnan");
119
120
        assert_eq!(booleans.len(), 4);
121
        assert!(!booleans.value(0));
122
        assert!(booleans.value(1));
123
        assert!(!booleans.value(2));
124
        assert!(booleans.value(3));
125
    }
126
}