[Enhancement](function) Support date_trunc(date) and use it in auto partition (#24341)

Support date_trunc(date) and use it in auto partition
This commit is contained in:
zclllyybb
2023-09-14 16:53:09 +08:00
committed by GitHub
parent f303a99250
commit 4fbb25bc55
19 changed files with 371 additions and 44 deletions

View File

@ -33,6 +33,7 @@
#include "runtime/types.h"
#include "udf/udf.h"
#include "util/binary_cast.hpp"
#include "util/datetype_cast.hpp"
#include "util/time_lut.h"
#include "vec/aggregate_functions/aggregate_function.h"
#include "vec/columns/column.h"
@ -368,26 +369,22 @@ private:
}
};
template <typename DateValueType, typename ArgType>
template <typename DateType>
struct DateTrunc {
static constexpr auto name = "date_trunc";
using ColumnType = date_cast::DateToColumnV<DateType>;
using DateValueType = date_cast::DateToDateValueTypeV<DateType>;
using ArgType = date_cast::ValueTypeOfDateColumnV<ColumnType>;
static bool is_variadic() { return true; }
static DataTypes get_variadic_argument_types() {
if constexpr (std::is_same_v<DateValueType, VecDateTimeValue>) {
return {std::make_shared<DataTypeDateTime>(), std::make_shared<DataTypeString>()};
} else {
return {std::make_shared<DataTypeDateTimeV2>(), std::make_shared<DataTypeString>()};
}
return {std::make_shared<DateType>(), std::make_shared<DataTypeString>()};
}
static DataTypePtr get_return_type_impl(const DataTypes& arguments) {
if constexpr (std::is_same_v<DateValueType, VecDateTimeValue>) {
return make_nullable(std::make_shared<DataTypeDateTime>());
} else {
return make_nullable(std::make_shared<DataTypeDateTimeV2>());
}
return make_nullable(std::make_shared<DateType>());
}
static Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
@ -404,22 +401,21 @@ struct DateTrunc {
std::tie(argument_columns[1], col_const[1]) =
unpack_if_const(block.get_by_position(arguments[1]).column);
auto datetime_column = static_cast<const ColumnVector<ArgType>*>(argument_columns[0].get());
auto datetime_column = static_cast<const ColumnType*>(argument_columns[0].get());
auto str_column = static_cast<const ColumnString*>(argument_columns[1].get());
auto& rdata = str_column->get_chars();
auto& roffsets = str_column->get_offsets();
ColumnPtr res = ColumnVector<ArgType>::create();
ColumnPtr res = ColumnType::create();
if (col_const[1]) {
execute_impl_right_const(
datetime_column->get_data(), str_column->get_data_at(0),
static_cast<ColumnVector<ArgType>*>(res->assume_mutable().get())->get_data(),
static_cast<ColumnType*>(res->assume_mutable().get())->get_data(),
null_map->get_data(), input_rows_count);
} else {
execute_impl(
datetime_column->get_data(), rdata, roffsets,
static_cast<ColumnVector<ArgType>*>(res->assume_mutable().get())->get_data(),
null_map->get_data(), input_rows_count);
execute_impl(datetime_column->get_data(), rdata, roffsets,
static_cast<ColumnType*>(res->assume_mutable().get())->get_data(),
null_map->get_data(), input_rows_count);
}
block.get_by_position(result).column = ColumnNullable::create(res, std::move(null_map));
@ -1269,16 +1265,19 @@ public:
using FunctionStrToDate = FunctionOtherTypesToDateType<StrToDate>;
using FunctionMakeDate = FunctionOtherTypesToDateType<MakeDateImpl>;
using FunctionDateTrunc = FunctionOtherTypesToDateType<DateTrunc<VecDateTimeValue, Int64>>;
using FunctionDateTruncV2 =
FunctionOtherTypesToDateType<DateTrunc<DateV2Value<DateTimeV2ValueType>, UInt64>>;
using FunctionDateTruncDate = FunctionOtherTypesToDateType<DateTrunc<DataTypeDate>>;
using FunctionDateTruncDateV2 = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateV2>>;
using FunctionDateTruncDatetime = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTime>>;
using FunctionDateTruncDatetimeV2 = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTimeV2>>;
void register_function_timestamp(SimpleFunctionFactory& factory) {
factory.register_function<FunctionStrToDate>();
factory.register_function<FunctionMakeDate>();
factory.register_function<FromDays>();
factory.register_function<FunctionDateTrunc>();
factory.register_function<FunctionDateTruncV2>();
factory.register_function<FunctionDateTruncDate>();
factory.register_function<FunctionDateTruncDateV2>();
factory.register_function<FunctionDateTruncDatetime>();
factory.register_function<FunctionDateTruncDatetimeV2>();
factory.register_function<FunctionUnixTimestamp<UnixTimeStampImpl>>();
factory.register_function<FunctionUnixTimestamp<UnixTimeStampDateImpl<DataTypeDate>>>();