Add support for hour(timestamp_ntz)

Author: Mariamalmesfer

velox/functions/sparksql/DateTimeFunctions.h

@@ -21,6 +21,7 @@
 #include "velox/functions/lib/DateTimeFormatter.h"
 #include "velox/functions/lib/TimeUtils.h"
 #include "velox/functions/sparksql/TimestampUtils.h"
+#include "velox/functions/sparksql/types/TimestampNTZType.h"
 #include "velox/type/TimestampConversion.h"
 #include "velox/type/tz/TimeZoneMap.h"
 
@@ -886,12 +887,28 @@ struct NextDayFunction {
 template <typename T>
 struct HourFunction : public InitSessionTimezone<T> {
   VELOX_DEFINE_FUNCTION_TYPES(T);
+  using InitSessionTimezone<T>::initialize;
+
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& /*config*/,
+      const arg_type<TimestampNTZ>* /*timestamp*/) {
+    // TIMESTAMP_NTZ represents local wall-clock time and must not use
+    // session timezone adjustment.
+    this->timeZone_ = nullptr;
+  }
 
   FOLLY_ALWAYS_INLINE void call(
       int32_t& result,
       const arg_type<Timestamp>& timestamp) {
     result = getDateTime(timestamp, this->timeZone_).tm_hour;
   }
+
+  FOLLY_ALWAYS_INLINE void call(
+      int32_t& result,
+      const arg_type<TimestampNTZ>& timestamp) {
+    result = getDateTime(Timestamp::fromMicros(timestamp), nullptr).tm_hour;
+  }
 };
 
 template <typename T>

velox/functions/sparksql/registration/Register.cpp

@@ -16,6 +16,7 @@
 #include "velox/functions/sparksql/registration/Register.h"
 #include "velox/expression/SimpleFunctionRegistry.h"
 #include "velox/expression/SpecialFormRegistry.h"
+#include "velox/functions/sparksql/types/TimestampNTZRegistration.h"
 
 namespace facebook::velox::functions::sparksql {
 
@@ -34,6 +35,7 @@ extern void registerStringFunctions(const std::string& prefix);
 extern void registerUrlFunctions(const std::string& prefix);
 
 void registerFunctions(const std::string& prefix) {
+  registerTimestampNTZType();
   registerArrayFunctions(prefix);
   registerBinaryFunctions(prefix);
   registerBitwiseFunctions(prefix);

velox/functions/sparksql/registration/RegisterDatetime.cpp

@@ -76,6 +76,7 @@ void registerDatetimeFunctions(const std::string& prefix) {
   registerFunction<GetTimestampFunction, Timestamp, Varchar, Varchar>(
       {prefix + "get_timestamp"});
   registerFunction<HourFunction, int32_t, Timestamp>({prefix + "hour"});
+  registerFunction<HourFunction, int32_t, TimestampNTZ>({prefix + "hour"});
   registerFunction<MinuteFunction, int32_t, Timestamp>({prefix + "minute"});
   registerFunction<SecondFunction, int32_t, Timestamp>({prefix + "second"});
   registerFunction<MakeYMIntervalFunction, IntervalYearMonth>(

velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp

@@ -16,6 +16,7 @@
 
 #include "velox/common/base/tests/GTestUtils.h"
 #include "velox/functions/sparksql/tests/SparkFunctionBaseTest.h"
+#include "velox/functions/sparksql/types/TimestampNTZType.h"
 #include "velox/type/Timestamp.h"
 #include "velox/type/tz/TimeZoneMap.h"
 
@@ -943,6 +944,29 @@ TEST_F(DateTimeFunctionsTest, hour) {
   EXPECT_EQ(2, hour("1969-01-01 13:23:00.001"));
 }
 
+TEST_F(DateTimeFunctionsTest, hourTimestampNTZ) {
+  const auto hourNtz = [&](const StringView timestampStr) {
+    auto micros = std::make_optional(parseTimestamp(timestampStr).toMicros());
+    return evaluateOnce<int32_t>("hour(c0)", TIMESTAMP_NTZ(), micros);
+  };
+
+  EXPECT_EQ(
+      std::nullopt,
+      evaluateOnce<int32_t>(
+          "hour(c0)", TIMESTAMP_NTZ(), std::optional<int64_t>{}));
+
+  EXPECT_EQ(0, hourNtz("2024-01-08 00:23:00.001"));
+  EXPECT_EQ(1, hourNtz("2024-01-08 01:23:00.001"));
+  EXPECT_EQ(13, hourNtz("2024-01-20 13:23:00.001"));
+
+  // TIMESTAMP_NTZ should not be affected by session timezone.
+  setQueryTimeZone("Pacific/Apia");
+
+  EXPECT_EQ(0, hourNtz("2024-01-08 00:23:00.001"));
+  EXPECT_EQ(1, hourNtz("2024-01-08 01:23:00.001"));
+  EXPECT_EQ(13, hourNtz("2024-01-20 13:23:00.001"));
+}
+
 TEST_F(DateTimeFunctionsTest, minute) {
   const auto minute = [&](const StringView& timestampStr) {
     const auto timeStamp = std::make_optional(parseTimestamp(timestampStr));