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

datafusion_functions/core/
coalesce.rs
