Support spark asinh, acosh, atanh, sec, csc math functions (facebooki…

…ncubator#274)
PHILO-HE · Jun 26, 2023 · 2f9c7e2 · 2f9c7e2
1 parent 49580f1
commit 2f9c7e2
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 0 deletions.
diff --git a/velox/functions/sparksql/Arithmetic.h b/velox/functions/sparksql/Arithmetic.h
@@ -165,4 +165,44 @@ struct FloorFunction {
   }
 };
 
+template <typename T>
+struct AcoshFunction {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void call(TInput& result, TInput a) {
+    result = std::acosh(a);
+  }
+};
+
+template <typename T>
+struct AsinhFunction {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void call(TInput& result, TInput a) {
+    result = std::asinh(a);
+  }
+};
+
+template <typename T>
+struct AtanhFunction {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void call(TInput& result, TInput a) {
+    result = std::atanh(a);
+  }
+};
+
+template <typename T>
+struct SecFunction {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void call(TInput& result, TInput a) {
+    result = 1 / std::cos(a);
+  }
+};
+
+template <typename T>
+struct CscFunction {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE void call(TInput& result, TInput a) {
+    result = 1 / std::sin(a);
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/RegisterArithmetic.cpp b/velox/functions/sparksql/RegisterArithmetic.cpp
@@ -31,6 +31,11 @@ void registerArithmeticFunctions(const std::string& prefix) {
   registerUnaryNumeric<UnaryMinusFunction>({prefix + "unaryminus"});
   // Math functions.
   registerUnaryNumeric<AbsFunction>({prefix + "abs"});
+  registerFunction<AcoshFunction, double, double>({prefix + "acosh"});
+  registerFunction<AsinhFunction, double, double>({prefix + "asinh"});
+  registerFunction<AtanhFunction, double, double>({prefix + "atanh"});
+  registerFunction<SecFunction, double, double>({prefix + "sec"});
+  registerFunction<CscFunction, double, double>({prefix + "csc"});
   registerFunction<ExpFunction, double, double>({prefix + "exp"});
   registerBinaryIntegral<PModFunction>({prefix + "pmod"});
   registerBinaryFloatingPoint<PModFloatFunction>({prefix + "pmod"});

diff --git a/velox/functions/sparksql/tests/ArithmeticTest.cpp b/velox/functions/sparksql/tests/ArithmeticTest.cpp
@@ -195,6 +195,64 @@ TEST_F(ArithmeticTest, Divide) {
   EXPECT_TRUE(std::isnan(divide(kInf, -kInf).value_or(0)));
 }
 
+TEST_F(ArithmeticTest, acosh) {
+  const auto acosh = [&](std::optional<double> a) {
+    return evaluateOnce<double>("acosh(c0)", a);
+  };
+
+  EXPECT_EQ(acosh(1), 0);
+  EXPECT_TRUE(std::isnan(acosh(0).value_or(0)));
+  EXPECT_EQ(acosh(kInf), kInf);
+  EXPECT_EQ(acosh(std::nullopt), std::nullopt);
+  EXPECT_TRUE(std::isnan(acosh(kNan).value_or(0)));
+}
+
+TEST_F(ArithmeticTest, asinh) {
+  const auto asinh = [&](std::optional<double> a) {
+    return evaluateOnce<double>("asinh(c0)", a);
+  };
+
+  EXPECT_EQ(asinh(0), 0);
+  EXPECT_EQ(asinh(kInf), kInf);
+  EXPECT_EQ(asinh(-kInf), -kInf);
+  EXPECT_EQ(asinh(std::nullopt), std::nullopt);
+  EXPECT_TRUE(std::isnan(asinh(kNan).value_or(0)));
+}
+
+TEST_F(ArithmeticTest, atanh) {
+  const auto atanh = [&](std::optional<double> a) {
+    return evaluateOnce<double>("atanh(c0)", a);
+  };
+
+  EXPECT_EQ(atanh(0), 0);
+  EXPECT_EQ(atanh(1), kInf);
+  EXPECT_EQ(atanh(-1), -kInf);
+  EXPECT_TRUE(std::isnan(atanh(1.1).value_or(0)));
+  EXPECT_TRUE(std::isnan(atanh(-1.1).value_or(0)));
+  EXPECT_EQ(atanh(std::nullopt), std::nullopt);
+  EXPECT_TRUE(std::isnan(atanh(kNan).value_or(0)));
+}
+
+TEST_F(ArithmeticTest, sec) {
+  const auto sec = [&](std::optional<double> a) {
+    return evaluateOnce<double>("sec(c0)", a);
+  };
+
+  EXPECT_EQ(sec(0), 1);
+  EXPECT_EQ(sec(std::nullopt), std::nullopt);
+  EXPECT_TRUE(std::isnan(sec(kNan).value_or(0)));
+}
+
+TEST_F(ArithmeticTest, csc) {
+  const auto csc = [&](std::optional<double> a) {
+    return evaluateOnce<double>("csc(c0)", a);
+  };
+
+  EXPECT_EQ(csc(0), kInf);
+  EXPECT_EQ(csc(std::nullopt), std::nullopt);
+  EXPECT_TRUE(std::isnan(csc(kNan).value_or(0)));
+}
+
 class CeilFloorTest : public SparkFunctionBaseTest {
  protected:
   template <typename T>