Add millisecond Presto function (#275)

Summary: Pull Request resolved: #275 Reviewed By: mbasmanova Differential Revision: D31122342 fbshipit-source-id: 844e7ad73e34846b4b2d34160b1c4596fc686068
facebookincubator · Sep 27, 2021 · 1e19e1e · 1e19e1e
1 parent e034dbc
commit 1e19e1e
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 0 deletions.
diff --git a/velox/docs/functions/datetime.rst b/velox/docs/functions/datetime.rst
@@ -11,6 +11,10 @@ Date and Time Functions
     Returns the UNIX timestamp ``unixtime`` as a timestamp with time zone
     using ``string`` for the time zone.
 
+.. function:: millisecond(timestamp) -> int64
+
+    Returns the millisecond of the second from ``timestamp``.
+
 .. function:: to_unixtime(timestamp) -> double
 
     Returns ``timestamp`` as a UNIX timestamp.
diff --git a/velox/functions/prestosql/CoreFunctions.cpp b/velox/functions/prestosql/CoreFunctions.cpp
@@ -54,6 +54,7 @@ void registerFunctions() {
   registerFunction<udf_to_unixtime, double, Timestamp>(
       {"to_unixtime", "to_unix_timestamp"});
   registerFunction<udf_from_unixtime, Timestamp, double>();
+  registerFunction<udf_millisecond, int64_t, Timestamp>();
 
   registerArithmeticFunctions();
   registerCheckedArithmeticFunctions();

diff --git a/velox/functions/prestosql/DateTimeImpl.h b/velox/functions/prestosql/DateTimeImpl.h
@@ -21,6 +21,7 @@
 namespace facebook::velox::functions {
 namespace {
 constexpr double kNanosecondsInSecond = 1'000'000'000;
+constexpr int64_t kNanosecondsInMilliseconds = 1'000'000;
 } // namespace
 
 FOLLY_ALWAYS_INLINE double toUnixtime(const Timestamp& timestamp) {
@@ -39,4 +40,13 @@ FOLLY_ALWAYS_INLINE std::optional<Timestamp> fromUnixtime(double unixtime) {
   return Timestamp(seconds, nanos * kNanosecondsInSecond);
 }
 
+VELOX_UDF_BEGIN(millisecond)
+FOLLY_ALWAYS_INLINE bool call(
+    int64_t& result,
+    const arg_type<Timestamp>& timestamp) {
+  result = timestamp.getNanos() / kNanosecondsInMilliseconds;
+  return true;
+}
+VELOX_UDF_END();
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
@@ -16,6 +16,7 @@
 
 #include "velox/functions/prestosql/TimestampWithTimeZoneType.h"
 #include "velox/functions/prestosql/tests/FunctionBaseTest.h"
+#include "velox/type/Timestamp.h"
 
 using namespace facebook::velox;
 
@@ -118,3 +119,14 @@ TEST_F(DateTimeFunctionsTest, fromUnixtime) {
   EXPECT_EQ(
       std::nullopt, fromUnixtime(std::numeric_limits<double>::quiet_NaN()));
 }
+
+TEST_F(DateTimeFunctionsTest, millisecond) {
+  const auto millisecond = [&](std::optional<Timestamp> timestamp) {
+    return evaluateOnce<int64_t>("millisecond(c0)", timestamp);
+  };
+  EXPECT_EQ(std::nullopt, millisecond(std::nullopt));
+  EXPECT_EQ(0, millisecond(Timestamp(0, 0)));
+  EXPECT_EQ(0, millisecond(Timestamp(4000000000, 0)));
+  EXPECT_EQ(123, millisecond(Timestamp(-1, 123000000)));
+  EXPECT_EQ(12300, millisecond(Timestamp(-1, 12300000000)));
+}