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

datafusion_functions/math/
nanvl.rs
