feat: add Spark-compatible atan2 function

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))
+}

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(),

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