diff --git a/src/Storages/MergeTree/KeyCondition.cpp b/src/Storages/MergeTree/KeyCondition.cpp
index 1ce64c52f6c..06fa8e83505 100644
--- a/src/Storages/MergeTree/KeyCondition.cpp
+++ b/src/Storages/MergeTree/KeyCondition.cpp
@@ -599,21 +599,29 @@ bool KeyCondition::canConstantBeWrappedByMonotonicFunctions(
     Field & out_value,
     DataTypePtr & out_type)
 {
-    auto adjusted_node = node->clone();
-    KeyDescription::moduloToModuloLegacyRecursive(adjusted_node);
+    const auto & sample_block = key_expr->getSampleBlock();
 
-    String adjusted_expr_name = adjusted_node->getColumnName();
     // Constant expr should use alias names if any
     String passed_expr_name = node->getColumnName();
     String expr_name;
 
-    const auto & sample_block = key_expr->getSampleBlock();
+    /// sample_block from key_expr cannot contain modulo and moduloLegacy at the same time.
+    /// For partition key it is always moduloLegacy.
     if (sample_block.has(passed_expr_name))
+    {
         expr_name = passed_expr_name;
-    else if (sample_block.has(adjusted_expr_name))
-        expr_name = adjusted_expr_name;
+    }
     else
-        return false;
+    {
+        auto adjusted_ast = node->clone();
+        KeyDescription::moduloToModuloLegacyRecursive(adjusted_ast);
+        String adjusted_expr_name = adjusted_ast->getColumnName();
+
+        if (!sample_block.has(adjusted_expr_name))
+            return false;
+
+        expr_name = adjusted_expr_name;
+    }
 
     /// TODO Nullable index is not yet landed.
     if (out_value.isNull())
@@ -677,21 +685,29 @@ bool KeyCondition::canConstantBeWrappedByMonotonicFunctions(
 bool KeyCondition::canConstantBeWrappedByFunctions(
     const ASTPtr & ast, size_t & out_key_column_num, DataTypePtr & out_key_column_type, Field & out_value, DataTypePtr & out_type)
 {
-    auto adjusted_ast = ast->clone();
-    KeyDescription::moduloToModuloLegacyRecursive(adjusted_ast);
+    const auto & sample_block = key_expr->getSampleBlock();
 
-    String adjusted_expr_name = adjusted_ast->getColumnName();
     // Constant expr should use alias names if any
     String passed_expr_name = ast->getColumnName();
     String expr_name;
 
-    const auto & sample_block = key_expr->getSampleBlock();
+    /// sample_block from key_expr cannot contain modulo and moduloLegacy at the same time.
+    /// For partition key it is always moduloLegacy.
     if (sample_block.has(passed_expr_name))
+    {
         expr_name = passed_expr_name;
-    else if (sample_block.has(adjusted_expr_name))
-        expr_name = adjusted_expr_name;
+    }
     else
-        return false;
+    {
+        auto adjusted_ast = ast->clone();
+        KeyDescription::moduloToModuloLegacyRecursive(adjusted_ast);
+        String adjusted_expr_name = adjusted_ast->getColumnName();
+
+        if (!sample_block.has(adjusted_expr_name))
+            return false;
+
+        expr_name = adjusted_expr_name;
+    }
 
     /// TODO Nullable index is not yet landed.
     if (out_value.isNull())
diff --git a/src/Storages/MergeTree/MergeTreePartition.h b/src/Storages/MergeTree/MergeTreePartition.h
index 3f0ee2864a0..541881633a0 100644
--- a/src/Storages/MergeTree/MergeTreePartition.h
+++ b/src/Storages/MergeTree/MergeTreePartition.h
@@ -47,7 +47,7 @@ public:
     static Block executePartitionByExpression(const StorageMetadataPtr & metadata_snapshot, Block & block, ContextPtr context);
 
     /// Make a modified partition key with substitution from modulo to moduloLegacy. Used in paritionPruner.
-    static KeyDescription adjustPartitionKey(const StorageMetadataPtr & partition_key, ContextPtr context);
+    static KeyDescription adjustPartitionKey(const StorageMetadataPtr & metadata_snapshot, ContextPtr context);
 };
 
 }