You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/05 09:51:48 UTC
(datafusion) 01/01: feat: add Spark-compatible atan2 function (#23962)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23962-426b351513be6317734aa2e425ac72ddde790fb4
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 9db5c5de5eb4d3f897dc4c575f505ecd69cdd12a
Author: Karpagam Balasubramaniam <kk...@trulyh.com>
AuthorDate: Wed Aug 5 02:51:25 2026 -0700
feat: add Spark-compatible atan2 function (#23962)
## Which issue does this PR close?
- Closes #23961
- Part of #15914
## Rationale for this change
Spark provides [`atan2(exprY,
exprX)`](https://spark.apache.org/docs/latest/api/sql/#atan2), which
returns the angle in radians between the positive x-axis and the point
given by the coordinates (`exprX`, `exprY`). It was not yet implemented
in `datafusion-spark` — only an auto-generated test stub existed at
`spark/math/atan2.slt` with its query commented out.
## What changes are included in this PR?
- Add `SparkAtan2` (implementing `ScalarUDFImpl`) in
`datafusion/spark/src/function/math/atan2.rs`, computing `y.atan2(x)`
element-wise via the Arrow `binary` kernel.
- Register it in `datafusion/spark/src/function/math/mod.rs`.
- Enable the `atan2.slt` sqllogictest.
The signature is `exact(Float64, Float64) -> Float64`, following the
`datafusion-spark` convention of only accepting Spark-supported types.
NULL in either argument propagates to NULL.
Note: DataFusion core already has an `atan2`, but this is implemented
self-contained rather than delegating — matching the crate convention,
where 14 of the 15 existing `math` functions reimplement rather than
wrap core (including `abs`, `ceil`, `floor`, `round`, `rint`, all of
which also exist in core; only `pow` delegates, as a special case). This
keeps the Spark function library complete and consistent.
## Are these changes tested?
Yes — `datafusion/sqllogictest/test_files/spark/math/atan2.slt` covers:
- standard angles and all four quadrants (atan2 is quadrant-aware via
the signs of both arguments),
- both axes and the negative x-axis (atan2's range extends to π, unlike
`atan`),
- NULL propagation (either argument) and NULL-beats-Infinity,
- NaN in every position, including `atan2(NaN, Infinity) = NaN` (for
atan2, NaN propagates even over Infinity — unlike `hypot`, where
Infinity dominates),
- all four infinity quadrants and the one-infinite-argument cases,
- signed zeros (`atan2(-0, -1) = -π`),
- the array path (including a NULL row).
## Are there any user-facing changes?
Yes — adds the Spark-compatible `atan2` scalar function to
`datafusion-spark`. No breaking changes to public APIs.
---
datafusion/spark/src/function/math/atan2.rs | 84 ++++++++++++
datafusion/spark/src/function/math/mod.rs | 4 +
.../sqllogictest/test_files/spark/math/atan2.slt | 152 ++++++++++++++++++++-
3 files changed, 236 insertions(+), 4 deletions(-)
diff --git a/datafusion/spark/src/function/math/atan2.rs b/datafusion/spark/src/function/math/atan2.rs
new file mode 100644
index 0000000000..70cc1ffeb2
--- /dev/null
+++ b/datafusion/spark/src/function/math/atan2.rs
@@ -0,0 +1,84 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray, Float64Array};
+use arrow::compute::kernels::arity::binary;
+use arrow::datatypes::{DataType, Float64Type};
+use datafusion_common::Result;
+use datafusion_common::utils::take_function_args;
+use datafusion_expr::{
+ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+use datafusion_functions::utils::make_scalar_function;
+
+/// Spark-compatible `atan2` function.
+///
+/// <https://spark.apache.org/docs/latest/api/sql/index.html#atan2>
+///
+/// `atan2(exprY, exprX)` returns the angle in radians between the positive
+/// x-axis and the point given by the coordinates (exprX, exprY).
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkAtan2 {
+ signature: Signature,
+}
+
+impl Default for SparkAtan2 {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkAtan2 {
+ pub fn new() -> Self {
+ Self {
+ // Spark only defines atan2 over doubles
+ signature: Signature::exact(
+ vec![DataType::Float64, DataType::Float64],
+ Volatility::Immutable,
+ ),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkAtan2 {
+ fn name(&self) -> &str {
+ "atan2"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Float64)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ make_scalar_function(spark_atan2, vec![])(&args.args)
+ }
+}
+
+fn spark_atan2(args: &[ArrayRef]) -> Result<ArrayRef> {
+ // Spark arg order is atan2(exprY, exprX); Rust computes y.atan2(x).
+ let [y, x] = take_function_args("atan2", args)?;
+ let y = y.as_primitive::<Float64Type>();
+ let x = x.as_primitive::<Float64Type>();
+ let result: Float64Array = binary(y, x, |y, x| y.atan2(x))?;
+ Ok(Arc::new(result))
+}
diff --git a/datafusion/spark/src/function/math/mod.rs b/datafusion/spark/src/function/math/mod.rs
index fb57b536f2..53cedaef91 100644
--- a/datafusion/spark/src/function/math/mod.rs
+++ b/datafusion/spark/src/function/math/mod.rs
@@ -16,6 +16,7 @@
// under the License.
pub mod abs;
+pub mod atan2;
pub mod bin;
pub mod ceil;
pub mod expm1;
@@ -37,6 +38,7 @@ use datafusion_functions::make_udf_function;
use std::sync::Arc;
make_udf_function!(abs::SparkAbs, abs);
+make_udf_function!(atan2::SparkAtan2, atan2);
make_udf_function!(ceil::SparkCeil, ceil);
make_udf_function!(expm1::SparkExpm1, expm1);
make_udf_function!(factorial::SparkFactorial, factorial);
@@ -59,6 +61,7 @@ pub mod expr_fn {
use datafusion_functions::export_functions;
export_functions!((abs, "Returns abs(expr)", arg1));
+ export_functions!((atan2, "Returns the angle in radians between the positive x-axis and the point (exprX, exprY).", arg1 arg2));
export_functions!((ceil, "Returns the ceiling of expr.", arg1));
export_functions!((expm1, "Returns exp(expr) - 1 as a Float64.", arg1));
export_functions!((
@@ -105,6 +108,7 @@ pub mod expr_fn {
pub fn functions() -> Vec<Arc<ScalarUDF>> {
vec![
abs(),
+ atan2(),
ceil(),
expm1(),
factorial(),
diff --git a/datafusion/sqllogictest/test_files/spark/math/atan2.slt b/datafusion/sqllogictest/test_files/spark/math/atan2.slt
index eb644854c4..11e7a90202 100644
--- a/datafusion/sqllogictest/test_files/spark/math/atan2.slt
+++ b/datafusion/sqllogictest/test_files/spark/math/atan2.slt
@@ -21,7 +21,151 @@
# For more information, please see:
# https://github.com/apache/datafusion/issues/15914
-## Original Query: SELECT atan2(0, 0);
-## PySpark 3.5.5 Result: {'ATAN2(0, 0)': 0.0, 'typeof(ATAN2(0, 0))': 'double', 'typeof(0)': 'int'}
-#query
-#SELECT atan2(0::int);
+# standard angles in radians
+query R
+SELECT atan2(0, 0);
+----
+0
+
+query R
+SELECT atan2(0, 1);
+----
+0
+
+# all four quadrants (atan2 is quadrant-aware via the signs of both arguments)
+query R
+SELECT atan2(1, 1);
+----
+0.785398163397448
+
+query R
+SELECT atan2(1, -1);
+----
+2.356194490192345
+
+query R
+SELECT atan2(-1, -1);
+----
+-2.356194490192345
+
+query R
+SELECT atan2(-1, 1);
+----
+-0.785398163397448
+
+# on the axes
+query R
+SELECT atan2(1, 0);
+----
+1.570796326794897
+
+# negative x-axis: atan2 range extends to pi (atan only reaches +/- pi/2)
+query R
+SELECT atan2(0, -1);
+----
+3.141592653589793
+
+# NULL if either argument is NULL
+query R
+SELECT atan2(NULL::double, 1.0::double);
+----
+NULL
+
+query R
+SELECT atan2(1.0::double, NULL::double);
+----
+NULL
+
+# NaN: any NaN input yields NaN (for atan2, NaN wins even over Infinity)
+query R
+SELECT atan2('NaN'::double, 1.0::double);
+----
+NaN
+
+query R
+SELECT atan2(1.0::double, 'NaN'::double);
+----
+NaN
+
+query R
+SELECT atan2('NaN'::double, 'NaN'::double);
+----
+NaN
+
+query R
+SELECT atan2('NaN'::double, 'Infinity'::double);
+----
+NaN
+
+# NULL beats every special value (validity is checked before the value)
+query R
+SELECT atan2(NULL::double, 'Infinity'::double);
+----
+NULL
+
+# both infinite: quadrant set by the signs (+/- pi/4, +/- 3pi/4)
+query R
+SELECT atan2('Infinity'::double, 'Infinity'::double);
+----
+0.785398163397448
+
+query R
+SELECT atan2('-Infinity'::double, 'Infinity'::double);
+----
+-0.785398163397448
+
+query R
+SELECT atan2('Infinity'::double, '-Infinity'::double);
+----
+2.356194490192345
+
+query R
+SELECT atan2('-Infinity'::double, '-Infinity'::double);
+----
+-2.356194490192345
+
+# one infinite argument
+query R
+SELECT atan2('Infinity'::double, 1.0::double);
+----
+1.570796326794897
+
+query R
+SELECT atan2('-Infinity'::double, 1.0::double);
+----
+-1.570796326794897
+
+query R
+SELECT atan2(1.0::double, 'Infinity'::double);
+----
+0
+
+query R
+SELECT atan2(1.0::double, '-Infinity'::double);
+----
+3.141592653589793
+
+query R
+SELECT atan2(-1.0::double, '-Infinity'::double);
+----
+-3.141592653589793
+
+# signed zeros: -0 flips the sign on the negative x-axis (atan2(+0, -1) = pi above; atan2(-0, -1) = -pi)
+query R
+SELECT atan2(-0.0::double, -1.0::double);
+----
+-3.141592653589793
+
+# -0 in the first argument still returns 0 on the positive x-axis
+query R
+SELECT atan2(-0.0::double, 1.0::double);
+----
+0
+
+# array path, including a NULL row
+query R
+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);
+----
+0
+0.785398163397448
+NULL
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org