fix: Support legacy date formatter for Spark date_format function

Author: zml1206

velox/functions/lib/DateTimeFormatter.cpp

@@ -2057,6 +2057,9 @@ Expected<std::shared_ptr<DateTimeFormatter>> buildSimpleDateTimeFormatter(
         case 'S':
           builder.appendFractionOfSecond(count);
           break;
+        case 'u':
+          builder.appendDayOfWeek1Based(count);
+          break;
         case 'w':
           builder.appendWeekOfWeekYear(count);
           break;

velox/functions/sparksql/DateTimeFunctions.h

@@ -1141,4 +1141,70 @@ struct MonthsBetweenFunction {
   const tz::TimeZone* sessionTimeZone_ = nullptr;
 };
 
+template <typename T>
+struct DateFormatFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config,
+      const arg_type<Timestamp>* /*timestamp*/,
+      const arg_type<Varchar>* formatString) {
+    legacyFormatter_ = config.sparkLegacyDateFormatter();
+    sessionTimeZone_ = getTimeZoneFromConfig(config);
+    if (formatString != nullptr) {
+      auto formatter = detail::initializeFormatter(
+          std::string_view(*formatString), legacyFormatter_);
+      if (formatter) {
+        formatter_ = formatter;
+        maxResultSize_ = formatter_->maxResultSize(sessionTimeZone_);
+      } else {
+        invalidFormat_ = true;
+      }
+      isConstFormat_ = true;
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Varchar>& result,
+      const arg_type<Timestamp>& timestamp,
+      const arg_type<Varchar>& formatString) {
+    if (invalidFormat_) {
+      return false;
+    }
+    if (!isConstFormat_) {
+      auto formatter = detail::initializeFormatter(
+          std::string_view(formatString), legacyFormatter_);
+      if (formatter) {
+        formatter_ = formatter;
+        maxResultSize_ = formatter_->maxResultSize(sessionTimeZone_);
+      } else {
+        return false;
+      }
+    }
+
+    format(timestamp, sessionTimeZone_, maxResultSize_, result);
+    return true;
+  }
+
+ private:
+  FOLLY_ALWAYS_INLINE void format(
+      const Timestamp& timestamp,
+      const tz::TimeZone* timeZone,
+      uint32_t maxResultSize,
+      out_type<Varchar>& result) const {
+    result.reserve(maxResultSize);
+    const auto resultSize =
+        formatter_->format(timestamp, timeZone, maxResultSize, result.data());
+    result.resize(resultSize);
+  }
+
+  const tz::TimeZone* sessionTimeZone_{nullptr};
+  std::shared_ptr<DateTimeFormatter> formatter_;
+  bool isConstFormat_{false};
+  bool legacyFormatter_{false};
+  bool invalidFormat_{false};
+  uint32_t maxResultSize_;
+};
+
 } // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/registration/RegisterDatetime.cpp

@@ -56,7 +56,7 @@ void registerDatetimeFunctions(const std::string& prefix) {
   registerFunction<DateAddFunction, Date, Date, int8_t>({prefix + "date_add"});
   registerFunction<DateAddFunction, Date, Date, int16_t>({prefix + "date_add"});
   registerFunction<DateAddFunction, Date, Date, int32_t>({prefix + "date_add"});
-  registerFunction<FormatDateTimeFunction, Varchar, Timestamp, Varchar>(
+  registerFunction<DateFormatFunction, Varchar, Timestamp, Varchar>(
       {prefix + "date_format"});
   registerFunction<DateFromUnixDateFunction, Date, int32_t>(
       {prefix + "date_from_unix_date"});

velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp

@@ -1858,5 +1858,30 @@ TEST_F(DateTimeFunctionsTest, monthsBetween) {
           parseTimestamp("1996-03-21 11:00:00"),
           true));
 }
+TEST_F(DateTimeFunctionsTest, dateFormat) {
+  const auto dateFormat = [&](std::optional<Timestamp> timestamp,
+                              std::optional<std::string> format) {
+    return evaluateOnce<std::string>("date_format(c0, c1)", timestamp, format);
+  };
+
+  // Check invalid format.
+  EXPECT_THROW(dateFormat(parseTimestamp("1970-01-01"), "u"), VeloxUserError);
+  EXPECT_THROW(
+      dateFormat(parseTimestamp("1970-01-01"), "'abcd"), VeloxUserError);
+
+  // Check Simple tests.
+  EXPECT_EQ("AD", dateFormat(parseTimestamp("1970-01-01"), "G"));
+  EXPECT_EQ("19", dateFormat(parseTimestamp("1900-01-01"), "C"));
+  EXPECT_EQ("2020", dateFormat(parseTimestamp("2020-01-01"), "Y"));
+  EXPECT_EQ("1", dateFormat(parseTimestamp("2022-01-01"), "D"));
+  EXPECT_EQ("1", dateFormat(parseTimestamp("2022-01-01"), "d"));
+  EXPECT_EQ("AM", dateFormat(parseTimestamp("2022-01-01 00:00:00"), "a"));
+  EXPECT_EQ(
+      "2022-01-01 00:00:00",
+      dateFormat(parseTimestamp("2022-01-01"), "yyyy-MM-dd HH:mm:ss"));
+
+  enableLegacyFormatter();
+  EXPECT_EQ("4", dateFormat(parseTimestamp("1970-01-01"), "u"));
+}
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test