Skip to content

Commit 5419c96

Browse files
KarpagamKarthikeyanKarpagamKarthikeyan
committed
feat: add Spark-compatible atan2 function
1 parent 74eebbb commit 5419c96

3 files changed

Lines changed: 236 additions & 4 deletions

File tree

‎datafusion/spark/src/function/math/atan2.rs‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
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 std::sync::Arc;
19+
20+
use arrow::array::{ArrayRef, AsArray, Float64Array};
21+
use arrow::compute::kernels::arity::binary;
22+
use arrow::datatypes::{DataType, Float64Type};
23+
use datafusion_common::Result;
24+
use datafusion_common::utils::take_function_args;
25+
use datafusion_expr::{
26+
ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
27+
};
28+
use datafusion_functions::utils::make_scalar_function;
29+
30+
/// Spark-compatible `atan2` function.
31+
///
32+
/// <https://spark.apache.org/docs/latest/api/sql/index.html#atan2>
33+
///
34+
/// `atan2(exprY, exprX)` returns the angle in radians between the positive
35+
/// x-axis and the point given by the coordinates (exprX, exprY).
36+
#[derive(Debug, PartialEq, Eq, Hash)]
37+
pub struct SparkAtan2 {
38+
signature: Signature,
39+
}
40+
41+
impl Default for SparkAtan2 {
42+
fn default() -> Self {
43+
Self::new()
44+
}
45+
}
46+
47+
impl SparkAtan2 {
48+
pub fn new() -> Self {
49+
Self {
50+
// Spark only defines atan2 over doubles
51+
signature: Signature::exact(
52+
vec![DataType::Float64, DataType::Float64],
53+
Volatility::Immutable,
54+
),
55+
}
56+
}
57+
}
58+
59+
impl ScalarUDFImpl for SparkAtan2 {
60+
fn name(&self) -> &str {
61+
"atan2"
62+
}
63+
64+
fn signature(&self) -> &Signature {
65+
&self.signature
66+
}
67+
68+
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
69+
Ok(DataType::Float64)
70+
}
71+
72+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
73+
make_scalar_function(spark_atan2, vec![])(&args.args)
74+
}
75+
}
76+
77+
fn spark_atan2(args: &[ArrayRef]) -> Result<ArrayRef> {
78+
// Spark arg order is atan2(exprY, exprX); Rust computes y.atan2(x).
79+
let [y, x] = take_function_args("atan2", args)?;
80+
let y = y.as_primitive::<Float64Type>();
81+
let x = x.as_primitive::<Float64Type>();
82+
let result: Float64Array = binary(y, x, |y, x| y.atan2(x))?;
83+
Ok(Arc::new(result))
84+
}

‎datafusion/spark/src/function/math/mod.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
// under the License.
1717

1818
pub mod abs;
19+
pub mod atan2;
1920
pub mod bin;
2021
pub mod ceil;
2122
pub mod expm1;
@@ -37,6 +38,7 @@ use datafusion_functions::make_udf_function;
3738
use std::sync::Arc;
3839

3940
make_udf_function!(abs::SparkAbs, abs);
41+
make_udf_function!(atan2::SparkAtan2, atan2);
4042
make_udf_function!(ceil::SparkCeil, ceil);
4143
make_udf_function!(expm1::SparkExpm1, expm1);
4244
make_udf_function!(factorial::SparkFactorial, factorial);
@@ -59,6 +61,7 @@ pub mod expr_fn {
5961
use datafusion_functions::export_functions;
6062

6163
export_functions!((abs, "Returns abs(expr)", arg1));
64+
export_functions!((atan2, "Returns the angle in radians between the positive x-axis and the point (exprX, exprY).", arg1 arg2));
6265
export_functions!((ceil, "Returns the ceiling of expr.", arg1));
6366
export_functions!((expm1, "Returns exp(expr) - 1 as a Float64.", arg1));
6467
export_functions!((
@@ -105,6 +108,7 @@ pub mod expr_fn {
105108
pub fn functions() -> Vec<Arc<ScalarUDF>> {
106109
vec![
107110
abs(),
111+
atan2(),
108112
ceil(),
109113
expm1(),
110114
factorial(),

‎datafusion/sqllogictest/test_files/spark/math/atan2.slt‎

Lines changed: 148 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,151 @@
2121
# For more information, please see:
2222
# https://github.com/apache/datafusion/issues/15914
2323

24-
## Original Query: SELECT atan2(0, 0);
25-
## PySpark 3.5.5 Result: {'ATAN2(0, 0)': 0.0, 'typeof(ATAN2(0, 0))': 'double', 'typeof(0)': 'int'}
26-
#query
27-
#SELECT atan2(0::int);
24+
# standard angles in radians
25+
query R
26+
SELECT atan2(0, 0);
27+
----
28+
0
29+
30+
query R
31+
SELECT atan2(0, 1);
32+
----
33+
0
34+
35+
# all four quadrants (atan2 is quadrant-aware via the signs of both arguments)
36+
query R
37+
SELECT atan2(1, 1);
38+
----
39+
0.785398163397448
40+
41+
query R
42+
SELECT atan2(1, -1);
43+
----
44+
2.356194490192345
45+
46+
query R
47+
SELECT atan2(-1, -1);
48+
----
49+
-2.356194490192345
50+
51+
query R
52+
SELECT atan2(-1, 1);
53+
----
54+
-0.785398163397448
55+
56+
# on the axes
57+
query R
58+
SELECT atan2(1, 0);
59+
----
60+
1.570796326794897
61+
62+
# negative x-axis: atan2 range extends to pi (atan only reaches +/- pi/2)
63+
query R
64+
SELECT atan2(0, -1);
65+
----
66+
3.141592653589793
67+
68+
# NULL if either argument is NULL
69+
query R
70+
SELECT atan2(NULL::double, 1.0::double);
71+
----
72+
NULL
73+
74+
query R
75+
SELECT atan2(1.0::double, NULL::double);
76+
----
77+
NULL
78+
79+
# NaN: any NaN input yields NaN (for atan2, NaN wins even over Infinity)
80+
query R
81+
SELECT atan2('NaN'::double, 1.0::double);
82+
----
83+
NaN
84+
85+
query R
86+
SELECT atan2(1.0::double, 'NaN'::double);
87+
----
88+
NaN
89+
90+
query R
91+
SELECT atan2('NaN'::double, 'NaN'::double);
92+
----
93+
NaN
94+
95+
query R
96+
SELECT atan2('NaN'::double, 'Infinity'::double);
97+
----
98+
NaN
99+
100+
# NULL beats every special value (validity is checked before the value)
101+
query R
102+
SELECT atan2(NULL::double, 'Infinity'::double);
103+
----
104+
NULL
105+
106+
# both infinite: quadrant set by the signs (+/- pi/4, +/- 3pi/4)
107+
query R
108+
SELECT atan2('Infinity'::double, 'Infinity'::double);
109+
----
110+
0.785398163397448
111+
112+
query R
113+
SELECT atan2('-Infinity'::double, 'Infinity'::double);
114+
----
115+
-0.785398163397448
116+
117+
query R
118+
SELECT atan2('Infinity'::double, '-Infinity'::double);
119+
----
120+
2.356194490192345
121+
122+
query R
123+
SELECT atan2('-Infinity'::double, '-Infinity'::double);
124+
----
125+
-2.356194490192345
126+
127+
# one infinite argument
128+
query R
129+
SELECT atan2('Infinity'::double, 1.0::double);
130+
----
131+
1.570796326794897
132+
133+
query R
134+
SELECT atan2('-Infinity'::double, 1.0::double);
135+
----
136+
-1.570796326794897
137+
138+
query R
139+
SELECT atan2(1.0::double, 'Infinity'::double);
140+
----
141+
0
142+
143+
query R
144+
SELECT atan2(1.0::double, '-Infinity'::double);
145+
----
146+
3.141592653589793
147+
148+
query R
149+
SELECT atan2(-1.0::double, '-Infinity'::double);
150+
----
151+
-3.141592653589793
152+
153+
# signed zeros: -0 flips the sign on the negative x-axis (atan2(+0, -1) = pi above; atan2(-0, -1) = -pi)
154+
query R
155+
SELECT atan2(-0.0::double, -1.0::double);
156+
----
157+
-3.141592653589793
158+
159+
# -0 in the first argument still returns 0 on the positive x-axis
160+
query R
161+
SELECT atan2(-0.0::double, 1.0::double);
162+
----
163+
0
164+
165+
# array path, including a NULL row
166+
query R
167+
SELECT atan2(a, b) FROM (VALUES (0.0::double, 1.0::double), (1.0::double, 1.0::double), (NULL::double, 1.0::double)) AS t(a, b);
168+
----
169+
0
170+
0.785398163397448
171+
NULL

0 commit comments

Comments
 (0)