support last_day and dayofmonth pushdown to tiflash (#4183)

close #4149

dbms/CMakeLists.txt

@@ -177,6 +177,7 @@ target_link_libraries (clickhouse_common_io
     cpptoml
 )
 target_include_directories (clickhouse_common_io BEFORE PRIVATE ${kvClient_SOURCE_DIR}/include)
+target_include_directories (clickhouse_common_io BEFORE PUBLIC ${kvproto_SOURCE_DIR} ${tipb_SOURCE_DIR} ${Protobuf_INCLUDE_DIR} ${gRPC_INCLUDE_DIRS})
 
 target_link_libraries (dbms
     clickhouse_parsers

dbms/src/Common/MyTime.cpp

@@ -523,15 +523,7 @@ bool checkTimeValid(Int32 year, Int32 month, Int32 day, Int32 hour, Int32 minute
     {
         return false;
     }
-    static int days_of_month_table[] = {0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31};
-    if (month != 2)
-        return day <= days_of_month_table[month];
-    bool is_leap_year = false;
-    if ((year & 0b0011) == 0)
-    {
-        is_leap_year = year % 100 != 0 || year % 400 == 0;
-    }
-    return day <= (is_leap_year ? 29 : 28);
+    return day <= getLastDay(year, month);
 }
 
 std::pair<Field, bool> parseMyDateTimeAndJudgeIsDate(const String & str, int8_t fsp, bool needCheckTimeValid)

dbms/src/Common/MyTime.h

@@ -192,4 +192,21 @@ bool isValidSeperator(char c, int previous_parts);
 // Note that this function will not check if the input is logically a valid datetime value.
 bool toCoreTimeChecked(const UInt64 & year, const UInt64 & month, const UInt64 & day, const UInt64 & hour, const UInt64 & minute, const UInt64 & second, const UInt64 & microsecond, MyDateTime & result);
 
+inline bool isLeapYear(UInt16 year)
+{
+    return ((year % 4 == 0) && (year % 100 != 0)) || (year % 400 == 0);
+}
+
+// Get last day of a month. Return 0 if month if invalid.
+inline UInt8 getLastDay(UInt16 year, UInt8 month)
+{
+    static constexpr UInt8 days_of_month_table[] = {0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31};
+    UInt8 last_day = 0;
+    if (month > 0 && month <= 12)
+        last_day = days_of_month_table[month];
+    if (month == 2 && isLeapYear(year))
+        last_day = 29;
+    return last_day;
+}
+
 } // namespace DB

dbms/src/Flash/Coprocessor/DAGContext.h

@@ -231,6 +231,27 @@ class DAGContext
         return (flags & f);
     }
 
+    UInt64 getSQLMode() const
+    {
+        return sql_mode;
+    }
+    void setSQLMode(UInt64 f)
+    {
+        sql_mode = f;
+    }
+    void addSQLMode(UInt64 f)
+    {
+        sql_mode |= f;
+    }
+    void delSQLMode(UInt64 f)
+    {
+        sql_mode &= (~f);
+    }
+    bool hasSQLMode(UInt64 f) const
+    {
+        return sql_mode & f;
+    }
+
     void initExchangeReceiverIfMPP(Context & context, size_t max_streams);
     const std::unordered_map<String, std::shared_ptr<ExchangeReceiver>> & getMPPExchangeReceiverMap() const;
 

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -560,7 +560,7 @@ const std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
     //{tipb::ScalarFuncSig::Timestamp2Args, "cast"},
     //{tipb::ScalarFuncSig::TimestampLiteral, "cast"},
 
-    //{tipb::ScalarFuncSig::LastDay, "cast"},
+    {tipb::ScalarFuncSig::LastDay, "tidbLastDay"},
     {tipb::ScalarFuncSig::StrToDateDate, "strToDateDate"},
     {tipb::ScalarFuncSig::StrToDateDatetime, "strToDateDatetime"},
     // {tipb::ScalarFuncSig::StrToDateDuration, "cast"},

dbms/src/Functions/FunctionsDateTime.cpp

@@ -121,6 +121,7 @@ void registerFunctionsDateTime(FunctionFactory & factory)
     factory.registerFunction<FunctionTiDBDateDiff>();
 
     factory.registerFunction<FunctionToTimeZone>();
+    factory.registerFunction<FunctionToLastDay>();
 }
 
 } // namespace DB

dbms/src/Functions/FunctionsDateTime.h

@@ -15,6 +15,7 @@
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeString.h>
 #include <DataTypes/DataTypesNumber.h>
+#include <Flash/Coprocessor/DAGContext.h>
 #include <Functions/FunctionHelpers.h>
 #include <Functions/IFunction.h>
 #include <IO/WriteHelpers.h>
@@ -3241,6 +3242,118 @@ class FunctionDateTimeToString : public IFunction
     const Context & context;
 };
 
+template <typename ToFieldType>
+struct TiDBLastDayTransformerImpl
+{
+    static_assert(std::is_same_v<ToFieldType, DataTypeMyDate::FieldType>);
+    static constexpr auto name = "tidbLastDay";
+
+    static void execute(const Context & context,
+                        const ColumnVector<DataTypeMyTimeBase::FieldType>::Container & vec_from,
+                        typename ColumnVector<ToFieldType>::Container & vec_to,
+                        typename ColumnVector<UInt8>::Container & vec_null_map)
+    {
+        for (size_t i = 0; i < vec_from.size(); ++i)
+        {
+            bool is_null = false;
+            MyTimeBase val(vec_from[i]);
+            vec_to[i] = execute(context, val, is_null);
+            vec_null_map[i] = is_null;
+        }
+    }
+
+    static ToFieldType execute(const Context & context, const MyTimeBase & val, bool & is_null)
+    {
+        // TiDB also considers NO_ZERO_DATE sql_mode. But sql_mode is not handled by TiFlash for now.
+        if (val.month == 0 || val.day == 0)
+        {
+            context.getDAGContext()->handleInvalidTime(
+                fmt::format("Invalid time value: month({}) or day({}) is zero", val.month, val.day),
+                Errors::Types::WrongValue);
+            is_null = true;
+            return 0;
+        }
+        UInt8 last_day = getLastDay(val.year, val.month);
+        return MyDate(val.year, val.month, last_day).toPackedUInt();
+    }
+};
+
+// Similar to FunctionDateOrDateTimeToSomething, but also handle nullable result and mysql sql mode.
+template <typename ToDataType, template <typename> class Transformer, bool return_nullable>
+class FunctionMyDateOrMyDateTimeToSomething : public IFunction
+{
+private:
+    const Context & context;
+
+public:
+    using ToFieldType = typename ToDataType::FieldType;
+    static constexpr auto name = Transformer<ToFieldType>::name;
+
+    explicit FunctionMyDateOrMyDateTimeToSomething(const Context & context)
+        : context(context)
+    {}
+    static FunctionPtr create(const Context & context) { return std::make_shared<FunctionMyDateOrMyDateTimeToSomething>(context); };
+
+    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(
+                fmt::format("Illegal type {} of argument of function {}. Should be a date or a date with time", arguments[0].type->getName(), getName()),
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+
+        DataTypePtr return_type = std::make_shared<ToDataType>();
+        if constexpr (return_nullable)
+            return_type = makeNullable(return_type);
+        return return_type;
+    }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) const override
+    {
+        const DataTypePtr & from_type = block.getByPosition(arguments[0]).type;
+
+        if (from_type->isMyDateOrMyDateTime())
+        {
+            using FromFieldType = typename DataTypeMyTimeBase::FieldType;
+
+            const ColumnVector<FromFieldType> * col_from
+                = checkAndGetColumn<ColumnVector<FromFieldType>>(block.getByPosition(arguments[0]).column.get());
+            const typename ColumnVector<FromFieldType>::Container & vec_from = col_from->getData();
+
+            const size_t size = vec_from.size();
+            auto col_to = ColumnVector<ToFieldType>::create(size);
+            typename ColumnVector<ToFieldType>::Container & vec_to = col_to->getData();
+
+            if constexpr (return_nullable)
+            {
+                ColumnUInt8::MutablePtr col_null_map = ColumnUInt8::create(size, 0);
+                ColumnUInt8::Container & vec_null_map = col_null_map->getData();
+                Transformer<ToFieldType>::execute(context, vec_from, vec_to, vec_null_map);
+                block.getByPosition(result).column = ColumnNullable::create(std::move(col_to), std::move(col_null_map));
+            }
+            else
+            {
+                Transformer<ToFieldType>::execute(context, vec_from, vec_to);
+                block.getByPosition(result).column = std::move(col_to);
+            }
+        }
+        else
+            throw Exception(
+                fmt::format("Illegal type {} of argument of function {}", block.getByPosition(arguments[0]).type->getName(), getName()),
+                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
+    }
+};
+
+static constexpr bool return_nullable = true;
+static constexpr bool return_not_null = false;
+
 using FunctionToYear = FunctionDateOrDateTimeToSomething<DataTypeUInt16, ToYearImpl>;
 using FunctionToQuarter = FunctionDateOrDateTimeToSomething<DataTypeUInt8, ToQuarterImpl>;
 using FunctionToMonth = FunctionDateOrDateTimeToSomething<DataTypeUInt8, ToMonthImpl>;
@@ -3259,6 +3372,7 @@ using FunctionToStartOfFiveMinute = FunctionDateOrDateTimeToSomething<DataTypeDa
 using FunctionToStartOfFifteenMinutes = FunctionDateOrDateTimeToSomething<DataTypeDateTime, ToStartOfFifteenMinutesImpl>;
 using FunctionToStartOfHour = FunctionDateOrDateTimeToSomething<DataTypeDateTime, ToStartOfHourImpl>;
 using FunctionToTime = FunctionDateOrDateTimeToSomething<DataTypeDateTime, ToTimeImpl>;
+using FunctionToLastDay = FunctionMyDateOrMyDateTimeToSomething<DataTypeMyDate, TiDBLastDayTransformerImpl, return_nullable>;
 
 using FunctionToRelativeYearNum = FunctionDateOrDateTimeToSomething<DataTypeUInt16, ToRelativeYearNumImpl>;
 using FunctionToRelativeQuarterNum = FunctionDateOrDateTimeToSomething<DataTypeUInt32, ToRelativeQuarterNumImpl>;

dbms/src/Functions/tests/gtest_datetime_last_day.cpp

@@ -0,0 +1,103 @@
+#include <Columns/ColumnConst.h>
+#include <Common/Exception.h>
+#include <Functions/FunctionFactory.h>
+#include <Functions/FunctionsDateTime.h>
+#include <Interpreters/Context.h>
+#include <TestUtils/FunctionTestUtils.h>
+#include <TestUtils/TiFlashTestBasic.h>
+
+#include <string>
+#include <vector>
+
+namespace DB
+{
+namespace tests
+{
+class TestLastDay : public DB::tests::FunctionTest
+{
+};
+
+TEST_F(TestLastDay, BasicTest)
+try
+{
+    const String func_name = TiDBLastDayTransformerImpl<DataTypeMyDate::FieldType>::name;
+
+    // Ignore invalid month error
+    DAGContext * dag_context = context.getDAGContext();
+    UInt64 ori_flags = dag_context->getFlags();
+    dag_context->addFlag(TiDBSQLFlags::TRUNCATE_AS_WARNING);
+    dag_context->clearWarnings();
+
+    // nullable column test
+    ASSERT_COLUMN_EQ(
+        createColumn<Nullable<MyDate>>({MyDate{2001, 2, 28}.toPackedUInt(),
+                                        MyDate{2000, 2, 29}.toPackedUInt(),
+                                        MyDate{2000, 6, 30}.toPackedUInt(),
+                                        MyDate{2000, 5, 31}.toPackedUInt(),
+                                        {}}),
+        executeFunction(func_name,
+                        {createColumn<MyDate>({MyDate{2001, 2, 10}.toPackedUInt(),
+                                               MyDate{2000, 2, 10}.toPackedUInt(),
+                                               MyDate{2000, 6, 10}.toPackedUInt(),
+                                               MyDate{2000, 5, 10}.toPackedUInt(),
+                                               MyDate{2000, 0, 10}.toPackedUInt()})}));
+
+    ASSERT_COLUMN_EQ(
+        createColumn<Nullable<MyDate>>({MyDate{2001, 2, 28}.toPackedUInt(),
+                                        MyDate{2000, 2, 29}.toPackedUInt(),
+                                        MyDate{2000, 6, 30}.toPackedUInt(),
+                                        MyDate{2000, 5, 31}.toPackedUInt(),
+                                        {}}),
+        executeFunction(func_name,
+                        {createColumn<MyDateTime>({MyDateTime{2001, 2, 10, 10, 10, 10, 0}.toPackedUInt(),
+                                                   MyDateTime{2000, 2, 10, 10, 10, 10, 0}.toPackedUInt(),
+                                                   MyDateTime{2000, 6, 10, 10, 10, 10, 0}.toPackedUInt(),
+                                                   MyDateTime{2000, 5, 10, 10, 10, 10, 0}.toPackedUInt(),
+                                                   MyDateTime{2000, 0, 10, 10, 10, 10, 0}.toPackedUInt()})}));
+
+    // const test
+    UInt64 input[] = {
+        MyDateTime{2001, 2, 10, 10, 10, 10, 0}.toPackedUInt(),
+        MyDateTime{2000, 2, 10, 10, 10, 10, 0}.toPackedUInt(),
+        MyDateTime{2000, 6, 10, 10, 10, 10, 0}.toPackedUInt(),
+        MyDateTime{2000, 5, 10, 10, 10, 10, 0}.toPackedUInt(),
+    };
+
+    UInt64 output[] = {
+        MyDate{2001, 2, 28}.toPackedUInt(),
+        MyDate{2000, 2, 29}.toPackedUInt(),
+        MyDate{2000, 6, 30}.toPackedUInt(),
+        MyDate{2000, 5, 31}.toPackedUInt(),
+    };
+
+    for (size_t i = 0; i < sizeof(input) / sizeof(UInt64); ++i)
+    {
+        ASSERT_COLUMN_EQ(
+            createConstColumn<Nullable<MyDate>>(3, output[i]),
+            executeFunction(func_name,
+                            {createConstColumn<MyDate>(3, input[i])}));
+
+        ASSERT_COLUMN_EQ(
+            createConstColumn<Nullable<MyDate>>(3, output[i]),
+            executeFunction(func_name,
+                            {createConstColumn<Nullable<MyDate>>(3, input[i])}));
+    }
+
+    // const nullable test
+    ASSERT_COLUMN_EQ(
+        createConstColumn<Nullable<MyDate>>(3, {}),
+        executeFunction(func_name,
+                        {createConstColumn<Nullable<MyDate>>(3, {})}));
+
+    // special const test, month is zero.
+    ASSERT_COLUMN_EQ(
+        createConstColumn<Nullable<MyDate>>(3, {}),
+        executeFunction(func_name,
+                        {createConstColumn<MyDate>(3, MyDateTime{2000, 0, 10, 10, 10, 10, 0}.toPackedUInt())}));
+
+    dag_context->setFlags(ori_flags);
+}
+CATCH
+
+} // namespace tests
+} // namespace DB

tests/fullstack-test/expr/day_of_month.test

@@ -0,0 +1,27 @@
+mysql> drop table if exists test.t1;
+mysql> create table test.t1 (c_str varchar(100), c_datetime datetime(4), c_date date);
+mysql> insert into test.t1 values('', '1999-10-10 10:10:10.123', '1999-01-10'), ('200', '1999-02-10 10:10:10.123', '1999-11-10'), ('1999-30-10', '1999-10-10 10:10:10.123', '1999-01-10'), ('1999-01-10', '1999-10-10 10:10:10.123', '1999-01-10');
+mysql> alter table test.t1 set tiflash replica 1;
+func> wait_table test t1
+
+mysql> set @@tidb_isolation_read_engines='tiflash'; set @@tidb_enforce_mpp = 1; select dayofmonth();
+ERROR 1582 (42000) at line 1: Incorrect parameter count in the call to native function 'dayofmonth'
+
+# invalid input
+mysql> set @@tidb_isolation_read_engines='tiflash'; set @@tidb_enforce_mpp = 1; select dayofmonth(''), dayofmonth('1'), dayofmonth('1999-30-01'), dayofmonth(null);
++----------------+-----------------+--------------------------+------------------+
+| dayofmonth('') | dayofmonth('1') | dayofmonth('1999-30-01') | dayofmonth(null) |
++----------------+-----------------+--------------------------+------------------+
+|           NULL |            NULL |                     NULL |             NULL |
++----------------+-----------------+--------------------------+------------------+
+
+# got bug: https://github.com/pingcap/tics/issues/4186
+# mysql> set @@tidb_isolation_read_engines='tiflash'; set @@tidb_enforce_mpp = 1; select dayofmonth(c_str), dayofmonth(c_datetime), dayofmonth(c_date) from test.t1 order by 1, 2, 3;
+# +-------------------+------------------------+--------------------+
+# | dayofmonth(c_str) | dayofmonth(c_datetime) | dayofmonth(c_date) |
+# +-------------------+------------------------+--------------------+
+# |              NULL |                     10 |                 10 |
+# |              NULL |                     10 |                 10 |
+# |              NULL |                     10 |                 10 |
+# |                10 |                     10 |                 10 |
+# +-------------------+------------------------+--------------------+