datafusion_functions/core/
greatest.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 crate::core::greatest_least_utils::GreatestLeastOperator;
19use arrow::array::{Array, BooleanArray, make_comparator};
20use arrow::buffer::BooleanBuffer;
21use arrow::compute::SortOptions;
22use arrow::compute::kernels::cmp;
23use arrow::datatypes::DataType;
24use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err};
25use datafusion_doc::Documentation;
26use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
27use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
28use datafusion_macros::user_doc;
29use std::any::Any;
30
31const SORT_OPTIONS: SortOptions = SortOptions {
32    // We want greatest first
33    descending: false,
34
35    // NULL will be less than any other value
36    nulls_first: true,
37};
38
39#[user_doc(
40    doc_section(label = "Conditional Functions"),
41    description = "Returns the greatest value in a list of expressions. Returns _null_ if all expressions are _null_.",
42    syntax_example = "greatest(expression1[, ..., expression_n])",
43    sql_example = r#"```sql
44> select greatest(4, 7, 5);
45+---------------------------+
46| greatest(4,7,5)           |
47+---------------------------+
48| 7                         |
49+---------------------------+
50```"#,
51    argument(
52        name = "expression1, expression_n",
53        description = "Expressions to compare and return the greatest value.. Can be a constant, column, or function, and any combination of arithmetic operators. Pass as many expression arguments as necessary."
54    )
55)]
56#[derive(Debug, PartialEq, Eq, Hash)]
57pub struct GreatestFunc {
58    signature: Signature,
59}
60
61impl Default for GreatestFunc {
62    fn default() -> Self {
63        GreatestFunc::new()
64    }
65}
66
67impl GreatestFunc {
68    pub fn new() -> Self {
69        Self {
70            signature: Signature::user_defined(Volatility::Immutable),
71        }
72    }
73}
74
75impl GreatestLeastOperator for GreatestFunc {
76    const NAME: &'static str = "greatest";
77
78    fn keep_scalar<'a>(
79        lhs: &'a ScalarValue,
80        rhs: &'a ScalarValue,
81    ) -> Result<&'a ScalarValue> {
82        if !lhs.data_type().is_nested() {
83            return if lhs >= rhs { Ok(lhs) } else { Ok(rhs) };
84        }
85
86        // If complex type we can't compare directly as we want null values to be smaller
87        let cmp = make_comparator(
88            lhs.to_array()?.as_ref(),
89            rhs.to_array()?.as_ref(),
90            SORT_OPTIONS,
91        )?;
92
93        if cmp(0, 0).is_ge() { Ok(lhs) } else { Ok(rhs) }
94    }
95
96    /// Return boolean array where `arr[i] = lhs[i] >= rhs[i]` for all i, where `arr` is the result array
97    /// Nulls are always considered smaller than any other value
98    fn get_indexes_to_keep(lhs: &dyn Array, rhs: &dyn Array) -> Result<BooleanArray> {
99        // Fast path:
100        // If both arrays are not nested, have the same length and no nulls, we can use the faster vectorized kernel
101        // - If both arrays are not nested: Nested types, such as lists, are not supported as the null semantics are not well-defined.
102        // - both array does not have any nulls: cmp::gt_eq will return null if any of the input is null while we want to return false in that case
103        if !lhs.data_type().is_nested()
104            && lhs.logical_null_count() == 0
105            && rhs.logical_null_count() == 0
106        {
107            return cmp::gt_eq(&lhs, &rhs).map_err(|e| e.into());
108        }
109
110        let cmp = make_comparator(lhs, rhs, SORT_OPTIONS)?;
111
112        assert_eq_or_internal_err!(
113            lhs.len(),
114            rhs.len(),
115            "All arrays should have the same length for greatest comparison"
116        );
117
118        let values = BooleanBuffer::collect_bool(lhs.len(), |i| cmp(i, i).is_ge());
119
120        // No nulls as we only want to keep the values that are larger, its either true or false
121        Ok(BooleanArray::new(values, None))
122    }
123}
124
125impl ScalarUDFImpl for GreatestFunc {
126    fn as_any(&self) -> &dyn Any {
127        self
128    }
129
130    fn name(&self) -> &str {
131        "greatest"
132    }
133
134    fn signature(&self) -> &Signature {
135        &self.signature
136    }
137
138    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
139        Ok(arg_types[0].clone())
140    }
141
142    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
143        super::greatest_least_utils::execute_conditional::<Self>(&args.args)
144    }
145
146    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
147        let coerced_type =
148            super::greatest_least_utils::find_coerced_type::<Self>(arg_types)?;
149
150        Ok(vec![coerced_type; arg_types.len()])
151    }
152
153    fn documentation(&self) -> Option<&Documentation> {
154        self.doc()
155    }
156}
157
158#[cfg(test)]
159mod test {
160    use crate::core;
161    use arrow::datatypes::DataType;
162    use datafusion_expr::ScalarUDFImpl;
163
164    #[test]
165    fn test_greatest_return_types_without_common_supertype_in_arg_type() {
166        let greatest = core::greatest::GreatestFunc::new();
167        let return_type = greatest
168            .coerce_types(&[DataType::Decimal128(10, 3), DataType::Decimal128(10, 4)])
169            .unwrap();
170        assert_eq!(
171            return_type,
172            vec![DataType::Decimal128(11, 4), DataType::Decimal128(11, 4)]
173        );
174    }
175}
datafusion_functions/core/greatest.rs

datafusion_functions/core/
greatest.rs
