diff --git a/src/Storages/MergeTree/KeyCondition.cpp b/src/Storages/MergeTree/KeyCondition.cpp index 9e19ffe1b04..b8c03ba2cd6 100644 --- a/src/Storages/MergeTree/KeyCondition.cpp +++ b/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(&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, diff --git a/tests/queries/0_stateless/03272_partition_pruning_monotonic_func_bug.reference b/tests/queries/0_stateless/03272_partition_pruning_monotonic_func_bug.reference new file mode 100644 index 00000000000..cbb68477a29 --- /dev/null +++ b/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 diff --git a/tests/queries/0_stateless/03272_partition_pruning_monotonic_func_bug.sql b/tests/queries/0_stateless/03272_partition_pruning_monotonic_func_bug.sql new file mode 100644 index 00000000000..b69b7b8a754 --- /dev/null +++ b/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;