pushdown get_format into TiFlash

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -513,7 +513,7 @@ const std::unordered_map<tipb::ScalarFuncSig, String> scalar_func_map({
     //{tipb::ScalarFuncSig::YearWeekWithMode, "cast"},
     //{tipb::ScalarFuncSig::YearWeekWithoutMode, "cast"},
 
-    //{tipb::ScalarFuncSig::GetFormat, "cast"},
+    {tipb::ScalarFuncSig::GetFormat, "getFormat"},
     {tipb::ScalarFuncSig::SysDateWithFsp, "sysDateWithFsp"},
     {tipb::ScalarFuncSig::SysDateWithoutFsp, "sysDateWithoutFsp"},
     //{tipb::ScalarFuncSig::CurrentDate, "cast"},

dbms/src/Functions/FunctionsConversion.cpp

@@ -240,6 +240,7 @@ void registerFunctionsConversion(FunctionFactory & factory)
 
     factory.registerFunction<FunctionFromUnixTime>();
     factory.registerFunction<FunctionDateFormat>();
+    factory.registerFunction<FunctionGetFormat>();
     factory.registerFunction<FunctionTiDBUnixTimeStamp<NameTiDBUnixTimeStampInt>>();
     factory.registerFunction<FunctionTiDBUnixTimeStamp<NameTiDBUnixTimeStampDec>>();
     factory.registerFunction<FunctionStrToDate<NameStrToDateDate>>();

dbms/src/Functions/FunctionsConversion.h

@@ -1751,6 +1751,114 @@ class FunctionDateFormat : public IFunction
     }
 };
 
+class FunctionGetFormat : public IFunction
+{
+private:
+    static String get_format(const String & time_type, const String & location)
+    {
+        if (time_type == "DATE")
+        {
+            if (location == "USA")
+                return "%m.%d.%Y";
+            else if (location == "JIS")
+                return "%Y-%m-%d";
+            else if (location == "ISO")
+                return "%Y-%m-%d";
+            else if (location == "EUR")
+                return "%d.%m.%Y";
+            else if (location == "INTERNAL")
+                return "%Y%m%d";
+        }
+        else if (time_type == "DATETIME" || time_type == "TIMESTAMP")
+        {
+            if (location == "USA")
+                return "%Y-%m-%d %H.%i.%s";
+            else if (location == "JIS")
+                return "%Y-%m-%d %H:%i:%s";
+            else if (location == "ISO")
+                return "%Y-%m-%d %H:%i:%s";
+            else if (location == "EUR")
+                return "%Y-%m-%d %H.%i.%s";
+            else if (location == "INTERNAL")
+                return "%Y%m%d%H%i%s";
+        }
+        else if (time_type == "TIME")
+        {
+            if (location == "USA")
+                return "%h:%i:%s %p";
+            else if (location == "JIS")
+                return "%H:%i:%s";
+            else if (location == "ISO")
+                return "%H:%i:%s";
+            else if (location == "EUR")
+                return "%H.%i.%s";
+            else if (location == "INTERNAL")
+                return "%H%i%s";
+        }
+        return "";
+    }
+
+public:
+    static constexpr auto name = "getFormat";
+    static FunctionPtr create(const Context &) { return std::make_shared<FunctionGetFormat>(); };
+
+    String getName() const override
+    {
+        return name;
+    }
+
+    size_t getNumberOfArguments() const override { return 2; }
+    bool isInjective(const Block &) const override { return false; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (!arguments[0].type->isString())
+            throw Exception("First argument for function " + getName() + " must be String constant", ErrorCodes::ILLEGAL_COLUMN);
+        if (!arguments[1].type->isString())
+            throw Exception("Second argument for function " + getName() + " must be String constant", ErrorCodes::ILLEGAL_COLUMN);
+
+        return std::make_shared<DataTypeString>();
+    }
+
+    bool useDefaultImplementationForConstants() const override { return true; }
+    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {0}; }
+
+    void executeImpl(Block & block, const ColumnNumbers & arguments, const size_t result) const override
+    {
+        const auto * location_col = checkAndGetColumn<ColumnString>(block.getByPosition(arguments[1]).column.get());
+        size_t size = location_col->size();
+        const auto & time_type_col = block.getByPosition(arguments[0]).column;
+        auto col_to = ColumnString::create();
+
+        if (time_type_col->isColumnConst())
+        {
+            const auto & time_type_col_const = checkAndGetColumnConst<ColumnString>(time_type_col.get());
+            const auto & time_type = time_type_col_const->getValue<String>();
+
+            ColumnString::Chars_t & data_to = col_to->getChars();
+            ColumnString::Offsets & offsets_to = col_to->getOffsets();
+            auto max_length = 18;
+            data_to.resize(size * max_length);
+            offsets_to.resize(size);
+            WriteBufferFromVector<ColumnString::Chars_t> write_buffer(data_to);
+            for (size_t i = 0; i < size; ++i)
+            {
+                const auto & location = location_col->getDataAt(i).toString();
+                const auto & result = get_format(time_type, location);
+                write_buffer.write(result.c_str(), result.size());
+                writeChar(0, write_buffer);
+                offsets_to[i] = write_buffer.count();
+            }
+            data_to.resize(write_buffer.count());
+            block.getByPosition(result).column = std::move(col_to);
+        }
+        else
+        {
+            throw Exception("Second argument for function " + getName() + " must be String constant", ErrorCodes::ILLEGAL_COLUMN);
+        }
+    }
+};
+
 struct NameStrToDateDate
 {
     static constexpr auto name = "strToDateDate";

tests/fullstack-test/expr/get_format.test

@@ -0,0 +1,64 @@
+# 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(location varchar(10));
+mysql> insert into test.t values(''), ('USA'), ('JIS'), ('ISO'), ('EUR'), ('INTERNAL');
+mysql> alter table test.t set tiflash replica 1;
+func> wait_table test t
+mysql> set @@tidb_enforce_mpp=1; set @@tidb_isolation_read_engines='tiflash'; select GET_FORMAT(DATE, location) from test.t;
++----------------------------+
+| GET_FORMAT(DATE, location) |
++----------------------------+
+|                            |
+| %m.%d.%Y                   |
+| %Y-%m-%d                   |
+| %Y-%m-%d                   |
+| %d.%m.%Y                   |
+| %Y%m%d                     |
++----------------------------+
+mysql> set @@tidb_enforce_mpp=1; set @@tidb_isolation_read_engines='tiflash'; select GET_FORMAT(DATETIME, location) from test.t;
++--------------------------------+
+| GET_FORMAT(DATETIME, location) |
++--------------------------------+
+|                                |
+| %Y-%m-%d %H.%i.%s              |
+| %Y-%m-%d %H:%i:%s              |
+| %Y-%m-%d %H:%i:%s              |
+| %Y-%m-%d %H.%i.%s              |
+| %Y%m%d%H%i%s                   |
++--------------------------------+
+mysql> set @@tidb_enforce_mpp=1; set @@tidb_isolation_read_engines='tiflash'; select GET_FORMAT(TIMESTAMP, location) from test.t;
++---------------------------------+
+| GET_FORMAT(TIMESTAMP, location) |
++---------------------------------+
+|                                 |
+| %Y-%m-%d %H.%i.%s               |
+| %Y-%m-%d %H:%i:%s               |
+| %Y-%m-%d %H:%i:%s               |
+| %Y-%m-%d %H.%i.%s               |
+| %Y%m%d%H%i%s                    |
++---------------------------------+
+mysql> set @@tidb_enforce_mpp=1; set @@tidb_isolation_read_engines='tiflash'; select GET_FORMAT(TIME, location) from test.t;
++----------------------------+
+| GET_FORMAT(TIME, location) |
++----------------------------+
+|                            |
+| %h:%i:%s %p                |
+| %H:%i:%s                   |
+| %H:%i:%s                   |
+| %H.%i.%s                   |
+| %H%i%s                     |
++----------------------------+
+mysql> drop table if exists test.t;