Add next_day Spark function (apache#7333)

Summary:
Fixes facebookincubator/velox#4741

Pull Request resolved: facebookincubator/velox#7333

Reviewed By: xiaoxmeng, Yuhta, pedroerp

Differential Revision: D52070541

Pulled By: mbasmanova

fbshipit-source-id: 4b915883284a93dbe34995a6d2f4b0f06dbc3330

velox/docs/functions/spark/datetime.rst

@@ -96,7 +96,22 @@ These functions support TIMESTAMP and DATE input types.
     Returns the month of ``date``. ::
 
         SELECT month('2009-07-30'); -- 7
-       
+
+.. spark:function:: next_day(startDate, dayOfWeek) -> date
+
+    Returns the first date which is later than ``startDate`` and named as ``dayOfWeek``.
+    Returns null if ``dayOfWeek`` is invalid.
+    ``dayOfWeek`` is case insensitive and must be one of the following:
+    ``SU``, ``SUN``, ``SUNDAY``, ``MO``, ``MON``, ``MONDAY``, ``TU``, ``TUE``, ``TUESDAY``,
+    ``WE``, ``WED``, ``WEDNESDAY``, ``TH``, ``THU``, ``THURSDAY``, ``FR``, ``FRI``, ``FRIDAY``,
+    ``SA``, ``SAT``, ``SATURDAY``. ::
+
+        SELECT next_day('2015-07-23', "Mon"); -- '2015-07-27'
+        SELECT next_day('2015-07-23', "mo"); -- '2015-07-27'
+        SELECT next_day('2015-07-23', "Tue"); -- '2015-07-28'
+        SELECT next_day('2015-07-23', "tu"); -- '2015-07-28'
+        SELECT next_day('2015-07-23', "we"); -- '2015-07-29'
+
 .. spark:function:: to_unix_timestamp(string) -> integer
 
     Alias for ``unix_timestamp(string) -> integer``.

velox/functions/lib/CMakeLists.txt

@@ -30,7 +30,8 @@ add_library(
   Re2Functions.cpp
   StringEncodingUtils.cpp
   SubscriptUtil.cpp
-  CheckNestedNulls.cpp)
+  CheckNestedNulls.cpp
+  TimeUtils.cpp)
 
 target_link_libraries(velox_functions_lib velox_functions_util velox_vector
                       re2::re2 Folly::folly)

velox/functions/lib/TimeUtils.cpp

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/functions/lib/TimeUtils.h"
+
+namespace facebook::velox::functions {
+
+const folly::F14FastMap<std::string, int8_t> kDayOfWeekNames{
+    {"th", 0},       {"fr", 1},     {"sa", 2},       {"su", 3},
+    {"mo", 4},       {"tu", 5},     {"we", 6},       {"thu", 0},
+    {"fri", 1},      {"sat", 2},    {"sun", 3},      {"mon", 4},
+    {"tue", 5},      {"wed", 6},    {"thursday", 0}, {"friday", 1},
+    {"saturday", 2}, {"sunday", 3}, {"monday", 4},   {"tuesday", 5},
+    {"wednesday", 6}};
+} // namespace facebook::velox::functions

velox/functions/lib/TimeUtils.h

@@ -21,9 +21,10 @@
 #include "velox/functions/Macros.h"
 
 namespace facebook::velox::functions {
-namespace {
+
 inline constexpr int64_t kSecondsInDay = 86'400;
 inline constexpr int64_t kDaysInWeek = 7;
+extern const folly::F14FastMap<std::string, int8_t> kDayOfWeekNames;
 
 FOLLY_ALWAYS_INLINE const date::time_zone* getTimeZoneFromConfig(
     const core::QueryConfig& config) {
@@ -45,7 +46,6 @@ getSeconds(Timestamp timestamp, const date::time_zone* timeZone) {
     return timestamp.getSeconds();
   }
 }
-} // namespace
 
 FOLLY_ALWAYS_INLINE
 std::tm getDateTime(Timestamp timestamp, const date::time_zone* timeZone) {

velox/functions/sparksql/DateTimeFunctions.h

@@ -408,4 +408,62 @@ struct DayOfYearFunction {
     result = getDayOfYear(getDateTime(date));
   }
 };
+
+template <typename T>
+struct NextDayFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void initialize(
+      const core::QueryConfig& /*config*/,
+      const arg_type<Date>* /*startDate*/,
+      const arg_type<Varchar>* dayOfWeek) {
+    if (dayOfWeek != nullptr) {
+      weekDay_ = getDayOfWeekFromString(*dayOfWeek);
+      if (!weekDay_.has_value()) {
+        invalidFormat_ = true;
+      }
+    }
+  }
+
+  FOLLY_ALWAYS_INLINE bool call(
+      out_type<Date>& result,
+      const arg_type<Date>& startDate,
+      const arg_type<Varchar>& dayOfWeek) {
+    if (invalidFormat_) {
+      return false;
+    }
+    auto weekDay = weekDay_.has_value() ? weekDay_.value()
+                                        : getDayOfWeekFromString(dayOfWeek);
+    if (!weekDay.has_value()) {
+      return false;
+    }
+    auto nextDay = getNextDate(startDate, weekDay.value());
+    if (nextDay != (int32_t)nextDay) {
+      return false;
+    }
+    result = nextDay;
+    return true;
+  }
+
+ private:
+  static FOLLY_ALWAYS_INLINE std::optional<int8_t> getDayOfWeekFromString(
+      const StringView& dayOfWeek) {
+    std::string lowerDayOfWeek =
+        boost::algorithm::to_lower_copy(dayOfWeek.str());
+    auto it = kDayOfWeekNames.find(lowerDayOfWeek);
+    if (it != kDayOfWeekNames.end()) {
+      return it->second;
+    }
+    return std::nullopt;
+  }
+
+  static FOLLY_ALWAYS_INLINE int64_t
+  getNextDate(int64_t startDay, int8_t dayOfWeek) {
+    return startDay + 1 + ((dayOfWeek - 1 - startDay) % 7 + 7) % 7;
+  }
+
+  std::optional<int8_t> weekDay_;
+  bool invalidFormat_{false};
+};
+
 } // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/Register.cpp

@@ -285,6 +285,8 @@ void registerFunctions(const std::string& prefix) {
 
   registerFunction<MonthFunction, int32_t, Date>({prefix + "month"});
 
+  registerFunction<NextDayFunction, Date, Date, Varchar>({prefix + "next_day"});
+
   // Register bloom filter function
   registerFunction<BloomFilterMightContainFunction, bool, Varbinary, int64_t>(
       {prefix + "might_contain"});

velox/functions/sparksql/tests/DateTimeFunctionsTest.cpp

@@ -456,5 +456,52 @@ TEST_F(DateTimeFunctionsTest, quarterDate) {
   EXPECT_EQ(3, quarter("1954-08-08"));
 }
 
+TEST_F(DateTimeFunctionsTest, nextDay) {
+  const auto nextDay = [&](const std::string& date,
+                           const std::string& dayOfWeek) {
+    auto startDates =
+        makeNullableFlatVector<int32_t>({parseDate(date)}, DATE());
+    auto dayOfWeeks = makeNullableFlatVector<std::string>({dayOfWeek});
+
+    auto result = evaluateOnce<int32_t>(
+        fmt::format("next_day(c0, '{}')", dayOfWeek),
+        makeRowVector({startDates}));
+
+    auto anotherResult = evaluateOnce<int32_t>(
+        "next_day(c0, c1)", makeRowVector({startDates, dayOfWeeks}));
+
+    EXPECT_EQ(result, anotherResult);
+    std::optional<std::string> res;
+    if (result.has_value()) {
+      res = DATE()->toString(result.value());
+    }
+    return res;
+  };
+
+  EXPECT_EQ(nextDay("2015-07-23", "Mon"), "2015-07-27");
+  EXPECT_EQ(nextDay("2015-07-23", "mo"), "2015-07-27");
+  EXPECT_EQ(nextDay("2015-07-23", "monday"), "2015-07-27");
+  EXPECT_EQ(nextDay("2015-07-23", "Tue"), "2015-07-28");
+  EXPECT_EQ(nextDay("2015-07-23", "tu"), "2015-07-28");
+  EXPECT_EQ(nextDay("2015-07-23", "tuesday"), "2015-07-28");
+  EXPECT_EQ(nextDay("2015-07-23", "we"), "2015-07-29");
+  EXPECT_EQ(nextDay("2015-07-23", "wed"), "2015-07-29");
+  EXPECT_EQ(nextDay("2015-07-23", "wednesday"), "2015-07-29");
+  EXPECT_EQ(nextDay("2015-07-23", "Thu"), "2015-07-30");
+  EXPECT_EQ(nextDay("2015-07-23", "TH"), "2015-07-30");
+  EXPECT_EQ(nextDay("2015-07-23", "thursday"), "2015-07-30");
+  EXPECT_EQ(nextDay("2015-07-23", "Fri"), "2015-07-24");
+  EXPECT_EQ(nextDay("2015-07-23", "fr"), "2015-07-24");
+  EXPECT_EQ(nextDay("2015-07-23", "friday"), "2015-07-24");
+  EXPECT_EQ(nextDay("2015-07-31", "wed"), "2015-08-05");
+  EXPECT_EQ(nextDay("2015-07-23", "saturday"), "2015-07-25");
+  EXPECT_EQ(nextDay("2015-07-23", "sunday"), "2015-07-26");
+  EXPECT_EQ(nextDay("2015-12-31", "Fri"), "2016-01-01");
+
+  EXPECT_EQ(nextDay("2015-07-23", "xx"), std::nullopt);
+  EXPECT_EQ(nextDay("2015-07-23", "\"quote"), std::nullopt);
+  EXPECT_EQ(nextDay("2015-07-23", ""), std::nullopt);
+}
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test