Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

branch-2.1: [function](date) Support date trunc function #49540 #49661

Merged
merged 1 commit into from
Mar 29, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 37 additions & 14 deletions be/src/vec/functions/function_timestamp.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -403,7 +403,7 @@ struct DateTruncState {
Callback_function callback_function;
};

template <typename DateType>
template <typename DateType, bool DateArgIsFirst>
struct DateTrunc {
static constexpr auto name = "date_trunc";

Expand All @@ -416,7 +416,11 @@ struct DateTrunc {
static size_t get_number_of_arguments() { return 2; }

static DataTypes get_variadic_argument_types() {
return {std::make_shared<DateType>(), std::make_shared<DataTypeString>()};
if constexpr (DateArgIsFirst) {
return {std::make_shared<DateType>(), std::make_shared<DataTypeString>()};
} else {
return {std::make_shared<DataTypeString>(), std::make_shared<DateType>()};
}
}

static DataTypePtr get_return_type_impl(const DataTypes& arguments) {
Expand All @@ -427,11 +431,12 @@ struct DateTrunc {
if (scope != FunctionContext::THREAD_LOCAL) {
return Status::OK();
}
if (!context->is_col_constant(1)) {
if (!context->is_col_constant(DateArgIsFirst ? 1 : 0)) {
return Status::InvalidArgument(
"date_trunc function of time unit argument must be constant.");
}
const auto& data_str = context->get_constant_col(1)->column_ptr->get_data_at(0);
const auto& data_str =
context->get_constant_col(DateArgIsFirst ? 1 : 0)->column_ptr->get_data_at(0);
std::string lower_str(data_str.data, data_str.size);
std::transform(lower_str.begin(), lower_str.end(), lower_str.begin(),
[](unsigned char c) { return std::tolower(c); });
Expand Down Expand Up @@ -467,8 +472,8 @@ struct DateTrunc {
DCHECK_EQ(arguments.size(), 2);

auto null_map = ColumnUInt8::create(input_rows_count, 0);
const auto& datetime_column =
block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();
const auto& datetime_column = block.get_by_position(arguments[DateArgIsFirst ? 0 : 1])
.column->convert_to_full_column_if_const();
ColumnPtr res = ColumnType::create(input_rows_count);
auto* state = reinterpret_cast<DateTruncState*>(
context->get_function_state(FunctionContext::THREAD_LOCAL));
Expand Down Expand Up @@ -1354,10 +1359,14 @@ class FunctionOtherTypesToDateType : public IFunction {
}

Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
if constexpr (std::is_same_v<Impl, DateTrunc<DataTypeDate>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateV2>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateTime>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateTimeV2>>) {
if constexpr (std::is_same_v<Impl, DateTrunc<DataTypeDate, true>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateV2, true>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateTime, true>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateTimeV2, true>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDate, false>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateV2, false>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateTime, false>> ||
std::is_same_v<Impl, DateTrunc<DataTypeDateTimeV2, false>>) {
return Impl::open(context, scope);
} else {
return Status::OK();
Expand Down Expand Up @@ -1621,10 +1630,20 @@ using FunctionStrToDatetime = FunctionOtherTypesToDateType<StrToDate<DataTypeDat
using FunctionStrToDateV2 = FunctionOtherTypesToDateType<StrToDate<DataTypeDateV2>>;
using FunctionStrToDatetimeV2 = FunctionOtherTypesToDateType<StrToDate<DataTypeDateTimeV2>>;
using FunctionMakeDate = FunctionOtherTypesToDateType<MakeDateImpl>;
using FunctionDateTruncDate = FunctionOtherTypesToDateType<DateTrunc<DataTypeDate>>;
using FunctionDateTruncDateV2 = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateV2>>;
using FunctionDateTruncDatetime = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTime>>;
using FunctionDateTruncDatetimeV2 = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTimeV2>>;
using FunctionDateTruncDate = FunctionOtherTypesToDateType<DateTrunc<DataTypeDate, true>>;
using FunctionDateTruncDateV2 = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateV2, true>>;
using FunctionDateTruncDatetime = FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTime, true>>;
using FunctionDateTruncDatetimeV2 =
FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTimeV2, true>>;

using FunctionDateTruncDateWithCommonOrder =
FunctionOtherTypesToDateType<DateTrunc<DataTypeDate, false>>;
using FunctionDateTruncDateV2WithCommonOrder =
FunctionOtherTypesToDateType<DateTrunc<DataTypeDateV2, false>>;
using FunctionDateTruncDatetimeWithCommonOrder =
FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTime, false>>;
using FunctionDateTruncDatetimeV2WithCommonOrder =
FunctionOtherTypesToDateType<DateTrunc<DataTypeDateTimeV2, false>>;
using FunctionFromIso8601DateV2 = FunctionOtherTypesToDateType<FromIso8601DateV2>;

void register_function_timestamp(SimpleFunctionFactory& factory) {
Expand All @@ -1638,6 +1657,10 @@ void register_function_timestamp(SimpleFunctionFactory& factory) {
factory.register_function<FunctionDateTruncDateV2>();
factory.register_function<FunctionDateTruncDatetime>();
factory.register_function<FunctionDateTruncDatetimeV2>();
factory.register_function<FunctionDateTruncDateWithCommonOrder>();
factory.register_function<FunctionDateTruncDateV2WithCommonOrder>();
factory.register_function<FunctionDateTruncDatetimeWithCommonOrder>();
factory.register_function<FunctionDateTruncDatetimeV2WithCommonOrder>();
factory.register_function<FunctionFromIso8601DateV2>();

factory.register_function<FunctionUnixTimestamp<UnixTimeStampImpl>>();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1734,7 +1734,7 @@ && collectChildReturnTypes()[0].isDecimalV3()) {
final String constParam = ((StringLiteral) getChild(1)).getValue().toLowerCase();
if (!Lists.newArrayList("year", "quarter", "month", "week", "day", "hour", "minute", "second")
.contains(constParam)) {
throw new AnalysisException("date_trunc function second param only support argument is "
throw new AnalysisException("date_trunc function time unit param only support argument is "
+ "year|quarter|month|week|day|hour|minute|second");
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -358,6 +358,26 @@ public static Expression dateTrunc(DateV2Literal date, StringLikeLiteral trunc)
return DateV2Literal.fromJavaDateType(dateTruncHelper(date.toJavaDateType(), trunc.getValue()));
}

@ExecFunction(name = "date_trunc")
public static Expression dateTrunc(StringLikeLiteral trunc, DateTimeLiteral date) {
return DateTimeLiteral.fromJavaDateType(dateTruncHelper(date.toJavaDateType(), trunc.getValue()));
}

@ExecFunction(name = "date_trunc")
public static Expression dateTrunc(StringLikeLiteral trunc, DateTimeV2Literal date) {
return DateTimeV2Literal.fromJavaDateType(dateTruncHelper(date.toJavaDateType(), trunc.getValue()));
}

@ExecFunction(name = "date_trunc")
public static Expression dateTrunc(StringLikeLiteral trunc, DateLiteral date) {
return DateLiteral.fromJavaDateType(dateTruncHelper(date.toJavaDateType(), trunc.getValue()));
}

@ExecFunction(name = "date_trunc")
public static Expression dateTrunc(StringLikeLiteral trunc, DateV2Literal date) {
return DateV2Literal.fromJavaDateType(dateTruncHelper(date.toJavaDateType(), trunc.getValue()));
}

private static LocalDateTime dateTruncHelper(LocalDateTime dateTime, String trunc) {
int year = dateTime.getYear();
int month = dateTime.getMonthValue();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,38 +21,28 @@
import org.apache.doris.nereids.exceptions.AnalysisException;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.CustomSignature;
import org.apache.doris.nereids.trees.expressions.functions.Monotonic;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral;
import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.DateTimeType;
import org.apache.doris.nereids.types.DateTimeV2Type;
import org.apache.doris.nereids.types.DateType;
import org.apache.doris.nereids.types.DateV2Type;
import org.apache.doris.nereids.types.VarcharType;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;

import java.util.List;

/**
* ScalarFunction 'date_trunc'. This class is generated by GenerateFunction.
*/
public class DateTrunc extends ScalarFunction
implements BinaryExpression, ExplicitlyCastableSignature, AlwaysNullable, Monotonic {

public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
FunctionSignature.ret(DateTimeV2Type.SYSTEM_DEFAULT)
.args(DateTimeV2Type.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT),
FunctionSignature.ret(DateTimeType.INSTANCE).args(DateTimeType.INSTANCE, VarcharType.SYSTEM_DEFAULT),
FunctionSignature.ret(DateV2Type.INSTANCE)
.args(DateV2Type.INSTANCE, VarcharType.SYSTEM_DEFAULT),
FunctionSignature.ret(DateType.INSTANCE).args(DateType.INSTANCE, VarcharType.SYSTEM_DEFAULT)
);
implements BinaryExpression, AlwaysNullable, Monotonic, CustomSignature {
private static final List<String> LEGAL_TIME_UNIT =
ImmutableList.of("year", "quarter", "month", "week", "day", "hour", "minute", "second");

/**
* constructor with 2 arguments.
Expand All @@ -63,15 +53,27 @@ public DateTrunc(Expression arg0, Expression arg1) {

@Override
public void checkLegalityBeforeTypeCoercion() {
if (getArgument(1).isConstant() == false || !(getArgument(1) instanceof VarcharLiteral)) {
throw new AnalysisException("the second parameter of "
boolean firstArgIsStringLiteral =
getArgument(0).isConstant() && getArgument(0) instanceof StringLikeLiteral;
boolean secondArgIsStringLiteral =
getArgument(1).isConstant() && getArgument(1) instanceof StringLikeLiteral;
if (!firstArgIsStringLiteral && !secondArgIsStringLiteral) {
throw new AnalysisException("the time unit parameter of "
+ getName() + " function must be a string constant: " + toSql());
}
final String constParam = ((VarcharLiteral) getArgument(1)).getStringValue().toLowerCase();
if (!Lists.newArrayList("year", "quarter", "month", "week", "day", "hour", "minute", "second")
.contains(constParam)) {
throw new AnalysisException("date_trunc function second param only support argument is "
+ "year|quarter|month|week|day|hour|minute|second");
} else if (firstArgIsStringLiteral && secondArgIsStringLiteral) {
if (!LEGAL_TIME_UNIT.contains(((VarcharLiteral) getArgument(0)).getStringValue().toLowerCase())
&& !LEGAL_TIME_UNIT.contains(((VarcharLiteral) getArgument(1))
.getStringValue().toLowerCase())) {
throw new AnalysisException("date_trunc function time unit param only support argument is "
+ String.join("|", LEGAL_TIME_UNIT));
}
} else {
final String constParam = ((VarcharLiteral) getArgument(firstArgIsStringLiteral ? 0 : 1))
.getStringValue().toLowerCase();
if (!LEGAL_TIME_UNIT.contains(constParam)) {
throw new AnalysisException("date_trunc function time unit param only support argument is "
+ String.join("|", LEGAL_TIME_UNIT));
}
}
}

Expand All @@ -85,8 +87,36 @@ public DateTrunc withChildren(List<Expression> children) {
}

@Override
public List<FunctionSignature> getSignatures() {
return SIGNATURES;
public FunctionSignature customSignature() {
if (getArgument(0).getDataType().isDateLikeType()) {
return FunctionSignature.ret(getArgument(0).getDataType())
.args(getArgument(0).getDataType(), VarcharType.SYSTEM_DEFAULT);
} else if (getArgument(1).getDataType().isDateLikeType()) {
return FunctionSignature.ret(getArgument(1).getDataType())
.args(VarcharType.SYSTEM_DEFAULT, getArgument(1).getDataType());
}
boolean firstArgIsStringLiteral =
getArgument(0).isConstant() && getArgument(0) instanceof VarcharLiteral;
boolean secondArgIsStringLiteral =
getArgument(1).isConstant() && getArgument(1) instanceof VarcharLiteral;
if (firstArgIsStringLiteral && !secondArgIsStringLiteral) {
return FunctionSignature.ret(DateTimeV2Type.SYSTEM_DEFAULT)
.args(VarcharType.SYSTEM_DEFAULT, DateTimeV2Type.SYSTEM_DEFAULT);
} else if (!firstArgIsStringLiteral && secondArgIsStringLiteral) {
return FunctionSignature.ret(DateTimeV2Type.SYSTEM_DEFAULT)
.args(DateTimeV2Type.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT);
} else if (firstArgIsStringLiteral && secondArgIsStringLiteral) {
boolean timeUnitIsFirst = LEGAL_TIME_UNIT.contains(((VarcharLiteral) getArgument(0))
.getStringValue().toLowerCase());
return timeUnitIsFirst ? FunctionSignature.ret(DateTimeV2Type.SYSTEM_DEFAULT)
.args(VarcharType.SYSTEM_DEFAULT, DateTimeV2Type.SYSTEM_DEFAULT)
: FunctionSignature.ret(DateTimeV2Type.SYSTEM_DEFAULT)
.args(DateTimeV2Type.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT);
}
// if both of args are not constant, `checkLegalityBeforeTypeCoercion` will throw exception so just return
// a signature here.
return FunctionSignature.ret(DateTimeV2Type.SYSTEM_DEFAULT)
.args(VarcharType.SYSTEM_DEFAULT, DateTimeV2Type.SYSTEM_DEFAULT);
}

@Override
Expand All @@ -101,11 +131,12 @@ public boolean isPositive() {

@Override
public int getMonotonicFunctionChildIndex() {
return 0;
return getArgument(0).getDataType().isDateLikeType() ? 0 : 1;
}

@Override
public Expression withConstantArgs(Literal literal) {
return new DateTrunc(literal, child(1));
return getArgument(0).getDataType().isDateLikeType()
? new DateTrunc(literal, child(1)) : new DateTrunc(child(0), literal);
}
}
Loading
Loading