Skip to content

Commit

Permalink
function: Support push function year,day,datediff,datesub,castTimeAsS…
Browse files Browse the repository at this point in the history
…tring,concat_ws down to TiFlash. (pingcap#2084)
  • Loading branch information
LittleFall authored Jun 4, 2021
1 parent 1a16bdb commit 0ed0a07
Show file tree
Hide file tree
Showing 10 changed files with 800 additions and 19 deletions.
40 changes: 29 additions & 11 deletions dbms/src/Flash/Coprocessor/DAGExpressionAnalyzer.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -242,25 +242,41 @@ static String buildCastFunction(DAGExpressionAnalyzer * analyzer, const tipb::Ex
return buildCastFunctionInternal(analyzer, {name, type_expr_name}, false, expr.field_type(), actions);
}

static String buildDateAddFunction(DAGExpressionAnalyzer * analyzer, const tipb::Expr & expr, ExpressionActionsPtr & actions)
struct DateAdd
{
static constexpr auto name = "date_add";
static const std::unordered_map<String, String> unit_to_func_name_map;
};
const std::unordered_map<String, String> DateAdd::unit_to_func_name_map = {{"DAY", "addDays"}, {"WEEK", "addWeeks"}, {"MONTH", "addMonths"},
{"YEAR", "addYears"}, {"HOUR", "addHours"}, {"MINUTE", "addMinutes"}, {"SECOND", "addSeconds"}};
struct DateSub
{
static constexpr auto name = "date_sub";
static const std::unordered_map<String, String> unit_to_func_name_map;
};
const std::unordered_map<String, String> DateSub::unit_to_func_name_map
= {{"DAY", "subtractDays"}, {"WEEK", "subtractWeeks"}, {"MONTH", "subtractMonths"}, {"YEAR", "subtractYears"},
{"HOUR", "subtractHours"}, {"MINUTE", "subtractMinutes"}, {"SECOND", "subtractSeconds"}};

template <typename Impl>
static String buildDateAddOrSubFunction(DAGExpressionAnalyzer * analyzer, const tipb::Expr & expr, ExpressionActionsPtr & actions)
{

static const std::unordered_map<String, String> unit_to_func_name_map({{"DAY", "addDays"}, {"WEEK", "addWeeks"}, {"MONTH", "addMonths"},
{"YEAR", "addYears"}, {"HOUR", "addHours"}, {"MINUTE", "addMinutes"}, {"SECOND", "addSeconds"}});
if (expr.children_size() != 3)
{
throw TiFlashException("date add function requires three arguments", Errors::Coprocessor::BadRequest);
throw TiFlashException(std::string() + Impl::name + " function requires three arguments", Errors::Coprocessor::BadRequest);
}
String date_column = analyzer->getActions(expr.children(0), actions);
String delta_column = analyzer->getActions(expr.children(1), actions);
if (expr.children(2).tp() != tipb::ExprType::String)
{
throw TiFlashException("3rd argument of date add function must be string literal", Errors::Coprocessor::BadRequest);
throw TiFlashException(
std::string() + "3rd argument of " + Impl::name + " function must be string literal", Errors::Coprocessor::BadRequest);
}
String unit = expr.children(2).val();
if (unit_to_func_name_map.find(unit) == unit_to_func_name_map.end())
throw TiFlashException("date_add does not support unit " + unit + " yet.", Errors::Coprocessor::Unimplemented);
String func_name = unit_to_func_name_map.find(unit)->second;
if (Impl::unit_to_func_name_map.find(unit) == Impl::unit_to_func_name_map.end())
throw TiFlashException(
std::string() + Impl::name + " function does not support unit " + unit + " yet.", Errors::Coprocessor::Unimplemented);
String func_name = Impl::unit_to_func_name_map.find(unit)->second;
const auto & date_column_type = removeNullable(actions->getSampleBlock().getByName(date_column).type);
if (!date_column_type->isDateOrDateTime())
{
Expand Down Expand Up @@ -328,10 +344,12 @@ static String buildFunction(DAGExpressionAnalyzer * analyzer, const tipb::Expr &
static std::unordered_map<String, std::function<String(DAGExpressionAnalyzer *, const tipb::Expr &, ExpressionActionsPtr &)>>
function_builder_map({{"in", buildInFunction}, {"notIn", buildInFunction}, {"globalIn", buildInFunction},
{"globalNotIn", buildInFunction}, {"tidbIn", buildInFunction}, {"tidbNotIn", buildInFunction}, {"ifNull", buildIfNullFunction},
{"multiIf", buildMultiIfFunction}, {"tidb_cast", buildCastFunction}, {"date_add", buildDateAddFunction},
{"multiIf", buildMultiIfFunction}, {"tidb_cast", buildCastFunction},
{"and", buildLogicalFunction}, {"or", buildLogicalFunction}, {"xor", buildLogicalFunction}, {"not", buildLogicalFunction},
{"bitAnd", buildBitwiseFunction}, {"bitOr", buildBitwiseFunction}, {"bitXor", buildBitwiseFunction},
{"bitNot", buildBitwiseFunction}, {"leftUTF8", buildLeftUTF8Function}});
{"bitNot", buildBitwiseFunction}, {"leftUTF8", buildLeftUTF8Function},
{"date_add", buildDateAddOrSubFunction<DateAdd>}, {"date_sub", buildDateAddOrSubFunction<DateSub>}
});

DAGExpressionAnalyzer::DAGExpressionAnalyzer(std::vector<NameAndTypePair> && source_columns_, const Context & context_)
: source_columns(std::move(source_columns_)), context(context_), after_agg(false), implicit_cast_count(0)
Expand Down
12 changes: 6 additions & 6 deletions dbms/src/Flash/Coprocessor/DAGUtils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -798,7 +798,7 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({

{tipb::ScalarFuncSig::DateFormatSig, "dateFormat"},
//{tipb::ScalarFuncSig::DateLiteral, "cast"},
//{tipb::ScalarFuncSig::DateDiff, "cast"},
{tipb::ScalarFuncSig::DateDiff, "tidbDateDiff"},
//{tipb::ScalarFuncSig::NullTimeDiff, "cast"},
//{tipb::ScalarFuncSig::TimeStringTimeDiff, "cast"},
//{tipb::ScalarFuncSig::DurationDurationTimeDiff, "cast"},
Expand Down Expand Up @@ -838,7 +838,7 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
//{tipb::ScalarFuncSig::NowWithoutArg, "cast"},

//{tipb::ScalarFuncSig::DayName, "cast"},
//{tipb::ScalarFuncSig::DayOfMonth, "cast"},
{tipb::ScalarFuncSig::DayOfMonth, "toDayOfMonth"},
//{tipb::ScalarFuncSig::DayOfWeek, "cast"},
//{tipb::ScalarFuncSig::DayOfYear, "cast"},

Expand All @@ -847,7 +847,7 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
//{tipb::ScalarFuncSig::WeekDay, "cast"},
//{tipb::ScalarFuncSig::WeekOfYear, "cast"},

//{tipb::ScalarFuncSig::Year, "cast"},
{tipb::ScalarFuncSig::Year, "toYear"},
//{tipb::ScalarFuncSig::YearWeekWithMode, "cast"},
//{tipb::ScalarFuncSig::YearWeekWithoutMode, "cast"},

Expand Down Expand Up @@ -925,12 +925,12 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
{tipb::ScalarFuncSig::AddDateDatetimeInt, "date_add"},

//{tipb::ScalarFuncSig::SubDateStringString, "cast"},
//{tipb::ScalarFuncSig::SubDateStringInt, "cast"},
{tipb::ScalarFuncSig::SubDateStringInt, "date_sub"},
//{tipb::ScalarFuncSig::SubDateStringDecimal, "cast"},
//{tipb::ScalarFuncSig::SubDateIntString, "cast"},
//{tipb::ScalarFuncSig::SubDateIntInt, "cast"},
//{tipb::ScalarFuncSig::SubDateDatetimeString, "cast"},
//{tipb::ScalarFuncSig::SubDateDatetimeInt, "cast"},
{tipb::ScalarFuncSig::SubDateDatetimeInt, "date_sub"},

//{tipb::ScalarFuncSig::FromDays, "cast"},
//{tipb::ScalarFuncSig::TimeFormat, "cast"},
Expand All @@ -941,7 +941,7 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
//{tipb::ScalarFuncSig::ASCII, "cast"},
//{tipb::ScalarFuncSig::Char, "cast"},
{tipb::ScalarFuncSig::CharLengthUTF8, "lengthUTF8"}, {tipb::ScalarFuncSig::Concat, "tidbConcat"},
//{tipb::ScalarFuncSig::ConcatWS, "cast"},
{tipb::ScalarFuncSig::ConcatWS, "tidbConcatWS"},
//{tipb::ScalarFuncSig::Convert, "cast"},
//{tipb::ScalarFuncSig::Elt, "cast"},
//{tipb::ScalarFuncSig::ExportSet3Arg, "cast"},
Expand Down
1 change: 1 addition & 0 deletions dbms/src/Functions/FunctionsDateTime.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ void registerFunctionsDateTime(FunctionFactory & factory)
factory.registerFunction<FunctionDateDiff>(FunctionFactory::CaseInsensitive);
factory.registerFunction<FunctionTiDBTimestampDiff>();
factory.registerFunction<FunctionExtractMyDateTime>();
factory.registerFunction<FunctionTiDBDateDiff>();

factory.registerFunction<FunctionToTimeZone>();
}
Expand Down
174 changes: 172 additions & 2 deletions dbms/src/Functions/FunctionsDateTime.h
Original file line number Diff line number Diff line change
Expand Up @@ -434,8 +434,8 @@ struct ToDayOfMonthImpl
{
return time_zone.toDayOfMonth(DayNum_t(d));
}
static inline UInt8 execute(UInt64 , const DateLUTImpl & ) {
throw Exception("Illegal type MyTime of argument for function toDayOfMonth", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
static inline UInt8 execute(UInt64 t, const DateLUTImpl & ) {
return (UInt8)((t >> 41) & 31);
}

using FactorTransform = ToStartOfMonthImpl;
Expand Down Expand Up @@ -1654,6 +1654,176 @@ class FunctionTiDBTimestampDiff : public IFunction
}
};

/** TiDBDateDiff(t1, t2)
* Supports for tidb's dateDiff,
* returns t1 − t2 expressed as a value in days from one date to the other.
* Only the date parts of the values are used in the calculation.
*/
class FunctionTiDBDateDiff : public IFunction
{
public:
static constexpr auto name = "tidbDateDiff";
static FunctionPtr create(const Context &) { return std::make_shared<FunctionTiDBDateDiff>(); };

String getName() const override
{
return name;
}

bool isVariadic() const override { return false; }
size_t getNumberOfArguments() const override { return 2; }

DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
if(!removeNullable(arguments[0]).get()->isDateOrDateTime())
throw Exception("First argument for function " + getName() + " must be MyDate or MyDateTime",
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);

if(!removeNullable(arguments[1]).get()->isDateOrDateTime())
throw Exception("Second argument for function " + getName() + " must be MyDate or MyDateTime",
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);

// to align with tidb, dateDiff with zeroDate input should return null, so always return nullable type
return makeNullable(std::make_shared<DataTypeInt64>());
}

bool useDefaultImplementationForNulls() const override { return false; }
bool useDefaultImplementationForConstants() const override { return true; }

void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) override
{
bool has_nullable = false;
bool has_null_constant = false;
for(const auto & arg : arguments)
{
const auto & elem = block.getByPosition(arg);
has_nullable |= elem.type->isNullable();
has_null_constant |= elem.type->onlyNull();
}

if (has_null_constant)
{
block.getByPosition(result).column = block.getByPosition(result).type->createColumnConst(block.rows(), Null());
return;
}

ColumnPtr x_p = block.getByPosition(arguments[0]).column;
ColumnPtr y_p = block.getByPosition(arguments[1]).column;
if (has_nullable)
{
Block temporary_block = createBlockWithNestedColumns(block, arguments, result);
x_p = temporary_block.getByPosition(arguments[0]).column;
y_p = temporary_block.getByPosition(arguments[1]).column;
}

const IColumn & x = *x_p;
const IColumn & y = *y_p;

size_t rows = block.rows();
auto res = ColumnInt64::create(rows);
auto result_null_map = ColumnUInt8::create(rows);

dispatch(x, y, res->getData(), result_null_map->getData());

if (block.getByPosition(arguments[0]).type->isNullable()
|| block.getByPosition(arguments[1]).type->isNullable())
{
ColumnUInt8::Container &vec_result_null_map = result_null_map->getData();
ColumnPtr x_p = block.getByPosition(arguments[0]).column;
ColumnPtr y_p = block.getByPosition(arguments[1]).column;
for (size_t i = 0; i < rows; i++) {
vec_result_null_map[i] |= (x_p->isNullAt(i) || y_p->isNullAt(i));
}
}
block.getByPosition(result).column = ColumnNullable::create(std::move(res), std::move(result_null_map));
}
private:
void dispatch(const IColumn & x, const IColumn & y, ColumnInt64::Container & res, ColumnUInt8::Container & res_null_map)
{
auto * x_const = checkAndGetColumnConst<ColumnUInt64>(&x);
auto * y_const = checkAndGetColumnConst<ColumnUInt64>(&y);
if(x_const)
{
auto * y_vec = checkAndGetColumn<ColumnUInt64>(&y);
constant_vector(x_const->getValue<UInt64>(), *y_vec, res, res_null_map);
}
else if (y_const)
{
auto * x_vec = checkAndGetColumn<ColumnUInt64>(&x);
vector_constant(*x_vec, y_const->getValue<UInt64>(), res, res_null_map);
}
else
{
auto * x_vec = checkAndGetColumn<ColumnUInt64>(&x);
auto * y_vec = checkAndGetColumn<ColumnUInt64>(&y);
vector_vector(*x_vec, *y_vec, res, res_null_map);
}
}

void vector_vector(const ColumnVector<UInt64> & x, const ColumnVector<UInt64> & y,
ColumnInt64::Container & result, ColumnUInt8::Container & result_null_map)
{
const auto & x_data = x.getData();
const auto & y_data = y.getData();
for (size_t i = 0, size = x.size(); i < size; ++i)
{
result_null_map[i] = (x_data[i] == 0 || y_data[i] == 0);
if (!result_null_map[i])
result[i] = calculate(x_data[i], y_data[i]);
}
}

void vector_constant(const ColumnVector<UInt64> & x, UInt64 y,
ColumnInt64::Container & result, ColumnUInt8::Container & result_null_map)
{
const auto & x_data = x.getData();
if (y == 0)
{
for (size_t i = 0, size = x.size(); i < size; ++i)
result_null_map[i] = 1;
}
else
{
for (size_t i = 0, size = x.size(); i < size; ++i)
{
result_null_map[i] = (x_data[i] == 0);
if (!result_null_map[i])
result[i] = calculate(x_data[i], y);
}
}
}

void constant_vector(UInt64 x, const ColumnVector<UInt64> & y,
ColumnInt64::Container & result, ColumnUInt8::Container & result_null_map)
{
const auto & y_data = y.getData();
if (x == 0)
{
for (size_t i = 0, size = y.size(); i < size; ++i)
result_null_map[i] = 1;
}
else
{
for (size_t i = 0, size = y.size(); i < size; ++i) {
result_null_map[i] = (y_data[i] == 0);
if (!result_null_map[i])
result[i] = calculate(x, y_data[i]);
}
}
}

Int64 calculate(UInt64 x_packed, UInt64 y_packed)
{
MyDateTime x(x_packed);
MyDateTime y(y_packed);

Int64 days_x = calcDayNum(x.year, x.month, x.day);
Int64 days_y = calcDayNum(y.year, y.month, y.day);

return days_x - days_y;
}
};

/** dateDiff('unit', t1, t2, [timezone])
* t1 and t2 can be Date or DateTime
*
Expand Down
Loading

0 comments on commit 0ed0a07

Please sign in to comment.