zclllyybb commented on code in PR #49119:
URL: https://github.com/apache/doris/pull/49119#discussion_r2020370693


##########
be/src/vec/functions/function_date_or_datetime_computation.h:
##########
@@ -998,4 +998,97 @@ class CurrentDateFunctionBuilder : public 
FunctionBuilderImpl {
     }
 };
 
+class FunctionNextDay : public IFunction {
+public:
+    static constexpr auto name = "next_day";
+    static FunctionPtr create() { return std::make_shared<FunctionNextDay>(); }
+    String get_name() const override { return name; }
+    size_t get_number_of_arguments() const override { return 2; }
+    DataTypePtr get_return_type_impl(const ColumnsWithTypeAndName& arguments) 
const override {
+        return make_nullable(std::make_shared<DataTypeDateV2>());
+    }
+
+    Status execute_impl(FunctionContext* context, Block& block, const 
ColumnNumbers& arguments,
+                        uint32_t result, size_t input_rows_count) const 
override {
+        CHECK_EQ(arguments.size(), 2);
+        auto res = ColumnDateV2::create();
+        const auto& [left_col, left_const] =
+                unpack_if_const(block.get_by_position(arguments[0]).column);
+        const auto& [right_col, right_const] =
+                unpack_if_const(block.get_by_position(arguments[1]).column);
+        const auto& date_col = *assert_cast<const 
ColumnDateV2*>(left_col.get());
+        const auto& week_col = *assert_cast<const 
ColumnString*>(right_col.get());
+        Status status;
+        if (left_const && right_const) {
+            status = execute_vector<true, true>(input_rows_count, date_col, 
week_col, *res);
+        } else if (left_const) {
+            status = execute_vector<true, false>(input_rows_count, date_col, 
week_col, *res);
+        } else if (right_const) {
+            status = execute_vector<false, true>(input_rows_count, date_col, 
week_col, *res);
+        } else {
+            status = execute_vector<false, false>(input_rows_count, date_col, 
week_col, *res);
+        }
+        if (!status.ok()) {
+            return status;
+        }
+        block.replace_by_position(result, std::move(res));
+        return Status::OK();
+    }
+
+private:
+    static int day_of_week(const StringRef& weekday) {
+        static const std::unordered_map<std::string, int> weekday_map = {
+                {"SU", 1}, {"SUN", 1}, {"SUNDAY", 1},   {"MO", 2}, {"MON", 2}, 
{"MONDAY", 2},
+                {"TU", 3}, {"TUE", 3}, {"TUESDAY", 3},  {"WE", 4}, {"WED", 4}, 
{"WEDNESDAY", 4},
+                {"TH", 5}, {"THU", 5}, {"THURSDAY", 5}, {"FR", 6}, {"FRI", 6}, 
{"FRIDAY", 6},
+                {"SA", 7}, {"SAT", 7}, {"SATURDAY", 7},
+        };
+        auto weekday_upper = weekday.to_string();
+        std::transform(weekday_upper.begin(), weekday_upper.end(), 
weekday_upper.begin(),
+                       ::toupper);
+        auto it = weekday_map.find(weekday_upper);
+        if (it == weekday_map.end()) {
+            return 0;
+        }
+        return it->second;
+    }
+    static Status compute_next_day(DateV2Value<DateV2ValueType>& dtv, const 
int week_day) {
+        auto days_to_add = (week_day - dtv.day_of_week() + 7) % 7;
+        days_to_add = days_to_add == 0 ? 7 : days_to_add;
+        dtv.date_add_interval<TimeUnit::DAY>(TimeInterval(TimeUnit::DAY, 
days_to_add, false));
+        return Status::OK();
+    }
+
+    template <bool left_const, bool right_const>
+    static Status execute_vector(size_t input_rows_count, const ColumnDateV2& 
left_col,
+                                 const ColumnString& right_col, ColumnDateV2& 
res_col) {
+        DateV2Value<DateV2ValueType> dtv;
+        int week_day;
+        if constexpr (left_const) {
+            dtv = binary_cast<UInt32, 
DateV2Value<DateV2ValueType>>(left_col.get_element(0));
+        }
+        if constexpr (right_const) {
+            week_day = day_of_week(right_col.get_data_at(0));
+            if (week_day == 0) {
+                return Status::InvalidArgument("Invalid weekday: {}", 
right_col.get_data_at(0));
+            }
+        }
+
+        for (size_t i = 0; i < input_rows_count; ++i) {
+            if constexpr (!left_const) {
+                dtv = binary_cast<UInt32, 
DateV2Value<DateV2ValueType>>(left_col.get_element(i));
+            }
+            if constexpr (!right_const) {
+                week_day = day_of_week(right_col.get_data_at(i));
+                if (week_day == 0) {
+                    return Status::InvalidArgument("Invalid weekday: {}", 
right_col.get_data_at(i));
+                }
+            }
+            RETURN_IF_ERROR(compute_next_day(dtv, week_day));
+            res_col.insert_value(binary_cast<DateV2Value<DateV2ValueType>, 
UInt32>(dtv));

Review Comment:
   > Maybe the logic can be made simpler by calling the `reserve` function. 
`res->reserve(input_rows_count);`
   
   Directly using PODArray is more efficient because it skips checks in 
`insert_value`. but those checks take few time so use `reserve` is also ok.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@doris.apache.org
For additional commands, e-mail: commits-h...@doris.apache.org

Reply via email to