update

facebookincubator · Oct 15, 2024 · 8ee67ed · 8ee67ed
1 parent 1ee8e0b
commit 8ee67ed
Showing 1 changed file with 58 additions and 26 deletions.
diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -246,43 +246,59 @@ struct FromUnixtimeFunction {
     legacyFormatter_ = config.sparkLegacyDateFormatter();
     sessionTimeZone_ = getTimeZoneFromConfig(config);
     if (format != nullptr) {
-      setFormatter(*format);
+      if (!setFormatter(*format)) {
+        invalidFormat_ = true;
+      }
       isConstantTimeFormat_ = true;
     }
   }
 
-  FOLLY_ALWAYS_INLINE void call(
+  FOLLY_ALWAYS_INLINE bool call(
       out_type<Varchar>& result,
       const arg_type<int64_t>& second,
       const arg_type<Varchar>& format) {
+    if (invalidFormat_) {
+      return false;
+    }
     if (!isConstantTimeFormat_) {
-      setFormatter(format);
+      if (!setFormatter(format)) {
+        return false;
+      }
     }
     const Timestamp timestamp{second, 0};
     result.reserve(maxResultSize_);
     int32_t resultSize;
     resultSize = formatter_->format(
         timestamp, sessionTimeZone_, maxResultSize_, result.data(), true);
     result.resize(resultSize);
+    return true;
   }
 
  private:
-  FOLLY_ALWAYS_INLINE void setFormatter(const arg_type<Varchar>& format) {
-    formatter_ = detail::getDateTimeFormatter(
-                     std::string_view(format.data(), format.size()),
-                     legacyFormatter_ ? DateTimeFormatterType::STRICT_SIMPLE
-                                      : DateTimeFormatterType::JODA)
-                     .thenOrThrow(folly::identity, [&](const Status& status) {
-                       VELOX_USER_FAIL("{}", status.message());
-                     });
+  FOLLY_ALWAYS_INLINE bool setFormatter(const arg_type<Varchar>& format) {
+    auto formatter = detail::getDateTimeFormatter(
+        std::string_view(format.data(), format.size()),
+        legacyFormatter_ ? DateTimeFormatterType::STRICT_SIMPLE
+                         : DateTimeFormatterType::JODA);
+    // Legacy formatter returns null for invalid format but Joda formatter
+    // throws user error.
+    if (!legacyFormatter_ && formatter.hasError()) {
+      VELOX_USER_FAIL(formatter.error().message());
+    }
+    if (formatter.hasError()) {
+      return false;
+    }
+    formatter_ = formatter.value();
     maxResultSize_ = formatter_->maxResultSize(sessionTimeZone_);
+    return true;
   }
 
   const tz::TimeZone* sessionTimeZone_{nullptr};
   std::shared_ptr<DateTimeFormatter> formatter_;
   uint32_t maxResultSize_;
   bool isConstantTimeFormat_{false};
   bool legacyFormatter_{false};
+  bool invalidFormat_{false};
 };
 
 template <typename T>
@@ -362,29 +378,44 @@ struct GetTimestampFunction {
       sessionTimeZone_ = tz::locateZone(sessionTimezoneName);
     }
     if (format != nullptr) {
-      formatter_ = detail::getDateTimeFormatter(
-                       std::string_view(*format),
-                       legacyFormatter_ ? DateTimeFormatterType::STRICT_SIMPLE
-                                        : DateTimeFormatterType::JODA)
-                       .thenOrThrow(folly::identity, [&](const Status& status) {
-                         VELOX_USER_FAIL("{}", status.message());
-                       });
-      isConstantTimeFormat_ = true;
+      auto formatter = detail::getDateTimeFormatter(
+          std::string_view(*format),
+          legacyFormatter_ ? DateTimeFormatterType::STRICT_SIMPLE
+                           : DateTimeFormatterType::JODA);
+      // Legacy formatter returns null for invalid format but Joda formatter
+      // throws user error.
+      if (!legacyFormatter_ && formatter.hasError()) {
+        VELOX_USER_FAIL(formatter.error().message());
+      }
+      if (formatter.hasError()) {
+        invalidFormat_ = true;
+      } else {
+        formatter_ = formatter.value();
+      }
     }
   }
 
   FOLLY_ALWAYS_INLINE bool call(
       out_type<Timestamp>& result,
       const arg_type<Varchar>& input,
       const arg_type<Varchar>& format) {
+    if (invalidFormat_) {
+      return false;
+    }
     if (!isConstantTimeFormat_) {
-      formatter_ = detail::getDateTimeFormatter(
-                       std::string_view(format),
-                       legacyFormatter_ ? DateTimeFormatterType::STRICT_SIMPLE
-                                        : DateTimeFormatterType::JODA)
-                       .thenOrThrow(folly::identity, [&](const Status& status) {
-                         VELOX_USER_FAIL("{}", status.message());
-                       });
+      auto formatter = detail::getDateTimeFormatter(
+          std::string_view(format),
+          legacyFormatter_ ? DateTimeFormatterType::STRICT_SIMPLE
+                           : DateTimeFormatterType::JODA);
+      // Legacy formatter returns null for invalid format but Joda formatter
+      // throws user error.
+      if (!legacyFormatter_ && formatter.hasError()) {
+        VELOX_USER_FAIL(formatter.error().message());
+      }
+      if (formatter.hasError()) {
+        return false;
+      }
+      formatter_ = formatter.value();
     }
     auto dateTimeResult = formatter_->parse(std::string_view(input));
     // Null as result for parsing error.
@@ -408,6 +439,7 @@ struct GetTimestampFunction {
   bool isConstantTimeFormat_{false};
   const tz::TimeZone* sessionTimeZone_{tz::locateZone(0)}; // default to GMT.
   bool legacyFormatter_{false};
+  bool invalidFormat_{false};
 };
 
 template <typename T>