Support weekofyear pushdown to tiflash (#4791)

close #4677

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -516,7 +516,7 @@ const std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
     //{tipb::ScalarFuncSig::WeekWithMode, "cast"},
     //{tipb::ScalarFuncSig::WeekWithoutMode, "cast"},
     //{tipb::ScalarFuncSig::WeekDay, "cast"},
-    //{tipb::ScalarFuncSig::WeekOfYear, "cast"},
+    {tipb::ScalarFuncSig::WeekOfYear, "tidbWeekOfYear"},
 
     {tipb::ScalarFuncSig::Year, "toYear"},
     //{tipb::ScalarFuncSig::YearWeekWithMode, "cast"},

dbms/src/Functions/FunctionsDateTime.cpp

@@ -136,6 +136,7 @@ void registerFunctionsDateTime(FunctionFactory & factory)
     factory.registerFunction<FunctionTiDBDateDiff>();
     factory.registerFunction<FunctionToTiDBDayOfWeek>();
     factory.registerFunction<FunctionToTiDBDayOfYear>();
+    factory.registerFunction<FunctionToTiDBWeekOfYear>();
 
     factory.registerFunction<FunctionToTimeZone>();
     factory.registerFunction<FunctionToLastDay>();

dbms/src/Functions/FunctionsDateTime.h

@@ -3238,6 +3238,45 @@ struct TiDBDayOfYearTransformerImpl
         return static_cast<ToFieldType>(val.yearDay());
     }
 };
+
+template <typename ToFieldType>
+struct TiDBWeekOfYearTransformerImpl
+{
+    static constexpr auto name = "tidbWeekOfYear";
+
+    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)
+    {
+        bool is_null = false;
+        for (size_t i = 0; i < vec_from.size(); ++i)
+        {
+            MyTimeBase val(vec_from[i]);
+            vec_to[i] = execute(context, val, is_null);
+            vec_null_map[i] = is_null;
+            is_null = false;
+        }
+    }
+
+    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;
+        }
+        /// Behavior differences from TiDB:
+        /// for '0000-01-02', weekofyear is the same with MySQL, while TiDB is offset by one day
+        /// TiDB_weekofyear('0000-01-02') = 52, MySQL/TiFlash_weekofyear('0000-01-02') = 1
+        return static_cast<ToFieldType>(val.week(3));
+    }
+};
+
 // 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
@@ -3336,6 +3375,7 @@ using FunctionToTime = FunctionDateOrDateTimeToSomething<DataTypeDateTime, ToTim
 using FunctionToLastDay = FunctionMyDateOrMyDateTimeToSomething<DataTypeMyDate, TiDBLastDayTransformerImpl, return_nullable>;
 using FunctionToTiDBDayOfWeek = FunctionMyDateOrMyDateTimeToSomething<DataTypeUInt16, TiDBDayOfWeekTransformerImpl, return_nullable>;
 using FunctionToTiDBDayOfYear = FunctionMyDateOrMyDateTimeToSomething<DataTypeUInt16, TiDBDayOfYearTransformerImpl, return_nullable>;
+using FunctionToTiDBWeekOfYear = FunctionMyDateOrMyDateTimeToSomething<DataTypeUInt16, TiDBWeekOfYearTransformerImpl, return_nullable>;
 
 using FunctionToRelativeYearNum = FunctionDateOrDateTimeToSomething<DataTypeUInt16, ToRelativeYearNumImpl>;
 using FunctionToRelativeQuarterNum = FunctionDateOrDateTimeToSomething<DataTypeUInt32, ToRelativeQuarterNumImpl>;

dbms/src/Functions/tests/gtest_weekofyear.cpp

@@ -0,0 +1,121 @@
+// Copyright 2022 PingCAP, Ltd.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <Common/MyTime.h>
+#include <Functions/FunctionFactory.h>
+#include <TestUtils/FunctionTestUtils.h>
+#include <TestUtils/TiFlashTestBasic.h>
+
+#include <string>
+#include <vector>
+
+namespace DB::tests
+{
+class TestWeekOfYear : public DB::tests::FunctionTest
+{
+};
+
+TEST_F(TestWeekOfYear, WeekOfYear)
+try
+{
+    DAGContext * dag_context = context.getDAGContext();
+    UInt64 ori_flags = dag_context->getFlags();
+    dag_context->addFlag(TiDBSQLFlags::TRUNCATE_AS_WARNING);
+    /// ColumnVector(nullable)
+    const String func_name = "tidbWeekOfYear";
+    static auto const nullable_datetime_type_ptr = makeNullable(std::make_shared<DataTypeMyDateTime>(6));
+    static auto const datetime_type_ptr = std::make_shared<DataTypeMyDateTime>(6);
+    static auto const date_type_ptr = std::make_shared<DataTypeMyDate>();
+    auto data_col_ptr = createColumn<Nullable<DataTypeMyDateTime::FieldType>>(
+                            {
+                                {}, // Null
+                                // FIXME: https://github.com/pingcap/tiflash/issues/4186
+                                // MyDateTime(2022, 12, 0, 1, 1, 1, 1).toPackedUInt(),
+                                // MyDateTime(2022, 13, 31, 1, 1, 1, 1).toPackedUInt(),
+                                MyDateTime(0, 0, 0, 0, 0, 0, 0).toPackedUInt(),
+                                MyDateTime(0, 1, 1, 0, 0, 0, 0).toPackedUInt(),
+                                MyDateTime(0, 1, 2, 0, 0, 0, 0).toPackedUInt(),
+                                MyDateTime(0, 1, 3, 0, 0, 0, 0).toPackedUInt(),
+                                MyDateTime(1969, 1, 1, 1, 1, 1, 1).toPackedUInt(),
+                                MyDateTime(1969, 1, 6, 1, 1, 1, 1).toPackedUInt(),
+                                MyDateTime(2022, 4, 28, 6, 7, 8, 9).toPackedUInt(),
+                                MyDateTime(2022, 5, 2, 9, 8, 7, 6).toPackedUInt(),
+                                MyDateTime(2022, 5, 9, 9, 8, 7, 6).toPackedUInt(),
+                                MyDateTime(2022, 12, 31, 0, 0, 0, 0).toPackedUInt(),
+                                MyDateTime(2020, 12, 31, 0, 0, 0, 0).toPackedUInt(),
+                            })
+                            .column;
+    auto input_col = ColumnWithTypeAndName(data_col_ptr, nullable_datetime_type_ptr, "input");
+    auto output_col = createColumn<Nullable<UInt16>>({{}, {}, 52, 1, 1, 1, 2, 17, 18, 19, 52, 53});
+    ASSERT_COLUMN_EQ(output_col, executeFunction(func_name, input_col));
+
+    /// ColumnVector(non-null)
+    data_col_ptr = createColumn<DataTypeMyDateTime::FieldType>(
+                       {
+                           MyDateTime(0, 0, 0, 0, 0, 0, 0).toPackedUInt(),
+                           MyDateTime(0, 1, 1, 0, 0, 0, 0).toPackedUInt(),
+                           MyDateTime(0, 1, 2, 0, 0, 0, 0).toPackedUInt(),
+                           MyDateTime(0, 1, 3, 0, 0, 0, 0).toPackedUInt(),
+                           MyDateTime(1969, 1, 1, 1, 1, 1, 1).toPackedUInt(),
+                           MyDateTime(1969, 1, 6, 1, 1, 1, 1).toPackedUInt(),
+                           MyDateTime(2022, 4, 28, 6, 7, 8, 9).toPackedUInt(),
+                           MyDateTime(2022, 5, 2, 9, 8, 7, 6).toPackedUInt(),
+                           MyDateTime(2022, 5, 9, 9, 8, 7, 6).toPackedUInt(),
+                           MyDateTime(2022, 12, 31, 0, 0, 0, 0).toPackedUInt(),
+                           MyDateTime(2020, 12, 31, 0, 0, 0, 0).toPackedUInt(),
+                       })
+                       .column;
+    input_col = ColumnWithTypeAndName(data_col_ptr, datetime_type_ptr, "input");
+    output_col = createColumn<Nullable<UInt16>>({{}, 52, 1, 1, 1, 2, 17, 18, 19, 52, 53});
+    ASSERT_COLUMN_EQ(output_col, executeFunction(func_name, input_col));
+
+    /// ColumnConst(non-null)
+    input_col = ColumnWithTypeAndName(createConstColumn<DataTypeMyDateTime::FieldType>(1, MyDateTime(2022, 5, 9, 9, 8, 7, 6).toPackedUInt()).column, datetime_type_ptr, "input");
+    output_col = createConstColumn<Nullable<UInt16>>(1, {19});
+    ASSERT_COLUMN_EQ(output_col, executeFunction(func_name, input_col));
+
+    /// ColumnConst(nullable)
+    input_col = ColumnWithTypeAndName(createConstColumn<Nullable<DataTypeMyDateTime::FieldType>>(1, MyDateTime(2022, 5, 9, 9, 8, 7, 6).toPackedUInt()).column, nullable_datetime_type_ptr, "input");
+    output_col = createConstColumn<Nullable<UInt16>>(1, {19});
+    ASSERT_COLUMN_EQ(output_col, executeFunction(func_name, input_col));
+
+    /// ColumnConst(nullable(null))
+    input_col = ColumnWithTypeAndName(createConstColumn<Nullable<DataTypeMyDateTime::FieldType>>(1, {}).column, nullable_datetime_type_ptr, "input");
+    output_col = createConstColumn<Nullable<UInt16>>(1, {});
+    ASSERT_COLUMN_EQ(output_col, executeFunction(func_name, input_col));
+
+    /// MyDate ColumnVector(non-null)
+    data_col_ptr = createColumn<DataTypeMyDate::FieldType>(
+                       {
+                           MyDate(0, 0, 0).toPackedUInt(),
+                           MyDate(0, 1, 1).toPackedUInt(),
+                           MyDate(0, 1, 2).toPackedUInt(),
+                           MyDate(0, 1, 3).toPackedUInt(),
+                           MyDate(1969, 1, 1).toPackedUInt(),
+                           MyDate(1969, 1, 6).toPackedUInt(),
+                           MyDate(2022, 4, 28).toPackedUInt(),
+                           MyDate(2022, 5, 2).toPackedUInt(),
+                           MyDate(2022, 5, 9).toPackedUInt(),
+                           MyDate(2022, 12, 31).toPackedUInt(),
+                           MyDate(2020, 12, 31).toPackedUInt(),
+                       })
+                       .column;
+    input_col = ColumnWithTypeAndName(data_col_ptr, date_type_ptr, "input");
+    output_col = createColumn<Nullable<UInt16>>({{}, 52, 1, 1, 1, 2, 17, 18, 19, 52, 53});
+    ASSERT_COLUMN_EQ(output_col, executeFunction(func_name, input_col));
+    dag_context->setFlags(ori_flags);
+}
+CATCH
+
+} // namespace DB::tests

tests/fullstack-test/expr/week_of_year.test

@@ -0,0 +1,33 @@
+# Copyright 2022 PingCAP, Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+mysql> drop table if exists test.t;
+mysql> create table test.t(a char(20), b datetime, c date);
+mysql> insert into test.t values('', '1970-1-1 12:12:12', '1970-1-1');
+mysql> insert into test.t values('123', '1989-6-6 12:12:12', '1989-6-6');
+mysql> insert into test.t values('2022-3-10', '2000-3-4 12:12:12', '2000-3-4');
+mysql> alter table test.t set tiflash replica 1;
+
+func> wait_table test t
+
+mysql> set @@tidb_isolation_read_engines='tiflash'; set @@tidb_enforce_mpp = 1; select weekofyear(a), weekofyear(b), weekofyear(c) from test.t;
++---------------+---------------+---------------+
+| weekofyear(a) | weekofyear(b) | weekofyear(c) |
++---------------+---------------+---------------+
+|          NULL |             1 |             1 |
+|          NULL |            23 |            23 |
+|            10 |             9 |             9 |
++---------------+---------------+---------------+
+
+mysql> drop table if exists test.t;