datafusion_functions/core/
nvl2.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 arrow::datatypes::{DataType, Field, FieldRef};
19use datafusion_common::{Result, internal_err, utils::take_function_args};
20use datafusion_expr::{
21    ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
22    ScalarUDFImpl, Signature, Volatility,
23    conditional_expressions::CaseBuilder,
24    simplify::{ExprSimplifyResult, SimplifyContext},
25    type_coercion::binary::type_union_coercion,
26};
27use datafusion_macros::user_doc;
28
29#[user_doc(
30    doc_section(label = "Conditional Functions"),
31    description = "Returns _expression2_ if _expression1_ is not NULL; otherwise it returns _expression3_.",
32    syntax_example = "nvl2(expression1, expression2, expression3)",
33    sql_example = r#"```sql
34> select nvl2(null, 'a', 'b');
35+--------------------------------+
36| nvl2(NULL,Utf8("a"),Utf8("b")) |
37+--------------------------------+
38| b                              |
39+--------------------------------+
40> select nvl2('data', 'a', 'b');
41+----------------------------------------+
42| nvl2(Utf8("data"),Utf8("a"),Utf8("b")) |
43+----------------------------------------+
44| a                                      |
45+----------------------------------------+
46```
47"#,
48    argument(
49        name = "expression1",
50        description = "Expression to test for null. Can be a constant, column, or function, and any combination of operators."
51    ),
52    argument(
53        name = "expression2",
54        description = "Expression to return if expr1 is not null. Can be a constant, column, or function, and any combination of operators."
55    ),
56    argument(
57        name = "expression3",
58        description = "Expression to return if expr1 is null. Can be a constant, column, or function, and any combination of operators."
59    )
60)]
61#[derive(Debug, PartialEq, Eq, Hash)]
62pub struct NVL2Func {
63    signature: Signature,
64}
65
66impl Default for NVL2Func {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl NVL2Func {
73    pub fn new() -> Self {
74        Self {
75            signature: Signature::user_defined(Volatility::Immutable),
76        }
77    }
78}
79
80impl ScalarUDFImpl for NVL2Func {
81    fn name(&self) -> &str {
82        "nvl2"
83    }
84
85    fn signature(&self) -> &Signature {
86        &self.signature
87    }
88
89    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
90        Ok(arg_types[1].clone())
91    }
92
93    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
94        let nullable =
95            args.arg_fields[1].is_nullable() || args.arg_fields[2].is_nullable();
96        let return_type = args.arg_fields[1].data_type().clone();
97        Ok(Field::new(self.name(), return_type, nullable).into())
98    }
99
100    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
101        internal_err!("nvl2 should have been simplified to case")
102    }
103
104    fn simplify(
105        &self,
106        args: Vec<Expr>,
107        _info: &SimplifyContext,
108    ) -> Result<ExprSimplifyResult> {
109        let [test, if_non_null, if_null] = take_function_args(self.name(), args)?;
110
111        let expr = CaseBuilder::new(
112            None,
113            vec![test.is_not_null()],
114            vec![if_non_null],
115            Some(Box::new(if_null)),
116        )
117        .end()?;
118
119        Ok(ExprSimplifyResult::Simplified(expr))
120    }
121
122    fn short_circuits(&self) -> bool {
123        true
124    }
125
126    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
127        let [tested, if_non_null, if_null] = take_function_args(self.name(), arg_types)?;
128        let new_type =
129            [if_non_null, if_null]
130                .iter()
131                .try_fold(tested.clone(), |acc, x| {
132                    // `type_union_coercion` finds a loose common type; the actual
133                    // coercion is done by `maybe_data_types`.
134                    let coerced_type = type_union_coercion(&acc, x);
135                    if let Some(coerced_type) = coerced_type {
136                        Ok(coerced_type)
137                    } else {
138                        internal_err!("Coercion from {acc} to {x} failed.")
139                    }
140                })?;
141        Ok(vec![new_type; arg_types.len()])
142    }
143
144    fn documentation(&self) -> Option<&Documentation> {
145        self.doc()
146    }
147}
datafusion_functions/core/nvl2.rs

datafusion_functions/core/
nvl2.rs
