datafusion_functions/math/
lcm.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::array::ArrayRef;
19use arrow::datatypes::{
20    DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Int64Type,
21};
22
23use crate::math::common::{lcm_signed, lcm_signed_int};
24use crate::utils::{calculate_binary_decimal_math_cast, calculate_binary_math};
25use datafusion_common::utils::take_function_args;
26use datafusion_common::{Result, exec_err, plan_err};
27use datafusion_expr::{
28    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
29    Volatility,
30};
31use datafusion_expr_common::type_coercion::binary::decimal_coercion;
32use datafusion_macros::user_doc;
33
34#[user_doc(
35    doc_section(label = "Math Functions"),
36    description = "Returns the least common multiple of `expression_x` and `expression_y`. Returns 0 if either input is zero.",
37    syntax_example = "lcm(expression_x, expression_y)",
38    sql_example = r#"```sql
39> SELECT lcm(4, 5);
40+----------+
41| lcm(4,5) |
42+----------+
43| 20       |
44+----------+
45```"#,
46    standard_argument(name = "expression_x", prefix = "First numeric"),
47    standard_argument(name = "expression_y", prefix = "Second numeric")
48)]
49#[derive(Debug, PartialEq, Eq, Hash)]
50pub struct LcmFunc {
51    signature: Signature,
52}
53
54impl Default for LcmFunc {
55    fn default() -> Self {
56        LcmFunc::new()
57    }
58}
59
60impl LcmFunc {
61    pub fn new() -> Self {
62        Self {
63            signature: Signature::user_defined(Volatility::Immutable),
64        }
65    }
66}
67
68impl ScalarUDFImpl for LcmFunc {
69    fn name(&self) -> &str {
70        "lcm"
71    }
72
73    fn signature(&self) -> &Signature {
74        &self.signature
75    }
76
77    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
78        Ok(arg_types[0].clone())
79    }
80
81    fn is_strict(&self) -> bool {
82        true
83    }
84
85    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
86        let [arg1, arg2] = take_function_args(self.name(), arg_types)?;
87
88        let coerced_type = match (arg1, arg2) {
89            (DataType::Null, _) | (_, DataType::Null) => Ok(DataType::Int64),
90            (lhs, rhs) if lhs.is_integer() && rhs.is_integer() => Ok(DataType::Int64),
91            (lhs, rhs) if lhs.is_decimal() || rhs.is_decimal() => {
92                decimal_coercion(lhs, rhs).map(Ok).unwrap_or_else(|| {
93                    plan_err!(
94                        "Unsupported argument types {lhs:?} and {rhs:?} for function {}",
95                        self.name()
96                    )
97                })
98            }
99            (lhs, rhs) => {
100                plan_err!(
101                    "Unsupported argument types {lhs:?} and {rhs:?} for function {}",
102                    self.name()
103                )
104            }
105        }?;
106        Ok(vec![coerced_type.clone(), coerced_type])
107    }
108
109    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
110        let left = &args.args[0].to_array(args.number_rows)?;
111        let right = &args.args[1];
112
113        let arr: ArrayRef = match (left.data_type(), right.data_type()) {
114            (DataType::Int64, _) => calculate_binary_math::<
115                Int64Type,
116                Int64Type,
117                Int64Type,
118                _,
119            >(&left, right, lcm_signed_int)?,
120            (
121                lhs @ DataType::Decimal32(precision, scale),
122                rhs @ DataType::Decimal32(_, _),
123            ) if *lhs == rhs => {
124                calculate_binary_decimal_math_cast::<
125                    Decimal32Type,
126                    Decimal32Type,
127                    Decimal32Type,
128                    _,
129                >(&left, right, lcm_signed, *precision, *scale, lhs)?
130            }
131            (
132                lhs @ DataType::Decimal64(precision, scale),
133                rhs @ DataType::Decimal64(_, _),
134            ) if *lhs == rhs => {
135                calculate_binary_decimal_math_cast::<
136                    Decimal64Type,
137                    Decimal64Type,
138                    Decimal64Type,
139                    _,
140                >(&left, right, lcm_signed, *precision, *scale, lhs)?
141            }
142            (
143                lhs @ DataType::Decimal128(precision, scale),
144                rhs @ DataType::Decimal128(_, _),
145            ) if *lhs == rhs => {
146                calculate_binary_decimal_math_cast::<
147                    Decimal128Type,
148                    Decimal128Type,
149                    Decimal128Type,
150                    _,
151                >(&left, right, lcm_signed, *precision, *scale, lhs)?
152            }
153            (
154                lhs @ DataType::Decimal256(precision, scale),
155                rhs @ DataType::Decimal256(_, _),
156            ) if *lhs == rhs => {
157                calculate_binary_decimal_math_cast::<
158                    Decimal256Type,
159                    Decimal256Type,
160                    Decimal256Type,
161                    _,
162                >(&left, right, lcm_signed, *precision, *scale, lhs)?
163            }
164            (lhs, rhs) => {
165                return exec_err!(
166                    "Unsupported data types {lhs:?} and {rhs:?} for function {}",
167                    self.name()
168                );
169            }
170        };
171        Ok(ColumnarValue::Array(arr))
172    }
173
174    fn documentation(&self) -> Option<&Documentation> {
175        self.doc()
176    }
177}
datafusion_functions/math/lcm.rs

datafusion_functions/math/
lcm.rs
