diff --git a/datafusion/spark/src/function/math/atan2.rs b/datafusion/spark/src/function/math/atan2.rs new file mode 100644 index 0000000000000..70cc1ffeb25a1 --- /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. +/// +/// +/// +/// `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 { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(spark_atan2, vec![])(&args.args) + } +} + +fn spark_atan2(args: &[ArrayRef]) -> Result { + // 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::(); + let x = x.as_primitive::(); + 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 fb57b536f26ec..53cedaef9147c 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> { 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 eb644854c402d..11e7a90202ddc 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