expression: push unix_timestamp, concat to TiFlash (#2059)

dbms/src/Common/MyTime.cpp

@@ -797,26 +797,6 @@ String MyDateTime::toString(int fsp) const
 
 inline bool isZeroDate(UInt64 time) { return time == 0; }
 
-inline bool supportedByDateLUT(const MyDateTime & my_time) { return my_time.year >= 1970; }
-
-/// DateLUT only support time from year 1970, in some corner cases, the input date may be
-/// 1969-12-31, need extra logical to handle it
-inline time_t getEpochSecond(const MyDateTime & my_time, const DateLUTImpl & time_zone)
-{
-    if likely (supportedByDateLUT(my_time))
-        return time_zone.makeDateTime(my_time.year, my_time.month, my_time.day, my_time.hour, my_time.minute, my_time.second);
-    if likely (my_time.year == 1969 && my_time.month == 12 && my_time.day == 31)
-    {
-        /// - 3600 * 24 + my_time.hour * 3600 + my_time.minute * 60 + my_time.second is UTC based, need to adjust
-        /// the epoch according to the input time_zone
-        return -3600 * 24 + my_time.hour * 3600 + my_time.minute * 60 + my_time.second - time_zone.getOffsetAtStartEpoch();
-    }
-    else
-    {
-        throw Exception("Unsupported timestamp value , TiFlash only support timestamp after 1970-01-01 00:00:00 UTC)");
-    }
-}
-
 void convertTimeZone(UInt64 from_time, UInt64 & to_time, const DateLUTImpl & time_zone_from, const DateLUTImpl & time_zone_to)
 {
     if (isZeroDate(from_time))

dbms/src/Common/MyTime.h

@@ -163,6 +163,25 @@ int calcDayNum(int year, int month, int day);
 size_t maxFormattedDateTimeStringLength(const String & format);
 
 
+inline bool supportedByDateLUT(const MyDateTime & my_time) { return my_time.year >= 1970; }
+
+/// DateLUT only support time from year 1970, in some corner cases, the input date may be
+/// 1969-12-31, need extra logical to handle it
+inline time_t getEpochSecond(const MyDateTime & my_time, const DateLUTImpl & time_zone)
+{
+    if likely (supportedByDateLUT(my_time))
+        return time_zone.makeDateTime(my_time.year, my_time.month, my_time.day, my_time.hour, my_time.minute, my_time.second);
+    if likely (my_time.year == 1969 && my_time.month == 12 && my_time.day == 31)
+    {
+        /// - 3600 * 24 + my_time.hour * 3600 + my_time.minute * 60 + my_time.second is UTC based, need to adjust
+        /// the epoch according to the input time_zone
+        return -3600 * 24 + my_time.hour * 3600 + my_time.minute * 60 + my_time.second - time_zone.getOffsetAtStartEpoch();
+    }
+    else
+    {
+        throw Exception("Unsupported timestamp value , TiFlash only support timestamp after 1970-01-01 00:00:00 UTC)");
+    }
+}
 
 bool isPunctuation(char c);
 

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -888,8 +888,7 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
     //{tipb::ScalarFuncSig::SubDateAndString, "cast"},
 
     //{tipb::ScalarFuncSig::UnixTimestampCurrent, "cast"},
-    //{tipb::ScalarFuncSig::UnixTimestampInt, "cast"},
-    //{tipb::ScalarFuncSig::UnixTimestampDec, "cast"},
+    {tipb::ScalarFuncSig::UnixTimestampInt, "tidbUnixTimeStampInt"}, {tipb::ScalarFuncSig::UnixTimestampDec, "tidbUnixTimeStampDec"},
 
     //{tipb::ScalarFuncSig::ConvertTz, "cast"},
     //{tipb::ScalarFuncSig::MakeDate, "cast"},
@@ -941,8 +940,7 @@ std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
     //{tipb::ScalarFuncSig::Bin, "cast"},
     //{tipb::ScalarFuncSig::ASCII, "cast"},
     //{tipb::ScalarFuncSig::Char, "cast"},
-    {tipb::ScalarFuncSig::CharLengthUTF8, "lengthUTF8"},
-    //{tipb::ScalarFuncSig::Concat, "cast"},
+    {tipb::ScalarFuncSig::CharLengthUTF8, "lengthUTF8"}, {tipb::ScalarFuncSig::Concat, "tidbConcat"},
     //{tipb::ScalarFuncSig::ConcatWS, "cast"},
     //{tipb::ScalarFuncSig::Convert, "cast"},
     //{tipb::ScalarFuncSig::Elt, "cast"},

dbms/src/Functions/FunctionsConversion.cpp

@@ -27,6 +27,124 @@ void throwExceptionForIncompletelyParsedValue(
 }
 
 
+struct NameTiDBUnixTimeStampInt { static constexpr auto name = "tidbUnixTimeStampInt"; };
+struct NameTiDBUnixTimeStampDec { static constexpr auto name = "tidbUnixTimeStampDec"; };
+
+template <typename Name>
+class FunctionTiDBUnixTimeStamp : public IFunction
+{
+public:
+    static constexpr auto name = Name::name;
+    static FunctionPtr create(const Context & context) { return std::make_shared<FunctionTiDBUnixTimeStamp>(context); };
+    explicit FunctionTiDBUnixTimeStamp(const Context & context) : timezone_(context.getTimezoneInfo()){};
+
+    String getName() const override
+    {
+        return name;
+    }
+
+    size_t getNumberOfArguments() const override { return 1; }
+
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (!arguments[0].type->isMyDateOrMyDateTime())
+            throw Exception("The argument of function " + getName() + " must be date or datetime type", ErrorCodes::ILLEGAL_COLUMN);
+
+        if constexpr (std::is_same_v<Name, NameTiDBUnixTimeStampInt>)
+            return std::make_shared<DataTypeUInt64>();
+
+        int fsp = 0;
+        if (checkDataType<DataTypeMyDateTime>(arguments[0].type.get()))
+        {
+            auto & datetimeType = dynamic_cast<const DataTypeMyDateTime &>(*arguments[0].type);
+            fsp = datetimeType.getFraction();
+        }
+        return std::make_shared<DataTypeDecimal64>(12+fsp, fsp);
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, const size_t result) override
+    {
+        const auto & col_with_type_and_name = block.getByPosition(arguments[0]);
+
+        const auto * col_from = checkAndGetColumn<ColumnUInt64>(col_with_type_and_name.column.get());
+        const ColumnUInt64::Container & vec_from = col_from->getData();
+        size_t size = vec_from.size();
+
+        if constexpr (std::is_same_v<Name, NameTiDBUnixTimeStampInt>)
+        {
+            auto col_to = ColumnUInt64::create();
+            auto & vec_to = col_to->getData();
+            vec_to.resize(size);
+
+            for (size_t i = 0; i < size; i++)
+            {
+                UInt64 ret = 0;
+                if (getUnixTimeStampHelper(vec_from[i], ret))
+                    vec_to[i] = ret;
+                else
+                    vec_to[i] = 0;
+            }
+
+            block.getByPosition(result).column = std::move(col_to);
+        }
+        else /* if constexpr (std::is_same_v<Name, NameTiDBUnixTimeStampDec>) */
+        {
+            // Todo: speed up by `prepare`.
+            int fsp = 0, multiplier = 1, divider = 1'000'000;
+            if (checkDataType<DataTypeMyDateTime>(col_with_type_and_name.type.get()))
+            {
+                auto & datetimeType = dynamic_cast<const DataTypeMyDateTime &>(*col_with_type_and_name.type);
+                fsp = datetimeType.getFraction();
+            }
+            multiplier = getScaleMultiplier<Decimal64>(fsp);
+            divider = 1'000'000 / multiplier;
+
+            auto col_to = ColumnDecimal<Decimal64>::create(0, fsp);
+            auto & vec_to = col_to->getData();
+            vec_to.resize(size);
+
+            for (size_t i = 0; i < size; i++)
+            {
+                UInt64 ret = 0;
+                if (getUnixTimeStampHelper(vec_from[i], ret))
+                {
+                    MyDateTime datetime(vec_from[i]);
+                    vec_to[i] = ret * multiplier + datetime.micro_second / divider;
+                }
+                else
+                    vec_to[i] = 0;
+            }
+
+            block.getByPosition(result).column = std::move(col_to);
+        }
+    }
+
+private:
+    const TimezoneInfo & timezone_;
+
+    bool getUnixTimeStampHelper(UInt64 packed, UInt64 & ret)
+    {
+        static const auto lut_utc = DateLUT::instance("UTC");
+
+        if (timezone_.is_name_based)
+            convertTimeZone(packed, ret, *timezone_.timezone, lut_utc);
+        else
+            convertTimeZoneByOffset(packed, ret, -timezone_.timezone_offset, lut_utc);
+
+        try
+        {
+            ret = getEpochSecond(ret, lut_utc);
+        }
+        catch (...)
+        {
+            return false;
+        }
+        return true;
+    }
+};
+
 void registerFunctionsConversion(FunctionFactory & factory)
 {
     factory.registerFunction<FunctionToUInt8>();
@@ -91,6 +209,8 @@ void registerFunctionsConversion(FunctionFactory & factory)
 
     factory.registerFunction<FunctionFromUnixTime>();
     factory.registerFunction<FunctionDateFormat>();
+    factory.registerFunction<FunctionTiDBUnixTimeStamp<NameTiDBUnixTimeStampInt>>();
+    factory.registerFunction<FunctionTiDBUnixTimeStamp<NameTiDBUnixTimeStampDec>>();
     factory.registerFunction<FunctionStrToDate<NameStrToDateDate>>();
     factory.registerFunction<FunctionStrToDate<NameStrToDateDatetime>>();
 }

dbms/src/Functions/FunctionsConversion.h

@@ -1,44 +1,45 @@
 #pragma once
 
-#include <ext/enumerate.h>
-#include <ext/collection_cast.h>
-#include <ext/range.h>
-#include <type_traits>
-
-#include <IO/WriteBufferFromVector.h>
-#include <IO/ReadBufferFromMemory.h>
-#include <IO/Operators.h>
-#include <IO/parseDateTimeBestEffort.h>
-#include <DataTypes/DataTypeFactory.h>
-#include <DataTypes/DataTypesNumber.h>
-#include <DataTypes/DataTypeString.h>
-#include <DataTypes/DataTypeFixedString.h>
+#include <Columns/ColumnArray.h>
+#include <Columns/ColumnConst.h>
+#include <Columns/ColumnFixedString.h>
+#include <Columns/ColumnNullable.h>
+#include <Columns/ColumnString.h>
+#include <Columns/ColumnTuple.h>
+#include <Columns/ColumnsCommon.h>
+#include <Common/FieldVisitors.h>
+#include <Common/MyTime.h>
+#include <DataTypes/DataTypeArray.h>
 #include <DataTypes/DataTypeDate.h>
 #include <DataTypes/DataTypeDateTime.h>
+#include <DataTypes/DataTypeEnum.h>
+#include <DataTypes/DataTypeFactory.h>
+#include <DataTypes/DataTypeFixedString.h>
+#include <DataTypes/DataTypeInterval.h>
 #include <DataTypes/DataTypeMyDate.h>
 #include <DataTypes/DataTypeMyDateTime.h>
-#include <DataTypes/DataTypeEnum.h>
-#include <DataTypes/DataTypeArray.h>
-#include <DataTypes/DataTypeTuple.h>
-#include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeNothing.h>
+#include <DataTypes/DataTypeNullable.h>
+#include <DataTypes/DataTypeString.h>
+#include <DataTypes/DataTypeTuple.h>
 #include <DataTypes/DataTypeUUID.h>
-#include <DataTypes/DataTypeInterval.h>
-#include <Columns/ColumnString.h>
-#include <Columns/ColumnFixedString.h>
-#include <Columns/ColumnConst.h>
-#include <Columns/ColumnArray.h>
-#include <Columns/ColumnNullable.h>
-#include <Columns/ColumnTuple.h>
-#include <Columns/ColumnsCommon.h>
-#include <Common/FieldVisitors.h>
-#include <Interpreters/ExpressionActions.h>
-#include <Functions/IFunction.h>
-#include <Functions/FunctionsMiscellaneous.h>
-#include <Functions/FunctionsDateTime.h>
-#include <Functions/FunctionHelpers.h>
+#include <DataTypes/DataTypesNumber.h>
 #include <Functions/FunctionFactory.h>
+#include <Functions/FunctionHelpers.h>
+#include <Functions/FunctionsDateTime.h>
+#include <Functions/FunctionsMiscellaneous.h>
+#include <Functions/IFunction.h>
+#include <IO/Operators.h>
+#include <IO/ReadBufferFromMemory.h>
+#include <IO/WriteBufferFromVector.h>
+#include <IO/parseDateTimeBestEffort.h>
 #include <Interpreters/Context.h>
+#include <Interpreters/ExpressionActions.h>
+
+#include <ext/collection_cast.h>
+#include <ext/enumerate.h>
+#include <ext/range.h>
+#include <type_traits>
 
 
 namespace DB
@@ -1750,7 +1751,6 @@ class FunctionStrToDate : public IFunction
     }
 };
 
-
 /// Monotonicity.
 
 struct PositiveMonotonicity

dbms/src/Functions/FunctionsString.cpp

@@ -931,6 +931,65 @@ class ConcatImpl : public IFunction
     }
 };
 
+/** TiDB Function CONCAT(str1,str2,...)
+  * Returns the string that results from concatenating the arguments. May have one or more arguments.
+  * CONCAT() returns NULL if any argument is NULL.
+*/
+class FunctionTiDBConcat : public IFunction
+{
+private:
+    const Context & context;
+
+    struct NameTiDBConcat
+    {
+        static constexpr auto name = "tidbConcat";
+    };
+
+public:
+    static constexpr auto name = NameTiDBConcat::name;
+    FunctionTiDBConcat(const Context & context) : context(context) {}
+    static FunctionPtr create(const Context & context)
+    {
+        return std::make_shared<FunctionTiDBConcat>(context);
+    }
+
+    String getName() const override{ return name; }
+
+    bool isVariadic() const override{ return true; }
+    size_t getNumberOfArguments() const override{ return 0; }
+
+    bool useDefaultImplementationForNulls() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+    {
+        if (arguments.size() < 1)
+            throw Exception("Number of arguments for function " + getName() + " doesn't match: passed " + toString(arguments.size())
+                            + ", should be at least 1.",
+                            ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+
+        for (const auto arg_idx : ext::range(0, arguments.size()))
+        {
+            const auto & arg = arguments[arg_idx].get();
+            if (!arg->isStringOrFixedString())
+                throw Exception{
+                    "Illegal type " + arg->getName() + " of argument " + std::to_string(arg_idx + 1) + " of function " + getName(),
+                    ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+        }
+
+        return std::make_shared<DataTypeString>();
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, const size_t result) override
+    {
+        if (arguments.size() == 1)
+        {
+            const IColumn * c0 = block.getByPosition(arguments[0]).column.get();
+            block.getByPosition(result).column = c0->cloneResized(c0->size());
+        }
+        else
+            return ConcatImpl<NameTiDBConcat, false>(context).executeImpl(block, arguments, result);
+    }
+};
 
 class FunctionSubstring : public IFunction
 {
@@ -2753,6 +2812,7 @@ void registerFunctionsString(FunctionFactory & factory)
 	factory.registerFunction<FunctionRPadUTF8>();
     factory.registerFunction<FunctionConcat>();
     factory.registerFunction<FunctionConcatAssumeInjective>();
+    factory.registerFunction<FunctionTiDBConcat>();
     factory.registerFunction<FunctionSubstring>();
     factory.registerFunction<FunctionSubstringUTF8>();
     factory.registerFunction<FunctionAppendTrailingCharIfAbsent>();

dbms/src/Functions/IFunction.h

@@ -253,7 +253,7 @@ class IFunction : public std::enable_shared_from_this<IFunction>,
     /// TODO: make const
     void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) override = 0;
 
-    /// Override this functions to change default implementation behavior. See details in IMyFunction.
+    /// Override this functions to change default implementation behavior. See details in IPreparedFunction.
     bool useDefaultImplementationForNulls() const override { return true; }
     bool useDefaultImplementationForConstants() const override { return false; }
     ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {}; }