[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:
@ -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>>>();
|
||||
|
||||
Reference in New Issue
Block a user