Fix array_join(Timestamp). (#10881)

Summary: Pull Request resolved: #10881 array_join(Timestamp) was returning results different from Presto Java. This was because it was not using session timezone, like CAST() does. The change is to make array_join(Timestamp) behave like CAST(Timestamp AS Varchar). Reviewed By: Yuhta Differential Revision: D61940944 fbshipit-source-id: 00655606cd45fea038aa44b4c1d7a1aeb86addb9
facebookincubator · Aug 29, 2024 · 5f935a7 · 5f935a7
1 parent d16b195
commit 5f935a7
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 19 deletions.
diff --git a/velox/expression/tests/CastExprTest.cpp b/velox/expression/tests/CastExprTest.cpp
@@ -52,25 +52,6 @@ class CastExprTest : public functions::test::CastBaseTest {
         {"error_on_odd_else_unknown"});
   }
 
-  void setLegacyCast(bool value) {
-    queryCtx_->testingOverrideConfigUnsafe({
-        {core::QueryConfig::kLegacyCast, std::to_string(value)},
-    });
-  }
-
-  void setCastMatchStructByName(bool value) {
-    queryCtx_->testingOverrideConfigUnsafe({
-        {core::QueryConfig::kCastMatchStructByName, std::to_string(value)},
-    });
-  }
-
-  void setTimezone(const std::string& value) {
-    queryCtx_->testingOverrideConfigUnsafe({
-        {core::QueryConfig::kSessionTimezone, value},
-        {core::QueryConfig::kAdjustTimestampToTimezone, "true"},
-    });
-  }
-
   std::shared_ptr<core::ConstantTypedExpr> makeConstantNullExpr(TypeKind kind) {
     return std::make_shared<core::ConstantTypedExpr>(
         createType(kind, {}), variant(kind));

diff --git a/velox/functions/prestosql/ArrayFunctions.h b/velox/functions/prestosql/ArrayFunctions.h
@@ -18,6 +18,7 @@
 #include <type_traits>
 
 #include "velox/expression/ComplexViewTypes.h"
+#include "velox/expression/PrestoCastHooks.h"
 #include "velox/functions/Udf.h"
 #include "velox/functions/lib/CheckedArithmetic.h"
 #include "velox/type/Conversions.h"
@@ -160,12 +161,38 @@ template <typename TExecCtx, typename T>
 struct ArrayJoinFunction {
   VELOX_DEFINE_FUNCTION_TYPES(TExecCtx);
 
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& inputTypes,
+      const core::QueryConfig& config,
+      const arg_type<velox::Array<T>>* /*arr*/,
+      const arg_type<Varchar>* /*delimiter*/) {
+    initialize(inputTypes, config, nullptr, nullptr, nullptr);
+  }
+
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config,
+      const arg_type<velox::Array<T>>* /*arr*/,
+      const arg_type<Varchar>* /*delimiter*/,
+      const arg_type<Varchar>* /*nullReplacement*/) {
+    const exec::PrestoCastHooks hooks{config};
+    options_ = hooks.timestampToStringOptions();
+  }
+
   template <typename C>
   void writeValue(out_type<velox::Varchar>& result, const C& value) {
     // To VARCHAR converter never throws.
     result += util::Converter<TypeKind::VARCHAR>::tryCast(value).value();
   }
 
+  void writeValue(out_type<velox::Varchar>& result, const Timestamp& value) {
+    Timestamp inputValue{value};
+    if (options_.timeZone) {
+      inputValue.toTimezone(*(options_.timeZone));
+    }
+    result += inputValue.toString(options_);
+  }
+
   template <typename C>
   void writeOutput(
       out_type<velox::Varchar>& result,
@@ -214,6 +241,9 @@ struct ArrayJoinFunction {
     createOutputString(result, inputArray, delim, nullReplacement.getString());
     return true;
   }
+
+ private:
+  TimestampToStringOptions options_;
 };
 
 /// Function Signature: combinations(array(T), n) -> array(array(T))

diff --git a/velox/functions/prestosql/tests/ArrayJoinTest.cpp b/velox/functions/prestosql/tests/ArrayJoinTest.cpp
@@ -101,4 +101,40 @@ TEST_F(ArrayJoinTest, boolTest) {
       {true, std::nullopt, false}, ","_sv, "apple"_sv, "true,apple,false"_sv);
 }
 
+TEST_F(ArrayJoinTest, timestampTest) {
+  setLegacyCast(false);
+  testArrayJoinNoReplacement<Timestamp>(
+      {Timestamp{333183, 0}, std::nullopt, Timestamp{2925183, 0}},
+      "~"_sv,
+      "1970-01-04 20:33:03.000~1970-02-03 20:33:03.000"_sv);
+  testArrayJoinReplacement<Timestamp>(
+      {Timestamp{333183, 0}, std::nullopt, Timestamp{2925183, 0}},
+      "~"_sv,
+      "<n/a>"_sv,
+      "1970-01-04 20:33:03.000~<n/a>~1970-02-03 20:33:03.000"_sv);
+
+  setLegacyCast(true);
+  testArrayJoinNoReplacement<Timestamp>(
+      {Timestamp{333183, 0}, std::nullopt, Timestamp{2925183, 0}},
+      "~"_sv,
+      "1970-01-04T20:33:03.000~1970-02-03T20:33:03.000"_sv);
+  testArrayJoinReplacement<Timestamp>(
+      {Timestamp{333183, 0}, std::nullopt, Timestamp{2925183, 0}},
+      "~"_sv,
+      "<missing>"_sv,
+      "1970-01-04T20:33:03.000~<missing>~1970-02-03T20:33:03.000"_sv);
+
+  setLegacyCast(false);
+  setTimezone("America/Los_Angeles");
+  testArrayJoinNoReplacement<Timestamp>(
+      {Timestamp{333183, 0}, std::nullopt, Timestamp{2925183, 0}},
+      "~"_sv,
+      "1970-01-04 12:33:03.000~1970-02-03 12:33:03.000"_sv);
+  testArrayJoinReplacement<Timestamp>(
+      {Timestamp{333183, 0}, std::nullopt, Timestamp{2925183, 0}},
+      "~"_sv,
+      "<absent>"_sv,
+      "1970-01-04 12:33:03.000~<absent>~1970-02-03 12:33:03.000"_sv);
+}
+
 } // namespace
diff --git a/velox/functions/prestosql/tests/utils/FunctionBaseTest.h b/velox/functions/prestosql/tests/utils/FunctionBaseTest.h
@@ -42,6 +42,25 @@ class FunctionBaseTest : public testing::Test,
       DoubleType,
       RealType>;
 
+  void setLegacyCast(bool value) {
+    queryCtx_->testingOverrideConfigUnsafe({
+        {core::QueryConfig::kLegacyCast, std::to_string(value)},
+    });
+  }
+
+  void setCastMatchStructByName(bool value) {
+    queryCtx_->testingOverrideConfigUnsafe({
+        {core::QueryConfig::kCastMatchStructByName, std::to_string(value)},
+    });
+  }
+
+  void setTimezone(const std::string& value) {
+    queryCtx_->testingOverrideConfigUnsafe({
+        {core::QueryConfig::kSessionTimezone, value},
+        {core::QueryConfig::kAdjustTimestampToTimezone, "true"},
+    });
+  }
+
  protected:
   static void SetUpTestCase();