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 <[email protected]> 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: [email protected] For additional commands, e-mail: [email protected]
