Skip to content

Commit bece785

Browse files
authored
move Floor, Gcd, Lcm, Pi to datafusion-functions (#9976)
* move Floor, Gcd, Lcm, Pi to datafusion-functions
1 parent 8c9e567 commit bece785

File tree

21 files changed

+588
-344
lines changed

21 files changed

+588
-344
lines changed

datafusion/expr/src/built_in_function.rs

Lines changed: 1 addition & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -45,20 +45,12 @@ pub enum BuiltinScalarFunction {
4545
Exp,
4646
/// factorial
4747
Factorial,
48-
/// floor
49-
Floor,
50-
/// gcd, Greatest common divisor
51-
Gcd,
52-
/// lcm, Least common multiple
53-
Lcm,
5448
/// iszero
5549
Iszero,
5650
/// log, same as log10
5751
Log,
5852
/// nanvl
5953
Nanvl,
60-
/// pi
61-
Pi,
6254
/// power
6355
Power,
6456
/// round
@@ -135,13 +127,9 @@ impl BuiltinScalarFunction {
135127
BuiltinScalarFunction::Coalesce => Volatility::Immutable,
136128
BuiltinScalarFunction::Exp => Volatility::Immutable,
137129
BuiltinScalarFunction::Factorial => Volatility::Immutable,
138-
BuiltinScalarFunction::Floor => Volatility::Immutable,
139-
BuiltinScalarFunction::Gcd => Volatility::Immutable,
140130
BuiltinScalarFunction::Iszero => Volatility::Immutable,
141-
BuiltinScalarFunction::Lcm => Volatility::Immutable,
142131
BuiltinScalarFunction::Log => Volatility::Immutable,
143132
BuiltinScalarFunction::Nanvl => Volatility::Immutable,
144-
BuiltinScalarFunction::Pi => Volatility::Immutable,
145133
BuiltinScalarFunction::Power => Volatility::Immutable,
146134
BuiltinScalarFunction::Round => Volatility::Immutable,
147135
BuiltinScalarFunction::Cot => Volatility::Immutable,
@@ -183,13 +171,10 @@ impl BuiltinScalarFunction {
183171
BuiltinScalarFunction::InitCap => {
184172
utf8_to_str_type(&input_expr_types[0], "initcap")
185173
}
186-
BuiltinScalarFunction::Pi => Ok(Float64),
187174
BuiltinScalarFunction::Random => Ok(Float64),
188175
BuiltinScalarFunction::EndsWith => Ok(Boolean),
189176

190-
BuiltinScalarFunction::Factorial
191-
| BuiltinScalarFunction::Gcd
192-
| BuiltinScalarFunction::Lcm => Ok(Int64),
177+
BuiltinScalarFunction::Factorial => Ok(Int64),
193178

194179
BuiltinScalarFunction::Power => match &input_expr_types[0] {
195180
Int64 => Ok(Int64),
@@ -210,7 +195,6 @@ impl BuiltinScalarFunction {
210195

211196
BuiltinScalarFunction::Ceil
212197
| BuiltinScalarFunction::Exp
213-
| BuiltinScalarFunction::Floor
214198
| BuiltinScalarFunction::Round
215199
| BuiltinScalarFunction::Trunc
216200
| BuiltinScalarFunction::Cot => match input_expr_types[0] {
@@ -248,7 +232,6 @@ impl BuiltinScalarFunction {
248232
],
249233
self.volatility(),
250234
),
251-
BuiltinScalarFunction::Pi => Signature::exact(vec![], self.volatility()),
252235
BuiltinScalarFunction::Random => Signature::exact(vec![], self.volatility()),
253236
BuiltinScalarFunction::Power => Signature::one_of(
254237
vec![Exact(vec![Int64, Int64]), Exact(vec![Float64, Float64])],
@@ -289,12 +272,8 @@ impl BuiltinScalarFunction {
289272
BuiltinScalarFunction::Factorial => {
290273
Signature::uniform(1, vec![Int64], self.volatility())
291274
}
292-
BuiltinScalarFunction::Gcd | BuiltinScalarFunction::Lcm => {
293-
Signature::uniform(2, vec![Int64], self.volatility())
294-
}
295275
BuiltinScalarFunction::Ceil
296276
| BuiltinScalarFunction::Exp
297-
| BuiltinScalarFunction::Floor
298277
| BuiltinScalarFunction::Cot => {
299278
// math expressions expect 1 argument of type f64 or f32
300279
// priority is given to f64 because e.g. `sqrt(1i32)` is in IR (real numbers) and thus we
@@ -319,10 +298,8 @@ impl BuiltinScalarFunction {
319298
BuiltinScalarFunction::Ceil
320299
| BuiltinScalarFunction::Exp
321300
| BuiltinScalarFunction::Factorial
322-
| BuiltinScalarFunction::Floor
323301
| BuiltinScalarFunction::Round
324302
| BuiltinScalarFunction::Trunc
325-
| BuiltinScalarFunction::Pi
326303
) {
327304
Some(vec![Some(true)])
328305
} else if *self == BuiltinScalarFunction::Log {
@@ -339,13 +316,9 @@ impl BuiltinScalarFunction {
339316
BuiltinScalarFunction::Cot => &["cot"],
340317
BuiltinScalarFunction::Exp => &["exp"],
341318
BuiltinScalarFunction::Factorial => &["factorial"],
342-
BuiltinScalarFunction::Floor => &["floor"],
343-
BuiltinScalarFunction::Gcd => &["gcd"],
344319
BuiltinScalarFunction::Iszero => &["iszero"],
345-
BuiltinScalarFunction::Lcm => &["lcm"],
346320
BuiltinScalarFunction::Log => &["log"],
347321
BuiltinScalarFunction::Nanvl => &["nanvl"],
348-
BuiltinScalarFunction::Pi => &["pi"],
349322
BuiltinScalarFunction::Power => &["power", "pow"],
350323
BuiltinScalarFunction::Random => &["random"],
351324
BuiltinScalarFunction::Round => &["round"],

datafusion/expr/src/expr_fn.rs

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -297,11 +297,6 @@ pub fn concat_ws(sep: Expr, values: Vec<Expr>) -> Expr {
297297
))
298298
}
299299

300-
/// Returns an approximate value of π
301-
pub fn pi() -> Expr {
302-
Expr::ScalarFunction(ScalarFunction::new(BuiltinScalarFunction::Pi, vec![]))
303-
}
304-
305300
/// Returns a random value in the range 0.0 <= x < 1.0
306301
pub fn random() -> Expr {
307302
Expr::ScalarFunction(ScalarFunction::new(BuiltinScalarFunction::Random, vec![]))
@@ -537,12 +532,6 @@ macro_rules! nary_scalar_expr {
537532
// math functions
538533
scalar_expr!(Cot, cot, num, "cotangent of a number");
539534
scalar_expr!(Factorial, factorial, num, "factorial");
540-
scalar_expr!(
541-
Floor,
542-
floor,
543-
num,
544-
"nearest integer less than or equal to argument"
545-
);
546535
scalar_expr!(
547536
Ceil,
548537
ceil,
@@ -556,8 +545,7 @@ nary_scalar_expr!(
556545
"truncate toward zero, with optional precision"
557546
);
558547
scalar_expr!(Exp, exp, num, "exponential");
559-
scalar_expr!(Gcd, gcd, arg_1 arg_2, "greatest common divisor");
560-
scalar_expr!(Lcm, lcm, arg_1 arg_2, "least common multiple");
548+
561549
scalar_expr!(Power, power, base exponent, "`base` raised to the power of `exponent`");
562550
scalar_expr!(Log, log, base x, "logarithm of a `x` for a particular `base`");
563551

@@ -974,7 +962,6 @@ mod test {
974962
fn scalar_function_definitions() {
975963
test_unary_scalar_expr!(Cot, cot);
976964
test_unary_scalar_expr!(Factorial, factorial);
977-
test_unary_scalar_expr!(Floor, floor);
978965
test_unary_scalar_expr!(Ceil, ceil);
979966
test_nary_scalar_expr!(Round, round, input);
980967
test_nary_scalar_expr!(Round, round, input, decimal_places);
@@ -984,8 +971,6 @@ mod test {
984971
test_scalar_expr!(Nanvl, nanvl, x, y);
985972
test_scalar_expr!(Iszero, iszero, input);
986973

987-
test_scalar_expr!(Gcd, gcd, arg_1, arg_2);
988-
test_scalar_expr!(Lcm, lcm, arg_1, arg_2);
989974
test_scalar_expr!(InitCap, initcap, string);
990975
test_scalar_expr!(EndsWith, ends_with, string, characters);
991976
}

datafusion/functions/src/math/gcd.rs

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
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+
18+
use arrow::array::{ArrayRef, Int64Array};
19+
use std::any::Any;
20+
use std::mem::swap;
21+
use std::sync::Arc;
22+
23+
use arrow::datatypes::DataType;
24+
use arrow::datatypes::DataType::Int64;
25+
26+
use crate::utils::make_scalar_function;
27+
use datafusion_common::{exec_err, DataFusionError, Result};
28+
use datafusion_expr::ColumnarValue;
29+
use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
30+
31+
#[derive(Debug)]
32+
pub struct GcdFunc {
33+
signature: Signature,
34+
}
35+
36+
impl Default for GcdFunc {
37+
fn default() -> Self {
38+
Self::new()
39+
}
40+
}
41+
42+
impl GcdFunc {
43+
pub fn new() -> Self {
44+
use DataType::*;
45+
Self {
46+
signature: Signature::uniform(2, vec![Int64], Volatility::Immutable),
47+
}
48+
}
49+
}
50+
51+
impl ScalarUDFImpl for GcdFunc {
52+
fn as_any(&self) -> &dyn Any {
53+
self
54+
}
55+
56+
fn name(&self) -> &str {
57+
"gcd"
58+
}
59+
60+
fn signature(&self) -> &Signature {
61+
&self.signature
62+
}
63+
64+
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
65+
Ok(Int64)
66+
}
67+
68+
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
69+
make_scalar_function(gcd, vec![])(args)
70+
}
71+
}
72+
73+
/// Gcd SQL function
74+
fn gcd(args: &[ArrayRef]) -> Result<ArrayRef> {
75+
match args[0].data_type() {
76+
Int64 => Ok(Arc::new(make_function_inputs2!(
77+
&args[0],
78+
&args[1],
79+
"x",
80+
"y",
81+
Int64Array,
82+
Int64Array,
83+
{ compute_gcd }
84+
)) as ArrayRef),
85+
other => exec_err!("Unsupported data type {other:?} for function gcd"),
86+
}
87+
}
88+
89+
/// Computes greatest common divisor using Binary GCD algorithm.
90+
pub fn compute_gcd(x: i64, y: i64) -> i64 {
91+
let mut a = x.wrapping_abs();
92+
let mut b = y.wrapping_abs();
93+
94+
if a == 0 {
95+
return b;
96+
}
97+
if b == 0 {
98+
return a;
99+
}
100+
101+
let shift = (a | b).trailing_zeros();
102+
a >>= shift;
103+
b >>= shift;
104+
a >>= a.trailing_zeros();
105+
106+
loop {
107+
b >>= b.trailing_zeros();
108+
if a > b {
109+
swap(&mut a, &mut b);
110+
}
111+
112+
b -= a;
113+
114+
if b == 0 {
115+
return a << shift;
116+
}
117+
}
118+
}
119+
120+
#[cfg(test)]
121+
mod test {
122+
use std::sync::Arc;
123+
124+
use arrow::array::{ArrayRef, Int64Array};
125+
126+
use crate::math::gcd::gcd;
127+
use datafusion_common::cast::as_int64_array;
128+
129+
#[test]
130+
fn test_gcd_i64() {
131+
let args: Vec<ArrayRef> = vec![
132+
Arc::new(Int64Array::from(vec![0, 3, 25, -16])), // x
133+
Arc::new(Int64Array::from(vec![0, -2, 15, 8])), // y
134+
];
135+
136+
let result = gcd(&args).expect("failed to initialize function gcd");
137+
let ints = as_int64_array(&result).expect("failed to initialize function gcd");
138+
139+
assert_eq!(ints.len(), 4);
140+
assert_eq!(ints.value(0), 0);
141+
assert_eq!(ints.value(1), 1);
142+
assert_eq!(ints.value(2), 5);
143+
assert_eq!(ints.value(3), 8);
144+
}
145+
}

0 commit comments

Comments
 (0)