copy week_of_year from prestosql, but change return type to integer

velox/docs/functions/spark/datetime.rst

@@ -22,6 +22,13 @@ These functions support TIMESTAMP and DATE input types.
     ``day`` need to be from 1 to 31, and matches the number of days in each month.
     days of ``year-month-day - 1970-01-01`` need to be in the range of INTEGER type.
 
+.. function:: week(x) -> integer
+
+    Returns the `ISO-Week`_ of the year from x. The value ranges from ``1`` to ``53``.
+    A week is considered to start on a Monday and week 1 is the first week with >3 days.
+
+.. _ISO-Week: https://en.wikipedia.org/wiki/ISO_week_date
+
 .. spark:function:: to_unix_timestamp(string) -> integer
 
     Alias for ``unix_timestamp(string) -> integer``.

velox/functions/sparksql/DateTimeFunctions.h

@@ -39,6 +39,60 @@ struct YearFunction : public InitSessionTimezone<T> {
   }
 };
 
+template <typename T>
+struct WeekFunction : public InitSessionTimezone<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE int32_t getWeek(const std::tm& time) {
+    // The computation of ISO week from date follows the algorithm here:
+    // https://en.wikipedia.org/wiki/ISO_week_date
+    int32_t week = floor(
+                       10 + (time.tm_yday + 1) -
+                       (time.tm_wday ? time.tm_wday : kDaysInWeek)) /
+        kDaysInWeek;
+
+    if (week == 0) {
+      // Distance in days between the first day of the current year and the
+      // Monday of the current week.
+      auto mondayOfWeek =
+          time.tm_yday + 1 - (time.tm_wday + kDaysInWeek - 1) % kDaysInWeek;
+      // Distance in days between the first day and the first Monday of the
+      // current year.
+      auto firstMondayOfYear =
+          1 + (mondayOfWeek + kDaysInWeek - 1) % kDaysInWeek;
+
+      if ((util::isLeapYear(time.tm_year + 1900 - 1) &&
+           firstMondayOfYear == 2) ||
+          firstMondayOfYear == 3 || firstMondayOfYear == 4) {
+        week = 53;
+      } else {
+        week = 52;
+      }
+    } else if (week == 53) {
+      // Distance in days between the first day of the current year and the
+      // Monday of the current week.
+      auto mondayOfWeek =
+          time.tm_yday + 1 - (time.tm_wday + kDaysInWeek - 1) % kDaysInWeek;
+      auto daysInYear = util::isLeapYear(time.tm_year + 1900) ? 366 : 365;
+      if (daysInYear - mondayOfWeek < 3) {
+        week = 1;
+      }
+    }
+
+    return week;
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<Timestamp>& timestamp) {
+    result = getWeek(getDateTime(timestamp, this->timeZone_));
+  }
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = getWeek(getDateTime(date));
+  }
+};
+
 template <typename T>
 struct UnixTimestampFunction {
   // unix_timestamp();

velox/functions/sparksql/Register.cpp

@@ -189,6 +189,8 @@ void registerFunctions(const std::string& prefix) {
   // Register date functions.
   registerFunction<YearFunction, int32_t, Timestamp>({prefix + "year"});
   registerFunction<YearFunction, int32_t, Date>({prefix + "year"});
+  registerFunction<WeekFunction, int32_t, Timestamp>({prefix + "week_of_year"});
+  registerFunction<WeekFunction, int32_t, Date>({prefix + "week_of_year"});
 
   registerFunction<UnixTimestampFunction, int64_t>({prefix + "unix_timestamp"});
 

velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp

@@ -30,6 +30,10 @@ class DateTimeFunctionsTest : public SparkFunctionBaseTest {
         {core::QueryConfig::kAdjustTimestampToTimezone, "true"},
     });
   }
+
+  int32_t parseDate(const std::string& dateStr) {
+    return DATE()->toDays(dateStr);
+  }
 };
 
 TEST_F(DateTimeFunctionsTest, year) {
@@ -66,6 +70,26 @@ TEST_F(DateTimeFunctionsTest, yearDate) {
   EXPECT_EQ(1920, year(DATE()->toDays("1920-01-01")));
 }
 
+TEST_F(DateTimeFunctionsTest, weekDate) {
+  const auto weekDate = [&](const char* dateString) {
+    auto date = std::make_optional(parseDate(dateString));
+    return evaluateOnce<int32_t, int32_t>("week_of_year(c0)", {date}, {DATE()})
+        .value();
+  };
+
+  EXPECT_EQ(1, weekDate("1919-12-31"));
+  EXPECT_EQ(1, weekDate("1920-01-01"));
+  EXPECT_EQ(1, weekDate("1920-01-04"));
+  EXPECT_EQ(2, weekDate("1920-01-05"));
+  EXPECT_EQ(53, weekDate("1960-01-01"));
+  EXPECT_EQ(53, weekDate("1960-01-03"));
+  EXPECT_EQ(1, weekDate("1960-01-04"));
+  EXPECT_EQ(1, weekDate("1969-12-31"));
+  EXPECT_EQ(1, weekDate("1970-01-01"));
+  EXPECT_EQ(1, weekDate("0001-01-01"));
+  EXPECT_EQ(52, weekDate("9999-12-31"));
+}
+
 TEST_F(DateTimeFunctionsTest, unixTimestamp) {
   const auto unixTimestamp = [&](std::optional<StringView> dateStr) {
     return evaluateOnce<int64_t>("unix_timestamp(c0)", dateStr);