[function](date function) add new date function 'to_monday' #13707

This commit is contained in:
lsy3993
2022-10-28 08:41:16 +08:00
committed by GitHub
parent f51464af59
commit c108554f14
7 changed files with 256 additions and 0 deletions

View File

@ -717,6 +717,124 @@ struct LastDayImpl {
}
};
template <typename DateType>
struct MondayImpl {
static constexpr auto name = "to_monday";
static Status execute_impl(FunctionContext* context, Block& block,
const ColumnNumbers& arguments, size_t result,
size_t input_rows_count) {
auto null_map = ColumnUInt8::create(input_rows_count, 0);
ColumnPtr res_column;
ColumnPtr argument_column =
block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
if constexpr (std::is_same_v<DateType, DataTypeDateTime> ||
std::is_same_v<DateType, DataTypeDate>) {
auto data_col = assert_cast<const ColumnVector<Int64>*>(argument_column.get());
res_column = ColumnInt64::create(input_rows_count);
execute_straight<VecDateTimeValue, Int64, Int64>(
input_rows_count, null_map->get_data(), data_col->get_data(),
static_cast<ColumnVector<Int64>*>(res_column->assume_mutable().get())
->get_data());
} else if constexpr (std::is_same_v<DateType, DataTypeDateV2>) {
auto data_col = assert_cast<const ColumnVector<UInt32>*>(argument_column.get());
res_column = ColumnVector<UInt32>::create(input_rows_count);
execute_straight<DateV2Value<DateV2ValueType>, UInt32, UInt32>(
input_rows_count, null_map->get_data(), data_col->get_data(),
static_cast<ColumnVector<UInt32>*>(res_column->assume_mutable().get())
->get_data());
} else if constexpr (std::is_same_v<DateType, DataTypeDateTimeV2>) {
auto data_col = assert_cast<const ColumnVector<UInt64>*>(argument_column.get());
res_column = ColumnVector<UInt32>::create(input_rows_count);
execute_straight<DateV2Value<DateTimeV2ValueType>, UInt32, UInt64>(
input_rows_count, null_map->get_data(), data_col->get_data(),
static_cast<ColumnVector<UInt32>*>(res_column->assume_mutable().get())
->get_data());
}
block.replace_by_position(
result, ColumnNullable::create(std::move(res_column), std::move(null_map)));
return Status::OK();
}
template <typename DateValueType, typename ReturnType, typename InputDateType>
static void execute_straight(size_t input_rows_count, NullMap& null_map,
const PaddedPODArray<InputDateType>& data_col,
PaddedPODArray<ReturnType>& res_data) {
for (int i = 0; i < input_rows_count; i++) {
if constexpr (std::is_same_v<DateValueType, VecDateTimeValue>) {
const auto& cur_data = data_col[i];
auto ts_value = binary_cast<Int64, VecDateTimeValue>(cur_data);
if (!ts_value.is_valid_date()) {
null_map[i] = 1;
continue;
}
if (is_special_day(ts_value.year(), ts_value.month(), ts_value.day())) {
ts_value.set_time(ts_value.year(), ts_value.month(), 1, 0, 0, 0);
ts_value.set_type(TIME_DATE);
res_data[i] = binary_cast<VecDateTimeValue, Int64>(ts_value);
continue;
}
// day_of_week, from 1(Mon) to 7(Sun)
int day_of_week = ts_value.weekday() + 1;
int gap_of_monday = day_of_week - 1;
TimeInterval interval(DAY, gap_of_monday, true);
ts_value.template date_add_interval<DAY>(interval);
ts_value.set_type(TIME_DATE);
res_data[i] = binary_cast<VecDateTimeValue, Int64>(ts_value);
} else if constexpr (std::is_same_v<DateValueType, DateV2Value<DateV2ValueType>>) {
const auto& cur_data = data_col[i];
auto ts_value = binary_cast<UInt32, DateValueType>(cur_data);
if (!ts_value.is_valid_date()) {
null_map[i] = 1;
continue;
}
if (is_special_day(ts_value.year(), ts_value.month(), ts_value.day())) {
ts_value.template set_time_unit<TimeUnit::DAY>(1);
res_data[i] = binary_cast<DateValueType, UInt32>(ts_value);
continue;
}
// day_of_week, from 1(Mon) to 7(Sun)
int day_of_week = ts_value.weekday() + 1;
int gap_of_monday = day_of_week - 1;
TimeInterval interval(DAY, gap_of_monday, true);
ts_value.template date_add_interval<DAY>(interval);
res_data[i] = binary_cast<DateValueType, UInt32>(ts_value);
} else {
const auto& cur_data = data_col[i];
auto ts_value = binary_cast<UInt64, DateValueType>(cur_data);
if (!ts_value.is_valid_date()) {
null_map[i] = 1;
continue;
}
if (is_special_day(ts_value.year(), ts_value.month(), ts_value.day())) {
ts_value.set_time(ts_value.year(), ts_value.month(), 1, 0, 0, 0, 0);
UInt64 cast_value = binary_cast<DateValueType, UInt64>(ts_value);
DataTypeDateTimeV2::cast_to_date_v2(cast_value, res_data[i]);
continue;
}
// day_of_week, from 1(Mon) to 7(Sun)
int day_of_week = ts_value.weekday() + 1;
int gap_of_monday = day_of_week - 1;
TimeInterval interval(DAY, gap_of_monday, true);
ts_value.template date_add_interval<DAY>(interval);
ts_value.set_time(ts_value.year(), ts_value.month(), ts_value.day(), 0, 0, 0, 0);
UInt64 cast_value = binary_cast<DateValueType, UInt64>(ts_value);
DataTypeDateTimeV2::cast_to_date_v2(cast_value, res_data[i]);
}
}
}
// specially, 1970-01-01, 1970-01-02, 1970-01-03 and 1970-01-04 return 1970-01-01
static bool is_special_day(int year, int month, int day) {
return year == 1970 && month == 1 && day > 0 && day < 5;
}
};
template <typename Impl>
class FunctionOtherTypesToDateType : public IFunction {
public:
@ -771,6 +889,10 @@ void register_function_timestamp(SimpleFunctionFactory& factory) {
factory.register_function<FunctionDateOrDateTimeToDate<LastDayImpl, DataTypeDate>>();
factory.register_function<FunctionDateOrDateTimeToDate<LastDayImpl, DataTypeDateV2>>();
factory.register_function<FunctionDateOrDateTimeToDate<LastDayImpl, DataTypeDateTimeV2>>();
factory.register_function<FunctionDateOrDateTimeToDate<MondayImpl, DataTypeDateV2>>();
factory.register_function<FunctionDateOrDateTimeToDate<MondayImpl, DataTypeDateTimeV2>>();
factory.register_function<FunctionDateOrDateTimeToDate<MondayImpl, DataTypeDate>>();
factory.register_function<FunctionDateOrDateTimeToDate<MondayImpl, DataTypeDateTime>>();
}
} // namespace doris::vectorized