datafusion_functions/math/
nanvl.rs

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
18use std::sync::Arc;
19
20use arrow::array::{ArrayRef, AsArray, Float16Array, Float32Array, Float64Array};
21use arrow::datatypes::DataType::{Float16, Float32, Float64};
22use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type};
23use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args};
24use datafusion_expr::TypeSignature::Exact;
25use datafusion_expr::{
26    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
27    Volatility,
28};
29use datafusion_macros::user_doc;
30
31#[user_doc(
32    doc_section(label = "Math Functions"),
33    description = r#"Returns the first argument if it's not _NaN_.
34Returns the second argument otherwise."#,
35    syntax_example = "nanvl(expression_x, expression_y)",
36    sql_example = r#"```sql
37> SELECT nanvl(0, 5);
38+------------+
39| nanvl(0,5) |
40+------------+
41| 0          |
42+------------+
43```"#,
44    argument(
45        name = "expression_x",
46        description = "Numeric expression to return if it's not _NaN_. Can be a constant, column, or function, and any combination of arithmetic operators."
47    ),
48    argument(
49        name = "expression_y",
50        description = "Numeric expression to return if the first expression is _NaN_. Can be a constant, column, or function, and any combination of arithmetic operators."
51    )
52)]
53#[derive(Debug, PartialEq, Eq, Hash)]
54pub struct NanvlFunc {
55    signature: Signature,
56}
57
58impl Default for NanvlFunc {
59    fn default() -> Self {
60        NanvlFunc::new()
61    }
62}
63
64impl NanvlFunc {
65    pub fn new() -> Self {
66        Self {
67            signature: Signature::one_of(
68                vec![
69                    Exact(vec![Float16, Float16]),
70                    Exact(vec![Float32, Float32]),
71                    Exact(vec![Float64, Float64]),
72                ],
73                Volatility::Immutable,
74            ),
75        }
76    }
77}
78
79impl ScalarUDFImpl for NanvlFunc {
80    fn name(&self) -> &str {
81        "nanvl"
82    }
83
84    fn signature(&self) -> &Signature {
85        &self.signature
86    }
87
88    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
89        match &arg_types[0] {
90            Float16 => Ok(Float16),
91            Float32 => Ok(Float32),
92            _ => Ok(Float64),
93        }
94    }
95
96    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
97        let [x, y] = take_function_args(self.name(), args.args)?;
98
99        match (x, y) {
100            (ColumnarValue::Scalar(ScalarValue::Float16(Some(v))), y) if v.is_nan() => {
101                Ok(y)
102            }
103            (ColumnarValue::Scalar(ScalarValue::Float32(Some(v))), y) if v.is_nan() => {
104                Ok(y)
105            }
106            (ColumnarValue::Scalar(ScalarValue::Float64(Some(v))), y) if v.is_nan() => {
107                Ok(y)
108            }
109            (x @ ColumnarValue::Scalar(_), _) => Ok(x),
110            (x, y) => {
111                let args = ColumnarValue::values_to_arrays(&[x, y])?;
112                Ok(ColumnarValue::Array(nanvl(&args)?))
113            }
114        }
115    }
116
117    fn documentation(&self) -> Option<&Documentation> {
118        self.doc()
119    }
120}
121
122/// Nanvl SQL function
123///
124/// - x is NaN -> output is y (which may itself be NULL)
125/// - otherwise -> output is x (which may itself be NULL)
126fn nanvl(args: &[ArrayRef]) -> Result<ArrayRef> {
127    match args[0].data_type() {
128        Float64 => {
129            let x = args[0].as_primitive::<Float64Type>();
130            let y = args[1].as_primitive::<Float64Type>();
131            let result: Float64Array = x
132                .iter()
133                .zip(y.iter())
134                .map(|(x_value, y_value)| match x_value {
135                    Some(x_value) if x_value.is_nan() => y_value,
136                    _ => x_value,
137                })
138                .collect();
139            Ok(Arc::new(result) as ArrayRef)
140        }
141        Float32 => {
142            let x = args[0].as_primitive::<Float32Type>();
143            let y = args[1].as_primitive::<Float32Type>();
144            let result: Float32Array = x
145                .iter()
146                .zip(y.iter())
147                .map(|(x_value, y_value)| match x_value {
148                    Some(x_value) if x_value.is_nan() => y_value,
149                    _ => x_value,
150                })
151                .collect();
152            Ok(Arc::new(result) as ArrayRef)
153        }
154        Float16 => {
155            let x = args[0].as_primitive::<Float16Type>();
156            let y = args[1].as_primitive::<Float16Type>();
157            let result: Float16Array = x
158                .iter()
159                .zip(y.iter())
160                .map(|(x_value, y_value)| match x_value {
161                    Some(x_value) if x_value.is_nan() => y_value,
162                    _ => x_value,
163                })
164                .collect();
165            Ok(Arc::new(result) as ArrayRef)
166        }
167        other => exec_err!("Unsupported data type {other:?} for function nanvl"),
168    }
169}
170
171#[cfg(test)]
172mod test {
173    use std::sync::Arc;
174
175    use crate::math::nanvl::nanvl;
176
177    use arrow::array::{ArrayRef, Float32Array, Float64Array};
178    use datafusion_common::cast::{as_float32_array, as_float64_array};
179
180    #[test]
181    fn test_nanvl_f64() {
182        let args: Vec<ArrayRef> = vec![
183            Arc::new(Float64Array::from(vec![1.0, f64::NAN, 3.0, f64::NAN])), // x
184            Arc::new(Float64Array::from(vec![5.0, 6.0, f64::NAN, f64::NAN])), // y
185        ];
186
187        let result = nanvl(&args).expect("failed to initialize function nanvl");
188        let floats =
189            as_float64_array(&result).expect("failed to initialize function nanvl");
190
191        assert_eq!(floats.len(), 4);
192        assert_eq!(floats.value(0), 1.0);
193        assert_eq!(floats.value(1), 6.0);
194        assert_eq!(floats.value(2), 3.0);
195        assert!(floats.value(3).is_nan());
196    }
197
198    #[test]
199    fn test_nanvl_f32() {
200        let args: Vec<ArrayRef> = vec![
201            Arc::new(Float32Array::from(vec![1.0, f32::NAN, 3.0, f32::NAN])), // x
202            Arc::new(Float32Array::from(vec![5.0, 6.0, f32::NAN, f32::NAN])), // y
203        ];
204
205        let result = nanvl(&args).expect("failed to initialize function nanvl");
206        let floats =
207            as_float32_array(&result).expect("failed to initialize function nanvl");
208
209        assert_eq!(floats.len(), 4);
210        assert_eq!(floats.value(0), 1.0);
211        assert_eq!(floats.value(1), 6.0);
212        assert_eq!(floats.value(2), 3.0);
213        assert!(floats.value(3).is_nan());
214    }
215}
datafusion_functions/math/nanvl.rs

datafusion_functions/math/
nanvl.rs
