Add millisecond Presto function (#275)

Summary: Pull Request resolved: #275 Reviewed By: mbasmanova Differential Revision: D31122342 fbshipit-source-id: a137f16043ba217364d94df0c4342495cc58792e
facebookincubator · Sep 23, 2021 · 4e7bbec · 4e7bbec
1 parent 5b68ac5
commit 4e7bbec
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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)));
+}