Initial

facebookincubator · Dec 25, 2024 · 9e5d869 · 9e5d869
1 parent 8ebd3a8
commit 9e5d869
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 5 deletions.
diff --git a/velox/functions/sparksql/specialforms/SparkCastExpr.cpp b/velox/functions/sparksql/specialforms/SparkCastExpr.cpp
@@ -22,7 +22,7 @@ exec::ExprPtr SparkCastCallToSpecialForm::constructSpecialForm(
     const TypePtr& type,
     std::vector<exec::ExprPtr>&& compiledChildren,
     bool trackCpuUsage,
-    const core::QueryConfig& /*config*/) {
+    const core::QueryConfig& config) {
   VELOX_CHECK_EQ(
       compiledChildren.size(),
       1,
@@ -33,14 +33,14 @@ exec::ExprPtr SparkCastCallToSpecialForm::constructSpecialForm(
       std::move(compiledChildren[0]),
       trackCpuUsage,
       false,
-      std::make_shared<SparkCastHooks>());
+      std::make_shared<SparkCastHooks>(config));
 }
 
 exec::ExprPtr SparkTryCastCallToSpecialForm::constructSpecialForm(
     const TypePtr& type,
     std::vector<exec::ExprPtr>&& compiledChildren,
     bool trackCpuUsage,
-    const core::QueryConfig& /*config*/) {
+    const core::QueryConfig& config) {
   VELOX_CHECK_EQ(
       compiledChildren.size(),
       1,
@@ -51,6 +51,6 @@ exec::ExprPtr SparkTryCastCallToSpecialForm::constructSpecialForm(
       std::move(compiledChildren[0]),
       trackCpuUsage,
       true,
-      std::make_shared<SparkCastHooks>());
+      std::make_shared<SparkCastHooks>(config));
 }
 } // namespace facebook::velox::functions::sparksql
diff --git a/velox/functions/sparksql/specialforms/SparkCastHooks.cpp b/velox/functions/sparksql/specialforms/SparkCastHooks.cpp
@@ -17,9 +17,13 @@
 #include "velox/functions/sparksql/specialforms/SparkCastHooks.h"
 #include "velox/functions/lib/string/StringImpl.h"
 #include "velox/type/TimestampConversion.h"
+#include "velox/type/tz/TimeZoneMap.h"
 
 namespace facebook::velox::functions::sparksql {
 
+SparkCastHooks::SparkCastHooks(const velox::core::QueryConfig& config)
+    : config_(config) {}
+
 Expected<Timestamp> SparkCastHooks::castStringToTimestamp(
     const StringView& view) const {
   return util::fromTimestampString(
@@ -76,12 +80,13 @@ StringView SparkCastHooks::removeWhiteSpaces(const StringView& view) const {
 
 const TimestampToStringOptions& SparkCastHooks::timestampToStringOptions()
     const {
-  static constexpr TimestampToStringOptions options = {
+  static TimestampToStringOptions options = {
       .precision = TimestampToStringOptions::Precision::kMicroseconds,
       .leadingPositiveSign = true,
       .skipTrailingZeros = true,
       .zeroPaddingYear = true,
       .dateTimeSeparator = ' ',
+      .timeZone = tz::locateZone(config_.sessionTimezone()),
   };
   return options;
 }

diff --git a/velox/functions/sparksql/specialforms/SparkCastHooks.h b/velox/functions/sparksql/specialforms/SparkCastHooks.h
@@ -16,13 +16,16 @@
 
 #pragma once
 
+#include "velox/core/QueryCtx.h"
 #include "velox/expression/CastHooks.h"
 
 namespace facebook::velox::functions::sparksql {
 
 // This class provides cast hooks following Spark semantics.
 class SparkCastHooks : public exec::CastHooks {
  public:
+  explicit SparkCastHooks(const velox::core::QueryConfig& config);
+
   // TODO: Spark hook allows more string patterns than Presto.
   Expected<Timestamp> castStringToTimestamp(
       const StringView& view) const override;
@@ -59,5 +62,8 @@ class SparkCastHooks : public exec::CastHooks {
   }
 
   exec::PolicyType getPolicy() const override;
+
+ private:
+  const core::QueryConfig& config_;
 };
 } // namespace facebook::velox::functions::sparksql