implement cast from varchar to TimestampWithTimezone

svm1 · Apr 29, 2024 · 5fb0d3a · 5fb0d3a
1 parent 4a0383b
commit 5fb0d3a
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 4 deletions.
diff --git a/velox/functions/prestosql/tests/TimestampWithTimeZoneCastTest.cpp b/velox/functions/prestosql/tests/TimestampWithTimeZoneCastTest.cpp
@@ -16,6 +16,7 @@
 
 #include "velox/functions/prestosql/tests/CastBaseTest.h"
 #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
+#include "velox/type/tz/TimeZoneMap.h"
 
 using namespace facebook::velox;
 
@@ -64,6 +65,74 @@ TEST_F(TimestampWithTimeZoneCastTest, fromTimestamp) {
   testCast(tsVector, expected);
 }
 
+TEST_F(TimestampWithTimeZoneCastTest, fromVarchar) {
+  const auto stringVector = makeNullableFlatVector<StringView>(
+      {std::nullopt,
+       "2012-10-31 01:00:47 America/Denver",
+       "2012-10-31 01:00:47 -06:00",
+       "1994-05-06 15:49 Europe/Vienna",
+       "1979-02-24 08:33:31 Pacific/Chatham",
+       "1979-02-24 08:33:31 +13:45"});
+
+  // Varchar representations above hold local time (relative to specified time
+  // zone). For instance, the first string represents a wall-clock displaying
+  // '2012-10-31 01:00:47' in Denver. Below, we use UTC representations of the
+  // above local wall-clocks, to match the UTC timepoints held in the
+  // TimestampWithTimezone type.
+  auto denverUTC = util::fromTimestampString("2012-10-31 07:00:47").toMillis();
+  auto viennaUTC = util::fromTimestampString("1994-05-06 13:49:00").toMillis();
+  auto chathamUTC = util::fromTimestampString("1979-02-23 18:48:31").toMillis();
+
+  auto timestamps = std::vector<int64_t>{
+      0, denverUTC, denverUTC, viennaUTC, chathamUTC, chathamUTC};
+
+  auto timezones = std::vector<TimeZoneKey>{
+      {0,
+       (int16_t)util::getTimeZoneID("America/Denver"),
+       (int16_t)util::getTimeZoneID("-06:00"),
+       (int16_t)util::getTimeZoneID("Europe/Vienna"),
+       (int16_t)util::getTimeZoneID("Pacific/Chatham"),
+       (int16_t)util::getTimeZoneID("+13:45")}};
+
+  const auto expected = makeTimestampWithTimeZoneVector(
+      timestamps.size(),
+      [&](int32_t index) { return timestamps[index]; },
+      [&](int32_t index) { return timezones[index]; });
+  expected->setNull(0, true);
+
+  testCast(stringVector, expected);
+}
+
+TEST_F(TimestampWithTimeZoneCastTest, fromVarcharInvalidInput) {
+  const auto invalidStringVector1 = makeNullableFlatVector<StringView>(
+      {"2012-10-31 01:00:47fooAmerica/Los_Angeles"});
+
+  const auto invalidStringVector2 = makeNullableFlatVector<StringView>(
+      {"2012-10-31 01:00:47 America/California"});
+
+  const auto invalidStringVector3 = makeNullableFlatVector<StringView>(
+      {"2012-10-31foo01:00:47 America/Los_Angeles"});
+
+  const auto invalidStringVector4 = makeNullableFlatVector<StringView>(
+      {"2012-10-31 35:00:47 America/Los_Angeles"});
+
+  auto millis = util::fromTimestampString("2012-10-31 07:00:47").toMillis();
+  auto timestamps = std::vector<int64_t>{millis};
+
+  auto timezones =
+      std::vector<TimeZoneKey>{(int16_t)util::getTimeZoneID("America/Denver")};
+
+  const auto expected = makeTimestampWithTimeZoneVector(
+      timestamps.size(),
+      [&](int32_t index) { return timestamps[index]; },
+      [&](int32_t index) { return timezones[index]; });
+
+  EXPECT_THROW(testCast(invalidStringVector1, expected), VeloxUserError);
+  EXPECT_THROW(testCast(invalidStringVector2, expected), VeloxUserError);
+  EXPECT_THROW(testCast(invalidStringVector3, expected), VeloxUserError);
+  EXPECT_THROW(testCast(invalidStringVector4, expected), VeloxUserError);
+}
+
 TEST_F(TimestampWithTimeZoneCastTest, toTimestamp) {
   auto timestamps =
       std::vector<int64_t>{1996 * kMillisInSecond, 0, 19920 * kMillisInSecond};

diff --git a/velox/functions/prestosql/types/TimestampWithTimeZoneType.cpp b/velox/functions/prestosql/types/TimestampWithTimeZoneType.cpp
@@ -14,6 +14,12 @@
  * limitations under the License.
  */
 #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
+#include <time.h>
+#include <chrono>
+#include "velox/expression/StringWriter.h"
+#include "velox/external/date/tz.h"
+#include "velox/type/Timestamp.h"
+#include "velox/type/TimestampConversion.h"
 #include "velox/type/tz/TimeZoneMap.h"
 
 namespace facebook::velox {
@@ -53,15 +59,48 @@ void castFromTimestamp(
   });
 }
 
+void castFromString(
+    const SimpleVector<StringView>& inputVector,
+    exec::EvalCtx& context,
+    const SelectivityVector& rows,
+    VectorPtr& result) {
+  auto timestampWithTzVector = result->asFlatVector<int64_t>();
+  auto rawTimestampWithTzValues =
+      timestampWithTzVector->values()->asMutable<int64_t>();
+  timestampWithTzVector->clearNulls(rows);
+
+  context.applyToSelectedNoThrow(rows, [&](auto row) {
+    if (inputVector.isNullAt(row)) {
+      result->setNull(row, true);
+    } else {
+      auto tsTzStr = std::string_view(
+          inputVector.valueAt(row).data(), inputVector.valueAt(row).size());
+      auto [ts, tzID] =
+          util::fromTimestampWithTimezoneString(StringView(tsTzStr));
+      ts.toGMT(tzID);
+      int64_t outputMilliseconds = ts.getSeconds() * kMillisInSecond;
+      rawTimestampWithTzValues[row] = pack(outputMilliseconds, tzID);
+    }
+  });
+}
+
 template <TypeKind kind>
 void castToTimestampWithTimeZone(
     const BaseVector& input,
     exec::EvalCtx& context,
     const SelectivityVector& rows,
     VectorPtr& result) {
-  VELOX_CHECK_EQ(kind, TypeKind::TIMESTAMP)
-  const auto inputVector = input.as<SimpleVector<Timestamp>>();
-  castFromTimestamp(*inputVector, context, rows, result);
+  if (kind == TypeKind::TIMESTAMP) {
+    const auto inputVector = input.as<SimpleVector<Timestamp>>();
+    castFromTimestamp(*inputVector, context, rows, result);
+  } else if (kind == TypeKind::VARCHAR) {
+    const auto inputVector = input.as<SimpleVector<StringView>>();
+    castFromString(*inputVector, context, rows, result);
+  } else {
+    VELOX_UNSUPPORTED(
+        "Cast from {} to TIMESTAMP WITH TIME ZONE not yet supported",
+        mapTypeKindToName(kind));
+  }
 }
 
 void castToTimestamp(
@@ -107,7 +146,7 @@ bool TimestampWithTimeZoneCastOperator::isSupportedFromType(
     case TypeKind::TIMESTAMP:
       return true;
     case TypeKind::VARCHAR:
-      // TODO: support cast from VARCHAR
+      return true;
     case TypeKind::INTEGER:
       // TODO: support cast from DATE
     default: