From 4e7bbec42359e22b659a5d6182b0f43b7a668014 Mon Sep 17 00:00:00 2001
From: Aniket Mokashi <anim@fb.com>
Date: Thu, 23 Sep 2021 09:19:29 -0700
Subject: [PATCH] Add millisecond Presto function (#275)

Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/275

Reviewed By: mbasmanova

Differential Revision: D31122342

fbshipit-source-id: a137f16043ba217364d94df0c4342495cc58792e
---
 velox/docs/functions/datetime.rst                    |  4 ++++
 velox/functions/prestosql/CoreFunctions.cpp          |  1 +
 velox/functions/prestosql/DateTimeImpl.h             | 12 +++++++++++-
 .../prestosql/tests/DateTimeFunctionsTest.cpp        | 12 ++++++++++++
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/velox/docs/functions/datetime.rst b/velox/docs/functions/datetime.rst
index aa2d35ba09a64..2aff80d1d4b3a 100644
--- a/velox/docs/functions/datetime.rst
+++ b/velox/docs/functions/datetime.rst
@@ -14,3 +14,7 @@ Date and Time Functions
 .. function:: to_unixtime(timestamp) -> double
 
     Returns ``timestamp`` as a UNIX timestamp.
+
+.. function:: millisecond(timestamp) -> int64
+
+    Returns the millisecond of the second from ``timestamp``.
diff --git a/velox/functions/prestosql/CoreFunctions.cpp b/velox/functions/prestosql/CoreFunctions.cpp
index 38e5ae033cdb5..8908773f9bcb4 100644
--- 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
index 314f9716d4090..80c18f49e0b91 100644
--- a/velox/functions/prestosql/DateTimeImpl.h
+++ b/velox/functions/prestosql/DateTimeImpl.h
@@ -20,7 +20,8 @@
 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) {
   double result = timestamp.getSeconds();
@@ -34,4 +35,13 @@ FOLLY_ALWAYS_INLINE 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
index a07c9ec79540e..e28347e9fb6dd 100644
--- 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;
 
@@ -102,3 +103,14 @@ TEST_F(DateTimeFunctionsTest, fromUnixtimeWithTimeZone) {
     assertEqualVectors(expected, result);
   }
 }
+
+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)));
+}