implement to_iso8601() for timestamp input

facebookincubator · Mar 5, 2024 · 8404d66 · 8404d66
1 parent 42b10d9
commit 8404d66
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 2 deletions.
diff --git a/velox/docs/functions/presto/datetime.rst b/velox/docs/functions/presto/datetime.rst
@@ -106,6 +106,9 @@ Date and Time Functions
 
     Returns ``timestamp`` as a UNIX timestamp.
 
+.. function:: to_iso8601(x) -> varchar
+
+    Formats timestamp ``x`` as an ISO 8601 string.
 Truncation Function
 -------------------
 

diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -15,7 +15,13 @@
  */
 #pragma once
 
+#include <time.h>
+#include <chrono>
+#include <cmath>
+#include <string>
 #include <string_view>
+#include "velox/external/date/date.h"
+#include "velox/external/date/tz.h"
 #include "velox/functions/lib/DateTimeFormatter.h"
 #include "velox/functions/lib/TimeUtils.h"
 #include "velox/functions/prestosql/DateTimeImpl.h"
@@ -1399,4 +1405,41 @@ struct TimeZoneMinuteFunction : public TimestampWithTimezoneSupport<T> {
   }
 };
 
+template <typename T>
+struct ToISO8601FromTimestampFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  std::string sessionTzName = "";
+  const date::time_zone* sessionTimezone = nullptr;
+
+  FOLLY_ALWAYS_INLINE void initialize(
+      const core::QueryConfig& config,
+      const arg_type<Timestamp>* timestamp) {
+    sessionTimezone = getTimeZoneFromConfig(config);
+    sessionTzName = (sessionTimezone != NULL ? sessionTimezone->name() : NULL);
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<Varchar>& result,
+      const arg_type<Timestamp>& timestamp) {
+    std::string tzOffsetStr = "";
+    Timestamp ts = timestamp;
+    TimestampToStringOptions options;
+    options.mode = TimestampToStringOptions::Mode::kFull;
+    options.zeroPaddingYear = true;
+    options.precision = TimestampToStringOptions::Precision::kMilliseconds;
+    auto tsStr = ts.toString(options);
+
+    if (sessionTzName == "UTC") {
+      tzOffsetStr = " UTC";
+    } else if (sessionTimezone != NULL) {
+      auto formatStr = "%Ez";
+      date::local_time<std::chrono::nanoseconds> tp{
+          std::chrono::nanoseconds{timestamp.toNanos()}};
+      tzOffsetStr = format(formatStr, date::zoned_time{sessionTzName, tp});
+    }
+    result = tsStr + tzOffsetStr;
+  }
+};
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp
@@ -175,6 +175,8 @@ void registerSimpleFunctions(const std::string& prefix) {
   registerFunction<DateParseFunction, Timestamp, Varchar, Varchar>(
       {prefix + "date_parse"});
   registerFunction<CurrentDateFunction, Date>({prefix + "current_date"});
+  registerFunction<ToISO8601FromTimestampFunction, Varchar, Timestamp>(
+      {prefix + "to_iso8601"});
 }
 } // namespace
 

diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
@@ -3836,3 +3836,50 @@ TEST_F(DateTimeFunctionsTest, fromUnixtimeDouble) {
   });
   assertEqualVectors(expected, actual);
 }
+
+TEST_F(DateTimeFunctionsTest, toISO8601TestTimestamp) {
+  const auto toISO8601 = [&](std::optional<Timestamp> timestamp) {
+    return evaluateOnce<std::string>("to_iso8601(c0)", timestamp);
+  };
+
+  auto getTimestamp = [](const std::string& timestampStr) {
+    return util::fromTimestampString(StringView{timestampStr});
+  };
+
+  setQueryTimeZone("UTC");
+
+  EXPECT_EQ(std::nullopt, toISO8601(std::nullopt));
+
+  EXPECT_EQ(
+      "2017-07-17T19:54:57.828 UTC",
+      toISO8601(Timestamp(1500321297, 827910000)));
+
+  setQueryTimeZone("America/Denver");
+
+  EXPECT_EQ(
+      "2017-07-17T19:54:57.828-06:00",
+      toISO8601(Timestamp(1500321297, 827910000)));
+
+  EXPECT_EQ(
+      "2017-07-17T19:54:57.820-06:00",
+      toISO8601(Timestamp(1500321297, 820000000)));
+
+  setQueryTimeZone("Australia/Adelaide");
+
+  EXPECT_EQ(
+      "2017-07-17T19:54:57.127+09:30",
+      toISO8601(Timestamp(1500321297, 126548219)));
+
+  EXPECT_EQ(
+      "2017-07-17T19:54:57.095+09:30",
+      toISO8601(Timestamp(1500321297, 95421806)));
+
+  EXPECT_EQ(
+      "1970-01-01T12:00:12.002+09:30", toISO8601(Timestamp(43212, 1599247)));
+
+  setQueryTimeZone("America/Los_Angeles");
+
+  EXPECT_EQ(
+      "1880-10-30T06:21:47.001-07:52",
+      toISO8601(Timestamp(-2813938693, 818583)));
+}
diff --git a/velox/type/Timestamp.cpp b/velox/type/Timestamp.cpp
@@ -255,6 +255,8 @@ std::string Timestamp::tmToString(
     const TimestampToStringOptions& options) {
   VELOX_DCHECK_GE(nanos, 0);
   VELOX_DCHECK_LT(nanos, 1'000'000'000);
+  double kNanosecondsInMillisecond = 1'000'000;
+  double kNanosecondsInMicrosecond = 1'000;
   auto precisionWidth = static_cast<int8_t>(options.precision);
   std::string out;
   out.reserve(getCapacity(options));
@@ -298,10 +300,10 @@ std::string Timestamp::tmToString(
   out += ':';
   appendSmallInt(tmValue.tm_sec, out);
   if (options.precision == TimestampToStringOptions::Precision::kMilliseconds) {
-    nanos /= 1'000'000;
+    nanos = round(nanos / kNanosecondsInMillisecond);
   } else if (
       options.precision == TimestampToStringOptions::Precision::kMicroseconds) {
-    nanos /= 1'000;
+    nanos = round(nanos / kNanosecondsInMicrosecond);
   }
   if (options.skipTrailingZeros && nanos == 0) {
     return out;