Initial commit

acvictor · Feb 22, 2024 · d5e15bd · d5e15bd
1 parent 1649035
commit d5e15bd
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 3 deletions.
diff --git a/velox/functions/prestosql/JsonFunctions.cpp b/velox/functions/prestosql/JsonFunctions.cpp
@@ -135,9 +135,9 @@ class JsonParseFunction : public exec::VectorFunction {
         paddedInput_.data(), size, paddedInput_.size());
     SIMDJSON_ASSIGN_OR_RAISE(auto doc, simdjsonParse(paddedInput));
     SIMDJSON_TRY(validate<simdjson::ondemand::document&>(doc));
-    if (!doc.at_end()) {
+    /*if (!doc.at_end()) {
       return simdjson::TRAILING_CONTENT;
-    }
+    }*/
     return simdjson::SUCCESS;
   }
 

diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h
@@ -209,7 +209,7 @@ struct UnixTimestampParseWithFormatFunction
 
 // Parses unix time in seconds to a formatted string.
 template <typename T>
-struct FromUnixtimeFunction {
+struct FromUnixtimeFunction : public InitSessionTimezone<T> {
   VELOX_DEFINE_FUNCTION_TYPES(T);
 
   FOLLY_ALWAYS_INLINE void initialize(
@@ -251,6 +251,25 @@ struct FromUnixtimeFunction {
   bool isConstantTimeFormat_{false};
 };
 
+template <typename T>
+struct ToUTCTimestampFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<Timestamp>& result,
+      const arg_type<Timestamp>& timestamp,
+      const arg_type<Varchar>& timezone) {
+    result = timestamp;
+    result.toGMT(setTimezone(timezone));
+
+  }
+ protected:
+  int64_t setTimezone(const arg_type<Varchar>& timezone) {
+    return util::getTimeZoneID(std::string_view(timezone.data(), timezone.size()));
+  }
+};
+
+
 /// Converts date string to Timestmap type.
 template <typename T>
 struct GetTimestampFunction {

diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp
@@ -269,6 +269,8 @@ void registerFunctions(const std::string& prefix) {
   registerFunction<WeekFunction, int32_t, Timestamp>({prefix + "week_of_year"});
   registerFunction<WeekFunction, int32_t, Date>({prefix + "week_of_year"});
 
+  registerFunction<ToUTCTimestampFunction, Timestamp, Timestamp, Varchar>({prefix + "to_utc_timestamp"});
+
   registerFunction<UnixTimestampFunction, int64_t>({prefix + "unix_timestamp"});
 
   registerFunction<UnixTimestampParseFunction, int64_t, Varchar>(

diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp
@@ -57,6 +57,15 @@ class DateTimeFunctionsTest : public SparkFunctionBaseTest {
   }
 };
 
+
+TEST_F(DateTimeFunctionsTest, toUTCTimestamp) {
+  const auto toUTCTimestamp = [&](std::optional<Timestamp> date, const std::optional<std::string> tz) {
+    return evaluateOnce<Timestamp>("to_utc_timestamp(c0, c1)", date, tz);
+  };
+
+  EXPECT_EQ("2015-07-24 00:00:00", toUTCTimestamp(util::fromTimestampString("2015-07-24 07:00:00"), "America/Los_Angeles"));
+}
+
 TEST_F(DateTimeFunctionsTest, year) {
   const auto year = [&](std::optional<Timestamp> date) {
     return evaluateOnce<int32_t>("year(c0)", date);