From 7a34fbc5b2f322341ae6c378920670f6a2258698 Mon Sep 17 00:00:00 2001
From: Anton Popov <anton@clickhouse.com>
Date: Fri, 1 Nov 2024 14:20:37 +0000
Subject: [PATCH] allow to prewarm mark cache without enabled setting

---
 src/Interpreters/InterpreterSystemQuery.cpp   |  2 +-
 src/Storages/MergeTree/MergeTreeData.cpp      |  7 ++++-
 src/Storages/MergeTree/MergeTreeData.h        |  1 +
 src/Storages/StorageMergeTree.cpp             |  2 +-
 src/Storages/StorageReplicatedMergeTree.cpp   |  2 +-
 .../03254_system_prewarm_mark_cache.reference |  4 +++
 .../03254_system_prewarm_mark_cache.sql       | 27 +++++++++++++++++++
 7 files changed, 41 insertions(+), 4 deletions(-)
 create mode 100644 tests/queries/0_stateless/03254_system_prewarm_mark_cache.reference
 create mode 100644 tests/queries/0_stateless/03254_system_prewarm_mark_cache.sql

diff --git a/src/Interpreters/InterpreterSystemQuery.cpp b/src/Interpreters/InterpreterSystemQuery.cpp
index 45636ab40b9..4c875026ace 100644
--- a/src/Interpreters/InterpreterSystemQuery.cpp
+++ b/src/Interpreters/InterpreterSystemQuery.cpp
@@ -1310,7 +1310,7 @@ RefreshTaskList InterpreterSystemQuery::getRefreshTasks()
 void InterpreterSystemQuery::prewarmMarkCache()
 {
     if (table_id.empty())
-        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Table is not specified for prewarming marks cache");
+        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Table is not specified for PREWARM MARK CACHE command");
 
     getContext()->checkAccess(AccessType::SYSTEM_PREWARM_MARK_CACHE, table_id);
 
diff --git a/src/Storages/MergeTree/MergeTreeData.cpp b/src/Storages/MergeTree/MergeTreeData.cpp
index 4ed8c67469d..69979809c31 100644
--- a/src/Storages/MergeTree/MergeTreeData.cpp
+++ b/src/Storages/MergeTree/MergeTreeData.cpp
@@ -2343,11 +2343,16 @@ void MergeTreeData::stopOutdatedAndUnexpectedDataPartsLoadingTask()
     }
 }
 
-void MergeTreeData::prewarmMarkCache(ThreadPool & pool)
+void MergeTreeData::prewarmMarkCacheIfNeeded(ThreadPool & pool)
 {
     if (!(*getSettings())[MergeTreeSetting::prewarm_mark_cache])
         return;
 
+    prewarmMarkCache(pool);
+}
+
+void MergeTreeData::prewarmMarkCache(ThreadPool & pool)
+{
     auto * mark_cache = getContext()->getMarkCache().get();
     if (!mark_cache)
         return;
diff --git a/src/Storages/MergeTree/MergeTreeData.h b/src/Storages/MergeTree/MergeTreeData.h
index a32106f76bb..8da4329a93b 100644
--- a/src/Storages/MergeTree/MergeTreeData.h
+++ b/src/Storages/MergeTree/MergeTreeData.h
@@ -508,6 +508,7 @@ public:
 
     /// Prewarm mark cache for the most recent data parts.
     void prewarmMarkCache(ThreadPool & pool);
+    void prewarmMarkCacheIfNeeded(ThreadPool & pool);
 
     String getLogName() const { return log.loadName(); }
 
diff --git a/src/Storages/StorageMergeTree.cpp b/src/Storages/StorageMergeTree.cpp
index 40cd6e01dba..1ba0617d8ae 100644
--- a/src/Storages/StorageMergeTree.cpp
+++ b/src/Storages/StorageMergeTree.cpp
@@ -155,7 +155,7 @@ StorageMergeTree::StorageMergeTree(
 
     loadMutations();
     loadDeduplicationLog();
-    prewarmMarkCache(getActivePartsLoadingThreadPool().get());
+    prewarmMarkCacheIfNeeded(getActivePartsLoadingThreadPool().get());
 }
 
 
diff --git a/src/Storages/StorageReplicatedMergeTree.cpp b/src/Storages/StorageReplicatedMergeTree.cpp
index bbfedb2f355..15341cca976 100644
--- a/src/Storages/StorageReplicatedMergeTree.cpp
+++ b/src/Storages/StorageReplicatedMergeTree.cpp
@@ -509,7 +509,7 @@ StorageReplicatedMergeTree::StorageReplicatedMergeTree(
     }
 
     loadDataParts(skip_sanity_checks, expected_parts_on_this_replica);
-    prewarmMarkCache(getActivePartsLoadingThreadPool().get());
+    prewarmMarkCacheIfNeeded(getActivePartsLoadingThreadPool().get());
 
     if (LoadingStrictnessLevel::ATTACH <= mode)
     {
diff --git a/tests/queries/0_stateless/03254_system_prewarm_mark_cache.reference b/tests/queries/0_stateless/03254_system_prewarm_mark_cache.reference
new file mode 100644
index 00000000000..86674e7765a
--- /dev/null
+++ b/tests/queries/0_stateless/03254_system_prewarm_mark_cache.reference
@@ -0,0 +1,4 @@
+20000
+20000
+1
+0
diff --git a/tests/queries/0_stateless/03254_system_prewarm_mark_cache.sql b/tests/queries/0_stateless/03254_system_prewarm_mark_cache.sql
new file mode 100644
index 00000000000..f9e77365836
--- /dev/null
+++ b/tests/queries/0_stateless/03254_system_prewarm_mark_cache.sql
@@ -0,0 +1,27 @@
+-- Tags: no-parallel, no-shared-merge-tree
+
+DROP TABLE IF EXISTS t_prewarm_cache;
+
+CREATE TABLE t_prewarm_cache (a UInt64, b UInt64, c UInt64)
+ENGINE = ReplicatedMergeTree('/clickhouse/tables/{database}/03254_prewarm_mark_cache_smt/t_prewarm_cache', '1')
+ORDER BY a SETTINGS prewarm_mark_cache = 0;
+
+SYSTEM DROP MARK CACHE;
+
+INSERT INTO t_prewarm_cache SELECT number, rand(), rand() FROM numbers(20000);
+
+SELECT count() FROM t_prewarm_cache WHERE NOT ignore(*);
+
+SYSTEM DROP MARK CACHE;
+
+SYSTEM PREWARM MARK CACHE t_prewarm_cache;
+
+SELECT count() FROM t_prewarm_cache WHERE NOT ignore(*);
+
+SYSTEM FLUSH LOGS;
+
+SELECT ProfileEvents['LoadedMarksCount'] > 0 FROM system.query_log
+WHERE current_database = currentDatabase() AND type = 'QueryFinish' AND query LIKE 'SELECT count() FROM t_prewarm_cache%'
+ORDER BY event_time_microseconds;
+
+DROP TABLE IF EXISTS t_prewarm_cache;