diff --git a/velox/functions/sparksql/DateTimeFunctions.h b/velox/functions/sparksql/DateTimeFunctions.h index 0ba287089e4a..61860ab38c90 100644 --- a/velox/functions/sparksql/DateTimeFunctions.h +++ b/velox/functions/sparksql/DateTimeFunctions.h @@ -107,6 +107,31 @@ struct UnixDateFunction { } }; +template +struct CurrentTimestampFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void initialize( + const std::vector& /*inputTypes*/, + const core::QueryConfig& config) { + timeZone_ = getTimeZoneFromConfig(config); + } + + FOLLY_ALWAYS_INLINE void call(out_type& result) { + auto now = std::chrono::system_clock::now(); + auto epoch = std::chrono::duration_cast( + now.time_since_epoch()) + .count(); + result = Timestamp::fromMicros(epoch); + if (timeZone_ != nullptr) { + result.toTimezone(*timeZone_); + } + } + + private: + const date::time_zone* timeZone_ = nullptr; +}; + template struct UnixTimestampFunction { // unix_timestamp(); diff --git a/velox/functions/sparksql/Register.cpp b/velox/functions/sparksql/Register.cpp index 4a00146594f5..35e673e5ca8b 100644 --- a/velox/functions/sparksql/Register.cpp +++ b/velox/functions/sparksql/Register.cpp @@ -284,6 +284,8 @@ void registerFunctions(const std::string& prefix) { registerFunction( {prefix + "from_utc_timestamp"}); + registerFunction({prefix + "current_timestamp"}); + registerFunction({prefix + "unix_date"}); registerFunction({prefix + "unix_timestamp"}); diff --git a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp index 383f8bd48d64..73391566b32d 100644 --- a/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp +++ b/velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp @@ -15,6 +15,7 @@ */ #include "velox/common/base/tests/GTestUtils.h" +#include "velox/external/date/tz.h" #include "velox/functions/sparksql/tests/SparkFunctionBaseTest.h" #include "velox/type/Timestamp.h" #include "velox/type/tz/TimeZoneMap.h" @@ -41,6 +42,17 @@ class DateTimeFunctionsTest : public SparkFunctionBaseTest { static constexpr int8_t kMinTinyint = std::numeric_limits::min(); static constexpr int8_t kMaxTinyint = std::numeric_limits::max(); +int64_t getCurrentTimestamp(const std::optional& timeZone) { + auto now = std::chrono::system_clock::now(); + auto tp = timeZone.has_value() + ? date::make_zoned( + timeZone.value(), now).get_local_time().time_since_epoch() + : now.time_since_epoch(); + auto since_epoch = std::chrono::duration_cast(tp); + LOG(INFO) << since_epoch.count(); + return since_epoch.count(); +} + protected: void setQueryTimeZone(const std::string& timeZone) { queryCtx_->testingOverrideConfigUnsafe({ @@ -228,6 +240,32 @@ TEST_F(DateTimeFunctionsTest, unixDate) { EXPECT_EQ(unixDate("-5877641-06-23"), kMin); } +TEST_F(DateTimeFunctionsTest, currentTimestamp) { + auto emptyRowVector = makeRowVector(ROW({}), 1); + + auto timestampBefore = getCurrentTimestamp(std::nullopt); + auto result = evaluateOnce("current_timestamp()", emptyRowVector); + auto timestampAfter = getCurrentTimestamp(std::nullopt); + + EXPECT_TRUE(result.has_value()); + auto resultInInt = result.value().toMicros(); + EXPECT_LE(timestampBefore, resultInInt); + EXPECT_LE(resultInInt, timestampAfter); + EXPECT_LE(timestampAfter - timestampBefore, 300); + + auto tz = "America/Los_Angeles"; + setQueryTimeZone(tz); + timestampBefore = getCurrentTimestamp(tz); + result = evaluateOnce("current_timestamp()", emptyRowVector); + timestampAfter = getCurrentTimestamp(tz); + + EXPECT_TRUE(result.has_value()); + resultInInt = result.value().toMicros(); + EXPECT_LE(timestampBefore, resultInInt); + EXPECT_LE(resultInInt, timestampAfter); + EXPECT_LE(timestampAfter - timestampBefore, 300); +} + TEST_F(DateTimeFunctionsTest, unixTimestamp) { const auto unixTimestamp = [&](std::optional dateStr) { return evaluateOnce("unix_timestamp(c0)", dateStr);