From c4fcdb8b91b9d01a4e2d8fa3b90ac938927e01fa Mon Sep 17 00:00:00 2001
From: huangzhaowei <huangzhaowei@didiglobal.com>
Date: Tue, 13 Jul 2021 13:54:44 +0800
Subject: [PATCH] Make toTimeZone monotonicity when timeZone is a constant
 value

---
 src/Functions/toTimezone.cpp | 114 ++++++++++++++++++++++++++++-------
 1 file changed, 93 insertions(+), 21 deletions(-)

diff --git a/src/Functions/toTimezone.cpp b/src/Functions/toTimezone.cpp
index 551e07a8354..a47fde7e429 100644
--- a/src/Functions/toTimezone.cpp
+++ b/src/Functions/toTimezone.cpp
@@ -19,44 +19,116 @@ namespace ErrorCodes
 
 namespace
 {
-
-/// Just changes time zone information for data type. The calculation is free.
-class FunctionToTimezone : public IFunction
+class ExecutableFunctionToTimeZone : public IExecutableFunction
 {
 public:
-    static constexpr auto name = "toTimezone";
-    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionToTimezone>(); }
+    explicit ExecutableFunctionToTimeZone() = default;
 
-    String getName() const override
+    String getName() const override { return "toTimeZone"; }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & /*result_type*/, size_t /*input_rows_count*/) const override
     {
-        return name;
+        return arguments[0].column;
+    }
+};
+
+class FunctionBaseToTimeZone : public IFunctionBase
+{
+public:
+    FunctionBaseToTimeZone(bool is_constant_timezone_,
+                           DataTypes argument_types_,
+                           DataTypePtr return_type_):
+        is_constant_timezone(is_constant_timezone_),
+        argument_types(std::move(argument_types_)),
+        return_type(std::move(return_type_)) {}
+
+    String getName() const override { return "toTimeZone"; }
+
+    const DataTypes & getArgumentTypes() const override
+    {
+        return argument_types;
     }
 
-    size_t getNumberOfArguments() const override { return 2; }
+    const DataTypePtr & getResultType() const override
+    {
+        return return_type;
+    }
 
-    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    ExecutableFunctionPtr prepare(const ColumnsWithTypeAndName & /*arguments*/) const override
+    {
+        return std::make_unique<ExecutableFunctionToTimeZone>();
+    }
+
+    bool hasInformationAboutMonotonicity() const override { return is_constant_timezone; }
+
+    Monotonicity getMonotonicityForRange(const IDataType & /*type*/, const Field & /*left*/, const Field & /*right*/) const override
+    {
+        return {is_constant_timezone, is_constant_timezone, is_constant_timezone};
+    }
+
+private:
+    bool is_constant_timezone;
+    DataTypes argument_types;
+    DataTypePtr return_type;
+};
+
+/// Just changes time zone information for data type. The calculation is free.
+class ToTimeZoneOverloadResolver : public IFunctionOverloadResolver
+{
+public:
+    static constexpr auto name = "toTimeZone";
+
+    String getName() const override { return name; }
+
+    size_t getNumberOfArguments() const override { return 2; }
+    static FunctionOverloadResolverPtr create(ContextPtr) { return std::make_unique<ToTimeZoneOverloadResolver>(); }
+
+    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
+    {
+        const auto which_type = WhichDataType(arguments[0]);
+        if (!which_type.isDateTime() && !which_type.isDateTime64())
+            throw Exception{"Illegal type " + arguments[0]->getName() + " of argument of function " + getName() +
+                            ". Should be DateTime or DateTime64", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+
+        if (which_type.isDateTime())
+            return std::make_shared<DataTypeDateTime>();
+
+        const auto * date_time64 = assert_cast<const DataTypeDateTime64 *>(arguments[0].get());
+        return std::make_shared<DataTypeDateTime64>(date_time64->getScale());
+    }
+
+    FunctionBasePtr buildImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &) const override
     {
         if (arguments.size() != 2)
+        {
             throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
-                + toString(arguments.size()) + ", should be 2",
-                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+                            + toString(arguments.size()) + ", should be 2",
+                            ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+        }
 
         const auto which_type = WhichDataType(arguments[0].type);
         if (!which_type.isDateTime() && !which_type.isDateTime64())
             throw Exception{"Illegal type " + arguments[0].type->getName() + " of argument of function " + getName() +
-                ". Should be DateTime or DateTime64", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+                            ". Should be DateTime or DateTime64", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT};
+
+        bool is_constant_timezone = false;
+        if (arguments[1].column)
+        {
+            is_constant_timezone = isColumnConst(*arguments[1].column);
+        }
 
         String time_zone_name = extractTimeZoneNameFromFunctionArguments(arguments, 1, 0);
+        DataTypes data_types(arguments.size());
+        for (size_t i = 0; i < arguments.size(); ++i)
+            data_types[i] = arguments[i].type;
+
         if (which_type.isDateTime())
-            return std::make_shared<DataTypeDateTime>(time_zone_name);
+        {
+            return std::make_unique<FunctionBaseToTimeZone>(is_constant_timezone, data_types, std::make_shared<DataTypeDateTime>(time_zone_name));
+        }
 
         const auto * date_time64 = assert_cast<const DataTypeDateTime64 *>(arguments[0].type.get());
-        return std::make_shared<DataTypeDateTime64>(date_time64->getScale(), time_zone_name);
-    }
-
-    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t /*input_rows_count*/) const override
-    {
-        return arguments[0].column;
+        return std::make_unique<FunctionBaseToTimeZone>(is_constant_timezone, data_types, std::make_shared<DataTypeDateTime64>(date_time64->getScale(), time_zone_name));
     }
 };
 
@@ -64,8 +136,8 @@ public:
 
 void registerFunctionToTimeZone(FunctionFactory & factory)
 {
-    factory.registerFunction<FunctionToTimezone>();
-    factory.registerAlias("toTimeZone", "toTimezone");
+    factory.registerFunction<ToTimeZoneOverloadResolver>();
+    factory.registerAlias("toTimezone", "toTimeZone");
 }
 
 }