Initial commit

velox/functions/sparksql/DateTimeFunctions.h

@@ -98,6 +98,36 @@ struct WeekFunction : public InitSessionTimezone<T> {
   }
 };
 
+template <typename T>
+struct YearOfWeekSparkFunction : public InitSessionTimezone<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE int32_t computeYearOfWeek(const std::tm& dateTime) {
+    int isoWeekDay = dateTime.tm_wday == 0 ? 7 : dateTime.tm_wday;
+    // The last few days in December may belong to the next year if they are
+    // in the same week as the next January 1 and this January 1 is a Thursday
+    // or before.
+    if (UNLIKELY(
+            dateTime.tm_mon == 11 && dateTime.tm_mday >= 29 &&
+            dateTime.tm_mday - isoWeekDay >= 31 - 3)) {
+      return 1900 + dateTime.tm_year + 1;
+    }
+    // The first few days in January may belong to the last year if they are
+    // in the same week as January 1 and January 1 is a Friday or after.
+    else if (UNLIKELY(
+                 dateTime.tm_mon == 0 && dateTime.tm_mday <= 3 &&
+                 isoWeekDay - (dateTime.tm_mday - 1) >= 5)) {
+      return 1900 + dateTime.tm_year - 1;
+    } else {
+      return 1900 + dateTime.tm_year;
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE void call(int32_t& result, const arg_type<Date>& date) {
+    result = computeYearOfWeek(getDateTime(date));
+  }
+};
+
 template <typename T>
 struct UnixDateFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);

velox/functions/sparksql/Register.cpp

@@ -324,6 +324,8 @@ void registerFunctions(const std::string& prefix) {
   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<YearOfWeekSparkFunction, int32_t, Date>(
+      {prefix + "spark_year_of_week"});
 
   registerFunction<ToUtcTimestampFunction, Timestamp, Timestamp, Varchar>(
       {prefix + "to_utc_timestamp"});