From 6dab43a79b1c76f07ff11192510bb671902427f3 Mon Sep 17 00:00:00 2001
From: rui-mo <rui.mo@intel.com>
Date: Wed, 18 Oct 2023 14:52:42 +0800
Subject: [PATCH] Add dayofyear function support in sparksql (7116)

---
 velox/functions/sparksql/DateTimeFunctions.h  | 27 +++++++++++++++++++
 velox/functions/sparksql/Register.cpp         |  4 +--
 .../sparksql/tests/DateTimeFunctionsTest.cpp  | 25 ++++++++++++++++-
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
index 032fa402cfe3..dcb6564439f0 100644
--- a/velox/functions/sparksql/DateTimeFunctions.h
+++ b/velox/functions/sparksql/DateTimeFunctions.h
@@ -450,4 +450,31 @@ struct DayFunction : public InitSessionTimezone<T>,
   }
 };
 
+template <typename T>
+struct DayOfYearFunction : public InitSessionTimezone<T>,
+                           public TimestampWithTimezoneSupport<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE int32_t getDayOfYear(const std::tm& time) {
+    return time.tm_yday + 1;
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<Timestamp>& timestamp) {
+    result = getDayOfYear(getDateTime(timestamp, this->timeZone_));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getDayOfYear(getDateTime(date));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
+    auto timestamp = this->toTimestamp(timestampWithTimezone);
+    result = getDayOfYear(getDateTime(timestamp, nullptr));
+  }
+};
+
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
index 223aca5045c7..fa3a7902f7ba 100644
--- a/velox/functions/sparksql/Register.cpp
+++ b/velox/functions/sparksql/Register.cpp
@@ -279,9 +279,9 @@ void registerFunctions(const std::string& prefix) {
       {prefix + "day", prefix + "dayofmonth"});
   registerFunction<DayFunction, int32_t, Date>(
       {prefix + "day", prefix + "dayofmonth"});
-  registerFunction<DayOfYearFunction, int64_t, Timestamp>(
+  registerFunction<DayOfYearFunction, int32_t, Timestamp>(
       {prefix + "doy", prefix + "dayofyear"});
-  registerFunction<DayOfYearFunction, int64_t, Date>(
+  registerFunction<DayOfYearFunction, int32_t, Date>(
       {prefix + "doy", prefix + "dayofyear"});
 
   registerFunction<DayOfWeekFunction, int32_t, Timestamp>(
diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
index 45de8fd65fae..7f625b049a65 100644
--- a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
+++ b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
@@ -291,7 +291,7 @@ TEST_F(DateTimeFunctionsTest, dateSub) {
 
 TEST_F(DateTimeFunctionsTest, dayOfYear) {
   const auto day = [&](std::optional<int32_t> date) {
-    return evaluateOnce<int64_t, int32_t>("dayofyear(c0)", {date}, {DATE()});
+    return evaluateOnce<int32_t, int32_t>("dayofyear(c0)", {date}, {DATE()});
   };
   EXPECT_EQ(std::nullopt, day(std::nullopt));
   EXPECT_EQ(100, day(parseDate("2016-04-09")));
@@ -299,6 +299,29 @@ TEST_F(DateTimeFunctionsTest, dayOfYear) {
   EXPECT_EQ(1, day(parseDate("1970-01-01")));
 }
 
+TEST_F(DateTimeFunctionsTest, dayOfYearTimestamp) {
+  const auto day = [&](std::optional<Timestamp> date) {
+    return evaluateOnce<int32_t>("dayofyear(c0)", date);
+  };
+  EXPECT_EQ(std::nullopt, day(std::nullopt));
+  EXPECT_EQ(1, day(Timestamp(0, 0)));
+  EXPECT_EQ(365, day(Timestamp(-1, 9000)));
+  EXPECT_EQ(273, day(Timestamp(1632989700, 0)));
+  EXPECT_EQ(274, day(Timestamp(1633076100, 0)));
+  EXPECT_EQ(279, day(Timestamp(1633508100, 0)));
+  EXPECT_EQ(304, day(Timestamp(1635668100, 0)));
+
+  setQueryTimeZone("Pacific/Apia");
+
+  EXPECT_EQ(std::nullopt, day(std::nullopt));
+  EXPECT_EQ(365, day(Timestamp(0, 0)));
+  EXPECT_EQ(365, day(Timestamp(-1, 9000)));
+  EXPECT_EQ(273, day(Timestamp(1632989700, 0)));
+  EXPECT_EQ(274, day(Timestamp(1633076100, 0)));
+  EXPECT_EQ(279, day(Timestamp(1633508100, 0)));
+  EXPECT_EQ(304, day(Timestamp(1635668100, 0)));
+}
+
 TEST_F(DateTimeFunctionsTest, dayOfMonth) {
   const auto day = [&](std::optional<int32_t> date) {
     return evaluateOnce<int32_t, int32_t>("dayofmonth(c0)", {date}, {DATE()});