diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h index 0ba287089e4a..ddd7d723685a 100644 --- a/velox/functions/sparksql/DateTimeFunctions.h +++ b/velox/functions/sparksql/DateTimeFunctions.h @@ -98,6 +98,36 @@ struct WeekFunction : public InitSessionTimezone { } }; +template +struct YearOfWeekSparkFunction : public InitSessionTimezone { + 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) { + result = computeYearOfWeek(getDateTime(date)); + } +}; + template struct UnixDateFunction { VELOX_DEFINE_FUNCTION_TYPES(T); diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp index fa44e98c27fa..10f70a2c920d 100644 --- a/velox/functions/sparksql/Register.cpp +++ b/velox/functions/sparksql/Register.cpp @@ -324,6 +324,8 @@ void registerFunctions(const std::string& prefix) { registerFunction({prefix + "year"}); registerFunction({prefix + "week_of_year"}); registerFunction({prefix + "week_of_year"}); + registerFunction( + {prefix + "spark_year_of_week"}); registerFunction( {prefix + "to_utc_timestamp"});