datafusion_functions/core/
coalesce.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, exec_err, internal_err, plan_err};
20use datafusion_expr::binary::try_type_union_resolution;
21use datafusion_expr::conditional_expressions::CaseBuilder;
22use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext};
23use datafusion_expr::{
24    ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
25};
26use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
27use datafusion_macros::user_doc;
28use itertools::Itertools;
29
30#[user_doc(
31    doc_section(label = "Conditional Functions"),
32    description = "Returns the first of its arguments that is not _null_. Returns _null_ if all arguments are _null_. This function is often used to substitute a default value for _null_ values.",
33    syntax_example = "coalesce(expression1[, ..., expression_n])",
34    sql_example = r#"```sql
35> select coalesce(null, null, 'datafusion');
36+----------------------------------------+
37| coalesce(NULL,NULL,Utf8("datafusion")) |
38+----------------------------------------+
39| datafusion                             |
40+----------------------------------------+
41```"#,
42    argument(
43        name = "expression1, expression_n",
44        description = "Expression to use if previous expressions are _null_. Can be a constant, column, or function, and any combination of arithmetic operators. Pass as many expression arguments as necessary."
45    )
46)]
47#[derive(Debug, PartialEq, Eq, Hash)]
48pub struct CoalesceFunc {
49    pub(super) signature: Signature,
50}
51
52impl Default for CoalesceFunc {
53    fn default() -> Self {
54        CoalesceFunc::new()
55    }
56}
57
58impl CoalesceFunc {
59    pub fn new() -> Self {
60        Self {
61            signature: Signature::user_defined(Volatility::Immutable),
62        }
63    }
64}
65
66impl ScalarUDFImpl for CoalesceFunc {
67    fn name(&self) -> &str {
68        "coalesce"
69    }
70
71    fn signature(&self) -> &Signature {
72        &self.signature
73    }
74
75    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
76        internal_err!("return_field_from_args should be called instead")
77    }
78
79    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
80        // If any the arguments in coalesce is non-null, the result is non-null
81        let nullable = args.arg_fields.iter().all(|f| f.is_nullable());
82        let return_type = args
83            .arg_fields
84            .iter()
85            .map(|f| f.data_type())
86            .find_or_first(|d| !d.is_null())
87            .unwrap()
88            .clone();
89        Ok(Field::new(self.name(), return_type, nullable).into())
90    }
91
92    fn simplify(
93        &self,
94        args: Vec<Expr>,
95        _info: &SimplifyContext,
96    ) -> Result<ExprSimplifyResult> {
97        if args.is_empty() {
98            return plan_err!("coalesce must have at least one argument");
99        }
100        if args.len() == 1 {
101            return Ok(ExprSimplifyResult::Simplified(
102                args.into_iter().next().unwrap(),
103            ));
104        }
105
106        let n = args.len();
107        let (init, last_elem) = args.split_at(n - 1);
108        let whens = init
109            .iter()
110            .map(|x| x.clone().is_not_null())
111            .collect::<Vec<_>>();
112        let cases = init.to_vec();
113        Ok(ExprSimplifyResult::Simplified(
114            CaseBuilder::new(None, whens, cases, Some(Box::new(last_elem[0].clone())))
115                .end()?,
116        ))
117    }
118
119    /// coalesce evaluates to the first value which is not NULL
120    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
121        internal_err!("coalesce should have been simplified to case")
122    }
123
124    fn conditional_arguments<'a>(
125        &self,
126        args: &'a [Expr],
127    ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> {
128        let eager = vec![&args[0]];
129        let lazy = args[1..].iter().collect();
130        Some((eager, lazy))
131    }
132
133    fn short_circuits(&self) -> bool {
134        true
135    }
136
137    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
138        if arg_types.is_empty() {
139            return exec_err!("coalesce must have at least one argument");
140        }
141
142        try_type_union_resolution(arg_types)
143    }
144
145    fn documentation(&self) -> Option<&Documentation> {
146        self.doc()
147    }
148}
datafusion_functions/core/coalesce.rs

datafusion_functions/core/
coalesce.rs
