Add Spark atan2 function (7113)

oap-project · Oct 18, 2023 · 80deccd · 80deccd
1 parent bec32c0
commit 80deccd
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 1 deletion.
diff --git a/velox/docs/functions/spark/math.rst b/velox/docs/functions/spark/math.rst
@@ -18,6 +18,10 @@ Mathematical Functions
 
     Returns inverse hyperbolic sine of ``x``.
 
+.. spark:function:: atan2(x, y) -> double
+
+    Returns the angle in radians between the positive x-axis of a plane and the point given by the coordinates(x, y).
+
 .. spark:function:: atanh(x) -> double
 
     Returns inverse hyperbolic tangent of ``x``.

diff --git a/velox/functions/sparksql/Arithmetic.h b/velox/functions/sparksql/Arithmetic.h
@@ -280,4 +280,13 @@ struct Log10Function {
     return true;
   }
 };
+
+template <typename T>
+struct Atan2Function {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void call(TInput& result, TInput y, TInput x) {
+    result = std::atan2(y + 0.0, x + 0.0);
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/RegisterArithmetic.cpp b/velox/functions/sparksql/RegisterArithmetic.cpp
@@ -95,6 +95,7 @@ void registerArithmeticFunctions(const std::string& prefix) {
   VELOX_REGISTER_VECTOR_FUNCTION(udf_decimal_sub, prefix + "subtract");
   VELOX_REGISTER_VECTOR_FUNCTION(udf_decimal_mul, prefix + "multiply");
   VELOX_REGISTER_VECTOR_FUNCTION(udf_decimal_div, prefix + "divide");
+  registerFunction<Atan2Function, double, double, double>({prefix + "atan2"});
 }
 
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/tests/ArithmeticTest.cpp b/velox/functions/sparksql/tests/ArithmeticTest.cpp
@@ -376,6 +376,14 @@ TEST_F(ArithmeticTest, cot) {
   EXPECT_EQ(cot(0), 1 / std::tan(0));
 }
 
+TEST_F(ArithmeticTest, atan2) {
+  const auto atan2 = [&](std::optional<double> y, std::optional<double> x) {
+    return evaluateOnce<double>("atan2(c0, c1)", y, x);
+  };
+
+  EXPECT_EQ(atan2(0, 0), 0.0);
+}
+
 class LogNTest : public SparkFunctionBaseTest {
  protected:
   static constexpr float kInf = std::numeric_limits<double>::infinity();
@@ -400,6 +408,5 @@ TEST_F(LogNTest, log10) {
   EXPECT_EQ(log10(-1.0), std::nullopt);
   EXPECT_EQ(log10(kInf), kInf);
 }
-
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test