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