datafusion_functions/math/
gcd.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, AsArray, PrimitiveArray};
19use arrow::compute::try_binary;
20use arrow::datatypes::{
21    DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Int64Type,
22};
23use std::sync::Arc;
24
25use crate::math::common::{gcd_signed, gcd_signed_int, unsigned_gcd};
26use crate::utils::calculate_binary_decimal_math_cast;
27use datafusion_common::utils::take_function_args;
28use datafusion_common::{
29    Result, ScalarValue, exec_err, internal_datafusion_err, plan_err,
30};
31use datafusion_expr::{
32    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
33    Volatility,
34};
35use datafusion_expr_common::type_coercion::binary::decimal_coercion;
36use datafusion_macros::user_doc;
37
38#[user_doc(
39    doc_section(label = "Math Functions"),
40    description = "Returns the greatest common divisor of `expression_x` and `expression_y`. Returns 0 if both inputs are zero.",
41    syntax_example = "gcd(expression_x, expression_y)",
42    sql_example = r#"```sql
43> SELECT gcd(48, 18);
44+------------+
45| gcd(48,18) |
46+------------+
47| 6          |
48+------------+
49```"#,
50    standard_argument(name = "expression_x", prefix = "First numeric"),
51    standard_argument(name = "expression_y", prefix = "Second numeric")
52)]
53#[derive(Debug, PartialEq, Eq, Hash)]
54pub struct GcdFunc {
55    signature: Signature,
56}
57
58impl Default for GcdFunc {
59    fn default() -> Self {
60        Self::new()
61    }
62}
63
64impl GcdFunc {
65    pub fn new() -> Self {
66        Self {
67            signature: Signature::user_defined(Volatility::Immutable),
68        }
69    }
70}
71
72impl ScalarUDFImpl for GcdFunc {
73    fn name(&self) -> &str {
74        "gcd"
75    }
76
77    fn signature(&self) -> &Signature {
78        &self.signature
79    }
80
81    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
82        Ok(arg_types[0].clone())
83    }
84
85    fn is_strict(&self) -> bool {
86        true
87    }
88
89    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
90        let [arg1, arg2] = take_function_args(self.name(), arg_types)?;
91
92        let coerced_type = match (arg1, arg2) {
93            (DataType::Null, _) | (_, DataType::Null) => Ok(DataType::Int64),
94            (lhs, rhs) if lhs.is_integer() && rhs.is_integer() => Ok(DataType::Int64),
95            (lhs, rhs) if lhs.is_decimal() || rhs.is_decimal() => {
96                decimal_coercion(lhs, rhs).map(Ok).unwrap_or_else(|| {
97                    plan_err!(
98                        "Unsupported argument types {lhs:?} and {rhs:?} for function {}",
99                        self.name()
100                    )
101                })
102            }
103            (lhs, rhs) => {
104                plan_err!(
105                    "Unsupported argument types {lhs:?} and {rhs:?} for function {}",
106                    self.name()
107                )
108            }
109        }?;
110        Ok(vec![coerced_type.clone(), coerced_type])
111    }
112
113    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
114        let number_rows = args.number_rows;
115        let args: [ColumnarValue; 2] = args.args.try_into().map_err(|_| {
116            internal_datafusion_err!("Expected 2 arguments for function gcd")
117        })?;
118
119        if args[0].data_type() == DataType::Int64 {
120            // Optimized path for both integers
121            match args {
122                [ColumnarValue::Array(a), ColumnarValue::Array(b)] => {
123                    compute_gcd_for_arrays(&a, &b)
124                }
125                [
126                    ColumnarValue::Scalar(ScalarValue::Int64(a)),
127                    ColumnarValue::Scalar(ScalarValue::Int64(b)),
128                ] => match (a, b) {
129                    (Some(a), Some(b)) => Ok(ColumnarValue::Scalar(ScalarValue::Int64(
130                        Some(gcd_signed_int(a, b)?),
131                    ))),
132                    _ => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))),
133                },
134                [
135                    ColumnarValue::Array(a),
136                    ColumnarValue::Scalar(ScalarValue::Int64(b)),
137                ] => compute_gcd_with_scalar(&a, b),
138                [
139                    ColumnarValue::Scalar(ScalarValue::Int64(a)),
140                    ColumnarValue::Array(b),
141                ] => compute_gcd_with_scalar(&b, a),
142                _ => exec_err!("Unsupported argument types for function gcd"),
143            }
144        } else {
145            // Decimal path: convert left to array and use generic helper
146            let left = args[0].to_array(number_rows)?;
147            let right = &args[1];
148
149            let arr: ArrayRef = match (left.data_type(), right.data_type()) {
150                (
151                    lhs @ DataType::Decimal32(precision, scale),
152                    rhs @ DataType::Decimal32(_, _),
153                ) if *lhs == rhs => calculate_binary_decimal_math_cast::<
154                    Decimal32Type,
155                    Decimal32Type,
156                    Decimal32Type,
157                    _,
158                >(
159                    &left, right, gcd_signed, *precision, *scale, lhs
160                )?,
161                (
162                    lhs @ DataType::Decimal64(precision, scale),
163                    rhs @ DataType::Decimal64(_, _),
164                ) if *lhs == rhs => calculate_binary_decimal_math_cast::<
165                    Decimal64Type,
166                    Decimal64Type,
167                    Decimal64Type,
168                    _,
169                >(
170                    &left, right, gcd_signed, *precision, *scale, lhs
171                )?,
172                (
173                    lhs @ DataType::Decimal128(precision, scale),
174                    rhs @ DataType::Decimal128(_, _),
175                ) if *lhs == rhs => calculate_binary_decimal_math_cast::<
176                    Decimal128Type,
177                    Decimal128Type,
178                    Decimal128Type,
179                    _,
180                >(
181                    &left, right, gcd_signed, *precision, *scale, lhs
182                )?,
183                (
184                    lhs @ DataType::Decimal256(precision, scale),
185                    rhs @ DataType::Decimal256(_, _),
186                ) if *lhs == rhs => calculate_binary_decimal_math_cast::<
187                    Decimal256Type,
188                    Decimal256Type,
189                    Decimal256Type,
190                    _,
191                >(
192                    &left, right, gcd_signed, *precision, *scale, lhs
193                )?,
194                (lhs, rhs) => {
195                    exec_err!(
196                        "Unsupported data types {lhs:?} and {rhs:?} for function {}",
197                        self.name()
198                    )
199                }?,
200            };
201            Ok(ColumnarValue::Array(arr))
202        }
203    }
204
205    fn documentation(&self) -> Option<&Documentation> {
206        self.doc()
207    }
208}
209
210fn compute_gcd_for_arrays(a: &ArrayRef, b: &ArrayRef) -> Result<ColumnarValue> {
211    let a = a.as_primitive::<Int64Type>();
212    let b = b.as_primitive::<Int64Type>();
213    try_binary(a, b, gcd_signed_int)
214        .map(|arr: PrimitiveArray<Int64Type>| {
215            ColumnarValue::Array(Arc::new(arr) as ArrayRef)
216        })
217        .map_err(Into::into) // convert ArrowError to DataFusionError
218}
219
220fn compute_gcd_with_scalar(arr: &ArrayRef, scalar: Option<i64>) -> Result<ColumnarValue> {
221    let prim = arr.as_primitive::<Int64Type>();
222    match scalar {
223        Some(scalar_value) if scalar_value != 0 && scalar_value != i64::MIN => {
224            // The gcd result divides both inputs' absolute values. When the
225            // scalar is neither 0 nor i64::MIN, the gcd's absolute value fits
226            // in i64, so the cast to i64 below cannot overflow. This allows us
227            // to use `unary` instead of `try_unary`, which allows LLVM to
228            // vectorize more effectively.
229            let sv = scalar_value.unsigned_abs();
230            let result: PrimitiveArray<Int64Type> =
231                prim.unary(|val| unsigned_gcd(val.unsigned_abs(), sv) as i64);
232            Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
233        }
234        Some(scalar_value) => {
235            let result: PrimitiveArray<Int64Type> =
236                prim.try_unary(|val| gcd_signed_int(val, scalar_value))?;
237            Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
238        }
239        None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))),
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246
247    #[test]
248    fn test_coercion() {
249        let mut coerced = GcdFunc::new()
250            .coerce_types(&[DataType::Int64, DataType::Int32])
251            .expect("coercion should succeed");
252        assert_eq!(coerced, vec![DataType::Int64, DataType::Int64]);
253
254        coerced = GcdFunc::new()
255            .coerce_types(&[DataType::Decimal128(10, 2), DataType::Int32])
256            .expect("coercion should succeed");
257
258        assert_eq!(
259            coerced,
260            vec![DataType::Decimal128(12, 2), DataType::Decimal128(12, 2)]
261        );
262
263        coerced = GcdFunc::new()
264            .coerce_types(&[DataType::Decimal128(10, 2), DataType::Null])
265            .expect("coercion should succeed");
266
267        assert_eq!(coerced, vec![DataType::Int64, DataType::Int64]);
268    }
269}
datafusion_functions/math/gcd.rs

datafusion_functions/math/
gcd.rs
