Merge pull request #71966 from ClickHouse/fix-pk-const-monotonic-transform-for-2-args

Fix partition pruning with binary monotonic function by second arg.

src/Storages/MergeTree/KeyCondition.cpp

@@ -945,6 +945,8 @@ static FieldRef applyFunction(const FunctionBasePtr & func, const DataTypePtr &
     return {field.columns, field.row_idx, result_idx};
 }
 
+DataTypePtr getArgumentTypeOfMonotonicFunction(const IFunctionBase & func);
+
 /// Sequentially applies functions to the column, returns `true`
 /// if all function arguments are compatible with functions
 /// signatures, and none of the functions produce `NULL` output.
@@ -976,7 +978,7 @@ bool applyFunctionChainToColumn(
     }
 
     // And cast it to the argument type of the first function in the chain
-    auto in_argument_type = functions[0]->getArgumentTypes()[0];
+    auto in_argument_type = getArgumentTypeOfMonotonicFunction(*functions[0]);
     if (canBeSafelyCasted(result_type, in_argument_type))
     {
         result_column = castColumnAccurate({result_column, result_type, ""}, in_argument_type);
@@ -1005,7 +1007,7 @@ bool applyFunctionChainToColumn(
         if (func->getArgumentTypes().empty())
             return false;
 
-        auto argument_type = func->getArgumentTypes()[0];
+        auto argument_type = getArgumentTypeOfMonotonicFunction(*func);
         if (!canBeSafelyCasted(result_type, argument_type))
             return false;
 
@@ -1489,6 +1491,18 @@ private:
     Kind kind = Kind::NO_CONST;
 };
 
+DataTypePtr getArgumentTypeOfMonotonicFunction(const IFunctionBase & func)
+{
+    const auto & arg_types = func.getArgumentTypes();
+    if (const auto * func_ptr = typeid_cast<const FunctionWithOptionalConstArg *>(&func))
+    {
+        if (func_ptr->getKind() == FunctionWithOptionalConstArg::Kind::LEFT_CONST)
+            return arg_types.at(1);
+    }
+
+    return arg_types.at(0);
+}
+
 
 bool KeyCondition::isKeyPossiblyWrappedByMonotonicFunctions(
     const RPNBuilderTreeNode & node,

tests/queries/0_stateless/03272_partition_pruning_monotonic_func_bug.reference

@@ -0,0 +1,11 @@
+1
+2
+Expression
+  ReadFromMergeTree
+  Indexes:
+    PrimaryKey
+      Keys: 
+        dateTrunc(\'hour\', ts)
+      Condition: and((dateTrunc(\'hour\', ts) in (-Inf, 1731592800]), (dateTrunc(\'hour\', ts) in [1731506400, +Inf)))
+      Parts: 1/1
+      Granules: 1/1

tests/queries/0_stateless/03272_partition_pruning_monotonic_func_bug.sql

@@ -0,0 +1,19 @@
+SET session_timezone = 'Etc/UTC';
+
+DROP TABLE IF EXISTS tt;
+CREATE TABLE tt
+(
+    `id` Int64,
+    `ts` DateTime
+)
+ENGINE = MergeTree()
+ORDER BY dateTrunc('hour', ts)
+SETTINGS index_granularity = 8192;
+
+INSERT INTO tt VALUES (1, '2024-11-14 00:00:00'), (2, '2024-11-14 00:00:00');
+
+SELECT id FROM tt PREWHERE ts BETWEEN toDateTime(1731506400) AND toDateTime(1731594420);
+
+explain indexes=1, description=0 SELECT id FROM tt PREWHERE ts BETWEEN toDateTime(1731506400) AND toDateTime(1731594420);
+
+DROP TABLE tt;