From 1b8fde844b191be60932c32f0f3f0d75bbac4813 Mon Sep 17 00:00:00 2001
From: PHILO-HE <feilong.he@intel.com>
Date: Wed, 29 Mar 2023 10:34:23 +0800
Subject: [PATCH] Initial commit

---
 velox/functions/sparksql/Arithmetic.h            | 16 +++++++++++++++-
 velox/functions/sparksql/RegisterArithmetic.cpp  |  3 ++-
 .../functions/sparksql/tests/ArithmeticTest.cpp  | 10 ++++++++++
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/velox/functions/sparksql/Arithmetic.h b/velox/functions/sparksql/Arithmetic.h
index 62fdae2551c2..f2eb9aa990bd 100644
--- a/velox/functions/sparksql/Arithmetic.h
+++ b/velox/functions/sparksql/Arithmetic.h
@@ -25,7 +25,7 @@
 namespace facebook::velox::functions::sparksql {
 
 template <typename T>
-struct PModFunction {
+struct PModIntFunction {
   template <typename TInput>
   FOLLY_ALWAYS_INLINE bool call(TInput& result, const TInput a, const TInput n)
 #if defined(__has_feature)
@@ -43,6 +43,20 @@ struct PModFunction {
   }
 };
 
+template <typename T>
+struct PModFloatFunction {
+  template <typename TInput>
+  FOLLY_ALWAYS_INLINE bool call(TInput& result, const TInput a, const TInput n)
+  {
+    if (UNLIKELY(n == (TInput)0)) {
+      return false;
+    }
+    TInput r = fmod(a, n);
+    result = (r > 0) ? r : fmod(r + n, n);
+    return true;
+  }
+};
+
 template <typename T>
 struct RemainderFunction {
   template <typename TInput>
diff --git a/velox/functions/sparksql/RegisterArithmetic.cpp b/velox/functions/sparksql/RegisterArithmetic.cpp
index 4f72fdf59c30..a17c66cc79ec 100644
--- a/velox/functions/sparksql/RegisterArithmetic.cpp
+++ b/velox/functions/sparksql/RegisterArithmetic.cpp
@@ -32,7 +32,8 @@ void registerArithmeticFunctions(const std::string& prefix) {
   // Math functions.
   registerUnaryNumeric<AbsFunction>({prefix + "abs"});
   registerFunction<ExpFunction, double, double>({prefix + "exp"});
-  registerBinaryIntegral<PModFunction>({prefix + "pmod"});
+  registerBinaryIntegral<PModIntFunction>({prefix + "pmod"});
+  registerBinaryFloatingPoint<PModFloatFunction>({prefix + "pmod"});
   registerFunction<PowerFunction, double, double, double>({prefix + "power"});
   registerUnaryNumeric<RoundFunction>({prefix + "round"});
   registerFunction<RoundFunction, int8_t, int8_t, int32_t>({prefix + "round"});
diff --git a/velox/functions/sparksql/tests/ArithmeticTest.cpp b/velox/functions/sparksql/tests/ArithmeticTest.cpp
index 886123b2d05e..381c9043c155 100644
--- a/velox/functions/sparksql/tests/ArithmeticTest.cpp
+++ b/velox/functions/sparksql/tests/ArithmeticTest.cpp
@@ -67,6 +67,16 @@ TEST_F(PmodTest, int64) {
   EXPECT_EQ(INT64_MAX - 1, pmod<int64_t>(INT64_MIN, INT64_MAX));
 }
 
+TEST_F(PmodTest, float) {
+  EXPECT_FLOAT_EQ(0.2, pmod<float>(0.5, 0.3).value());
+  EXPECT_FLOAT_EQ(0.9, pmod<float>(-1.1, 2).value());
+}
+
+TEST_F(PmodTest, double) {
+  EXPECT_DOUBLE_EQ(0.2, pmod<double>(0.5, 0.3).value());
+  EXPECT_DOUBLE_EQ(0.9, pmod<double>(-1.1, 2).value());
+}
+
 class RemainderTest : public SparkFunctionBaseTest {
  protected:
   template <typename T>