Support week year

velox/functions/lib/DateTimeFormatter.cpp

@@ -24,6 +24,7 @@
 #include "velox/external/date/iso_week.h"
 #include "velox/external/date/tz.h"
 #include "velox/functions/lib/DateTimeFormatterBuilder.h"
+#include "velox/functions/lib/TimeUtils.h"
 #include "velox/type/TimestampConversion.h"
 #include "velox/type/tz/TimeZoneMap.h"
 
@@ -1013,6 +1014,7 @@ uint32_t DateTimeFormatter::maxResultSize(const tz::TimeZone* timezone) const {
         size += 2;
         break;
       case DateTimeFormatSpecifier::YEAR_OF_ERA:
+      case DateTimeFormatSpecifier::WEEK_YEAR:
         // Timestamp is in [-32767-01-01, 32767-12-31] range.
         size += std::max((int)token.pattern.minRepresentDigits, 6);
         break;
@@ -1067,7 +1069,6 @@ uint32_t DateTimeFormatter::maxResultSize(const tz::TimeZone* timezone) const {
         size += 9;
         break;
       // Not supported.
-      case DateTimeFormatSpecifier::WEEK_YEAR:
       default:
         VELOX_UNSUPPORTED(
             "Date format specifier is not supported: {}",
@@ -1328,7 +1329,23 @@ int32_t DateTimeFormatter::format(
               result);
           break;
         }
-        case DateTimeFormatSpecifier::WEEK_YEAR:
+        case DateTimeFormatSpecifier::WEEK_YEAR: {
+          auto year = getWeekYear(
+              static_cast<int>(calDate.year()),
+              static_cast<uint32_t>(calDate.month()),
+              static_cast<uint32_t>(calDate.day()),
+              2, // (ISO 8601) Monday = 2
+              4 // At least 4 days in first week
+          );
+
+          result += padContent(
+              static_cast<signed>(year),
+              '0',
+              token.pattern.minRepresentDigits,
+              maxResultEnd,
+              result);
+          break;
+        }
         default:
           VELOX_UNSUPPORTED(
               "format is not supported for specifier {}",

velox/functions/lib/TimeUtils.h

@@ -17,7 +17,9 @@
 
 #include <velox/type/Timestamp.h>
 #include "velox/core/QueryConfig.h"
+#include "velox/external/date/date.h"
 #include "velox/functions/Macros.h"
+#include "velox/type/TimestampConversion.h"
 #include "velox/type/tz/TimeZoneMap.h"
 
 namespace facebook::velox::functions {
@@ -104,4 +106,77 @@ struct InitSessionTimezone {
     timeZone_ = getTimeZoneFromConfig(config);
   }
 };
+
+// Return the first `fisrtDayOfWeek` in the year.
+// getFirstDayOfWeek() use 1-based weekday number:
+//   Sunday - 1, Monday - 2, ..., Saturday - 7
+FOLLY_ALWAYS_INLINE
+uint32_t getFirstDayOfWeek(int32_t y, uint32_t firstDayOfWeek) {
+  auto firstDay =
+      date::year_month_day(date::year(y), date::month(1), date::day(1));
+  auto weekday = date::weekday(firstDay).c_encoding() + 1;
+
+  int32_t delta = firstDayOfWeek - weekday;
+  if (delta < 0)
+    delta += 7;
+
+  return delta + 1;
+}
+
+// Return week-based-year of year month day.
+// The week containing Jan 1 has `minimalDaysInFirstWeek` or more days is
+// week 1. The week is defined via `firstDayOfWeek`. `firstDayOfWeek` is 1-based
+// weekday number starting with Sunday.
+//
+// For ISO 8601, `firstDayOfWeek` is 2 (Monday) and `minimalDaysInFirstWeek`
+// is 4. For Legacy Spark, `firstDayOfWeek` is 1 (Sunday) and
+// `minimalDaysInFisrtWeek` is 1 by default.
+//
+// getWeekYear only works with gregorian calendar due to limitations in the date
+// library. As a result, dates before the gregorian cutover would yield
+// mismatched results.
+//
+// The algorithm refers to the weekyear algorithm in jdk:
+// https://github.com/openjdk/jdk8/blob/6a383433a9f4661a96a90b2a4c7b5b9a85720031/jdk/src/share/classes/java/util/GregorianCalendar.java#L2077
+FOLLY_ALWAYS_INLINE
+int32_t getWeekYear(
+    int32_t y,
+    uint32_t m,
+    uint32_t d,
+    uint32_t firstDayOfWeek,
+    uint32_t minimalDaysInFirstWeek) {
+  auto year = y;
+  auto calDate =
+      date::year_month_day(date::year(y), date::month(m), date::day(d));
+  auto weekday = date::weekday(calDate).c_encoding();
+  auto firstDayOfTheYear =
+      date::year_month_day(calDate.year(), date::month(1), date::day(1));
+  auto dayOfYear =
+      (date::sys_days{calDate} - date::sys_days{firstDayOfTheYear}).count() + 1;
+  auto maxDayOfYear = util::isLeapYear(y) ? 366 : 365;
+
+  if (dayOfYear > minimalDaysInFirstWeek && dayOfYear < (maxDayOfYear - 6)) {
+    return year;
+  }
+
+  auto minDayOfYear = getFirstDayOfWeek(y, firstDayOfWeek);
+  if (dayOfYear < minDayOfYear) {
+    if (minDayOfYear <= minimalDaysInFirstWeek) {
+      --year;
+    }
+  } else {
+    auto minDayOfYear = getFirstDayOfWeek(y + 1, firstDayOfWeek) - 1;
+    if (minDayOfYear == 0) {
+      minDayOfYear = 7;
+    }
+    if (minDayOfYear >= minimalDaysInFirstWeek) {
+      int days = maxDayOfYear - dayOfYear + 1;
+      if (days <= (7 - minDayOfYear)) {
+        ++year;
+      }
+    }
+  }
+
+  return year;
+}
 } // namespace facebook::velox::functions

velox/functions/lib/tests/CMakeLists.txt

@@ -24,6 +24,7 @@ add_executable(
   Re2FunctionsTest.cpp
   RepeatTest.cpp
   Utf8Test.cpp
+  TimeUtilsTest.cpp
   ZetaDistributionTest.cpp)
 
 add_test(

velox/functions/lib/tests/DateTimeFormatterTest.cpp

@@ -1343,6 +1343,26 @@ TEST_F(JodaDateTimeFormatterTest, betterErrorMessaging) {
       "Value 429 for dayOfMonth must be in the range [1,365] for year 2057 and month 2.");
 }
 
+TEST_F(JodaDateTimeFormatterTest, formatWeekYear) {
+  DateTimeFormatterBuilder builder(10);
+  auto formatter =
+      builder.appendWeekYear(4).setType(DateTimeFormatterType::JODA).build();
+  auto* timezone = tz::locateZone("GMT");
+  const auto maxSize = formatter->maxResultSize(timezone);
+
+  auto weekYear = [&](const StringView& time) {
+    std::string result(maxSize, '\0');
+    auto resultSize = formatter->format(
+        fromTimestampString(time), timezone, maxSize, result.data());
+    result.resize(resultSize);
+    return result;
+  };
+
+  EXPECT_EQ(weekYear("2019-12-31 00:00:00"), "2020");
+  EXPECT_EQ(weekYear("2020-12-26 00:00:00"), "2020");
+  EXPECT_EQ(weekYear("2021-01-01 00:00:00"), "2020");
+}
+
 class MysqlDateTimeTest : public DateTimeFormatterTest {};
 
 TEST_F(MysqlDateTimeTest, validBuild) {