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 std::sync::Arc;
19
20use arrow::array::{ArrayRef, AsArray, PrimitiveArray};
21use arrow::compute::try_binary;
22use arrow::datatypes::DataType;
23use arrow::datatypes::DataType::Int64;
24use arrow::datatypes::Int64Type;
25
26use arrow::error::ArrowError;
27use datafusion_common::{Result, exec_err};
28use datafusion_expr::{
29    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
30    Volatility,
31};
32use datafusion_macros::user_doc;
33
34use super::gcd::unsigned_gcd;
35use crate::utils::make_scalar_function;
36
37#[user_doc(
38    doc_section(label = "Math Functions"),
39    description = "Returns the least common multiple of `expression_x` and `expression_y`. Returns 0 if either input is zero.",
40    syntax_example = "lcm(expression_x, expression_y)",
41    sql_example = r#"```sql
42> SELECT lcm(4, 5);
43+----------+
44| lcm(4,5) |
45+----------+
46| 20       |
47+----------+
48```"#,
49    standard_argument(name = "expression_x", prefix = "First numeric"),
50    standard_argument(name = "expression_y", prefix = "Second numeric")
51)]
52#[derive(Debug, PartialEq, Eq, Hash)]
53pub struct LcmFunc {
54    signature: Signature,
55}
56
57impl Default for LcmFunc {
58    fn default() -> Self {
59        LcmFunc::new()
60    }
61}
62
63impl LcmFunc {
64    pub fn new() -> Self {
65        use DataType::*;
66        Self {
67            signature: Signature::uniform(2, vec![Int64], Volatility::Immutable),
68        }
69    }
70}
71
72impl ScalarUDFImpl for LcmFunc {
73    fn name(&self) -> &str {
74        "lcm"
75    }
76
77    fn signature(&self) -> &Signature {
78        &self.signature
79    }
80
81    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
82        Ok(Int64)
83    }
84
85    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
86        make_scalar_function(lcm, vec![])(&args.args)
87    }
88
89    fn documentation(&self) -> Option<&Documentation> {
90        self.doc()
91    }
92}
93
94/// Lcm SQL function
95fn lcm(args: &[ArrayRef]) -> Result<ArrayRef> {
96    let compute_lcm = |x: i64, y: i64| -> Result<i64, ArrowError> {
97        if x == 0 || y == 0 {
98            return Ok(0);
99        }
100
101        // lcm(x, y) = |x| * |y| / gcd(|x|, |y|)
102        let a = x.unsigned_abs();
103        let b = y.unsigned_abs();
104        let gcd = unsigned_gcd(a, b);
105        // gcd is not zero since both a and b are not zero, so the division is safe.
106        (a / gcd)
107            .checked_mul(b)
108            .and_then(|v| i64::try_from(v).ok())
109            .ok_or_else(|| {
110                ArrowError::ComputeError(format!(
111                    "Signed integer overflow in LCM({x}, {y})"
112                ))
113            })
114    };
115
116    match args[0].data_type() {
117        Int64 => {
118            let arg1 = args[0].as_primitive::<Int64Type>();
119            let arg2 = args[1].as_primitive::<Int64Type>();
120
121            let result: PrimitiveArray<Int64Type> = try_binary(arg1, arg2, compute_lcm)?;
122            Ok(Arc::new(result) as ArrayRef)
123        }
124        other => exec_err!("Unsupported data type {other:?} for function lcm"),
125    }
126}
datafusion_functions/math/lcm.rs

datafusion_functions/math/
lcm.rs
