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