Merge pull request #68710 from ClickHouse/backport/24.8/68477

Backport #68477 to 24.8: Drop query cache by tag

docs/en/operations/query-cache.md

@@ -155,6 +155,8 @@ SELECT 1 SETTINGS use_query_cache = true, query_cache_tag = 'tag 1';
 SELECT 1 SETTINGS use_query_cache = true, query_cache_tag = 'tag 2';
 ```
 
+To remove only entries with tag `tag` from the query cache, you can use statement `SYSTEM DROP QUERY CACHE TAG 'tag'`.
+
 ClickHouse reads table data in blocks of [max_block_size](settings/settings.md#setting-max_block_size) rows. Due to filtering, aggregation,
 etc., result blocks are typically much smaller than 'max_block_size' but there are also cases where they are much bigger. Setting
 [query_cache_squash_partial_results](settings/settings.md#query-cache-squash-partial-results) (enabled by default) controls if result blocks

docs/en/sql-reference/statements/system.md

@@ -136,7 +136,13 @@ The compiled expression cache is enabled/disabled with the query/user/profile-le
 
 ## DROP QUERY CACHE
 
+```sql
+SYSTEM DROP QUERY CACHE;
+SYSTEM DROP QUERY CACHE TAG '<tag>'
+````
+
 Clears the [query cache](../../operations/query-cache.md).
+If a tag is specified, only query cache entries with the specified tag are deleted.
 
 ## DROP FORMAT SCHEMA CACHE {#system-drop-schema-format}
 

src/Common/CacheBase.h

@@ -197,6 +197,12 @@ public:
         cache_policy->remove(key);
     }
 
+    void remove(std::function<bool(const Key&, const MappedPtr &)> predicate)
+    {
+        std::lock_guard lock(mutex);
+        cache_policy->remove(predicate);
+    }
+
     size_t sizeInBytes() const
     {
         std::lock_guard lock(mutex);

src/Common/ICachePolicy.h

@@ -55,6 +55,7 @@ public:
     virtual void set(const Key & key, const MappedPtr & mapped) = 0;
 
     virtual void remove(const Key & key) = 0;
+    virtual void remove(std::function<bool(const Key & key, const MappedPtr & mapped)> predicate) = 0;
 
     virtual void clear() = 0;
     virtual std::vector<KeyMapped> dump() const = 0;

src/Common/LRUCachePolicy.h

@@ -79,6 +79,22 @@ public:
         cells.erase(it);
     }
 
+    void remove(std::function<bool(const Key &, const MappedPtr &)> predicate) override
+    {
+        for (auto it = cells.begin(); it != cells.end();)
+        {
+            if (predicate(it->first, it->second.value))
+            {
+                Cell & cell = it->second;
+                current_size_in_bytes -= cell.size;
+                queue.erase(cell.queue_iterator);
+                it = cells.erase(it);
+            }
+            else
+                ++it;
+        }
+    }
+
     MappedPtr get(const Key & key) override
     {
         auto it = cells.find(key);

src/Common/SLRUCachePolicy.h

@@ -95,6 +95,27 @@ public:
         cells.erase(it);
     }
 
+    void remove(std::function<bool(const Key &, const MappedPtr &)> predicate) override
+    {
+        for (auto it = cells.begin(); it != cells.end();)
+        {
+            if (predicate(it->first, it->second.value))
+            {
+                auto & cell = it->second;
+
+                current_size_in_bytes -= cell.size;
+                if (cell.is_protected)
+                    current_protected_size -= cell.size;
+
+                auto & queue = cell.is_protected ? protected_queue : probationary_queue;
+                queue.erase(cell.queue_iterator);
+                it = cells.erase(it);
+            }
+            else
+                ++it;
+        }
+    }
+
     MappedPtr get(const Key & key) override
     {
         auto it = cells.find(key);

src/Common/TTLCachePolicy.h

@@ -145,6 +145,23 @@ public:
         size_in_bytes -= sz;
     }
 
+    void remove(std::function<bool(const Key &, const MappedPtr &)> predicate) override
+    {
+        for (auto it = cache.begin(); it != cache.end();)
+        {
+            if (predicate(it->first, it->second))
+            {
+                size_t sz = weight_function(*it->second);
+                if (it->first.user_id.has_value())
+                    Base::user_quotas->decreaseActual(*it->first.user_id, sz);
+                it = cache.erase(it);
+                size_in_bytes -= sz;
+            }
+            else
+                ++it;
+        }
+    }
+
     MappedPtr get(const Key & key) override
     {
         auto it = cache.find(key);

src/Interpreters/Cache/QueryCache.cpp

@@ -619,9 +619,18 @@ QueryCache::Writer QueryCache::createWriter(const Key & key, std::chrono::millis
     return Writer(cache, key, max_entry_size_in_bytes, max_entry_size_in_rows, min_query_runtime, squash_partial_results, max_block_size);
 }
 
-void QueryCache::clear()
+void QueryCache::clear(const std::optional<String> & tag)
 {
-    cache.clear();
+    if (tag)
+    {
+        auto predicate = [tag](const Key & key, const Cache::MappedPtr &) { return key.tag == tag.value(); };
+        cache.remove(predicate);
+    }
+    else
+    {
+        cache.clear();
+    }
+
     std::lock_guard lock(mutex);
     times_executed.clear();
 }

src/Interpreters/Cache/QueryCache.h

@@ -211,7 +211,7 @@ public:
     Reader createReader(const Key & key);
     Writer createWriter(const Key & key, std::chrono::milliseconds min_query_runtime, bool squash_partial_results, size_t max_block_size, size_t max_query_cache_size_in_bytes_quota, size_t max_query_cache_entries_quota);
 
-    void clear();
+    void clear(const std::optional<String> & tag);
 
     size_t sizeInBytes() const;
     size_t count() const;

src/Interpreters/Context.cpp

@@ -3225,12 +3225,12 @@ QueryCachePtr Context::getQueryCache() const
     return shared->query_cache;
 }
 
-void Context::clearQueryCache() const
+void Context::clearQueryCache(const std::optional<String> & tag) const
 {
     std::lock_guard lock(shared->mutex);
 
     if (shared->query_cache)
-        shared->query_cache->clear();
+        shared->query_cache->clear(tag);
 }
 
 void Context::clearCaches() const

src/Interpreters/Context.h

@@ -1068,7 +1068,7 @@ public:
     void setQueryCache(size_t max_size_in_bytes, size_t max_entries, size_t max_entry_size_in_bytes, size_t max_entry_size_in_rows);
     void updateQueryCacheConfiguration(const Poco::Util::AbstractConfiguration & config);
     std::shared_ptr<QueryCache> getQueryCache() const;
-    void clearQueryCache() const;
+    void clearQueryCache(const std::optional<String> & tag) const;
 
     /** Clear the caches of the uncompressed blocks and marks.
       * This is usually done when renaming tables, changing the type of columns, deleting a table.

src/Interpreters/InterpreterSystemQuery.cpp

@@ -369,9 +369,12 @@ BlockIO InterpreterSystemQuery::execute()
             system_context->clearMMappedFileCache();
             break;
         case Type::DROP_QUERY_CACHE:
+        {
             getContext()->checkAccess(AccessType::SYSTEM_DROP_QUERY_CACHE);
-            getContext()->clearQueryCache();
+            getContext()->clearQueryCache(query.query_cache_tag);
             break;
+        }
+
         case Type::DROP_COMPILED_EXPRESSION_CACHE:
 #if USE_EMBEDDED_COMPILER
             getContext()->checkAccess(AccessType::SYSTEM_DROP_COMPILED_EXPRESSION_CACHE);

src/Parsers/ASTSystemQuery.h

@@ -130,6 +130,8 @@ public:
     String disk;
     UInt64 seconds{};
 
+    std::optional<String> query_cache_tag;
+
     String filesystem_cache_name;
     std::string key_to_drop;
     std::optional<size_t> offset_to_drop;

src/Parsers/CommonParsers.h

@@ -470,6 +470,7 @@ namespace DB
     MR_MACROS(TABLE_OVERRIDE, "TABLE OVERRIDE") \
     MR_MACROS(TABLE, "TABLE") \
     MR_MACROS(TABLES, "TABLES") \
+    MR_MACROS(TAG, "TAG") \
     MR_MACROS(TAGS, "TAGS") \
     MR_MACROS(TAGS_INNER_UUID, "TAGS INNER UUID") \
     MR_MACROS(TEMPORARY_TABLE, "TEMPORARY TABLE") \

src/Parsers/ParserSystemQuery.cpp

@@ -470,6 +470,16 @@ bool ParserSystemQuery::parseImpl(IParser::Pos & pos, ASTPtr & node, Expected &
             res->seconds = seconds->as<ASTLiteral>()->value.safeGet<UInt64>();
             break;
         }
+        case Type::DROP_QUERY_CACHE:
+        {
+            ParserLiteral tag_parser;
+            ASTPtr ast;
+            if (ParserKeyword{Keyword::TAG}.ignore(pos, expected) && tag_parser.parse(pos, ast, expected))
+                res->query_cache_tag = std::make_optional<String>(ast->as<ASTLiteral>()->value.safeGet<String>());
+            if (!parseQueryWithOnCluster(res, pos, expected))
+                return false;
+            break;
+        }
         case Type::DROP_FILESYSTEM_CACHE:
         {
             ParserLiteral path_parser;

tests/queries/0_stateless/02494_query_cache_drop_cache.reference

@@ -1,3 +1,17 @@
+Cache query result in query cache
 1
 1
+DROP entries with a certain tag, no entry will match
+1
+After a full DROP, the cache is empty now
+0
+Cache query result with different or no tag in query cache
+1
+1
+1
+2
+4
+DROP entries with certain tags
+2
+1
 0

tests/queries/0_stateless/02494_query_cache_drop_cache.sql

@@ -4,10 +4,31 @@
 -- (it's silly to use what will be tested below but we have to assume other tests cluttered the query cache)
 SYSTEM DROP QUERY CACHE;
 
--- Cache query result in query cache
+SELECT 'Cache query result in query cache';
 SELECT 1 SETTINGS use_query_cache = true;
 SELECT count(*) FROM system.query_cache;
 
--- No query results are cached after DROP
+SELECT 'DROP entries with a certain tag, no entry will match';
+SYSTEM DROP QUERY CACHE TAG 'tag';
+SELECT count(*) FROM system.query_cache;
+
+SELECT 'After a full DROP, the cache is empty now';
 SYSTEM DROP QUERY CACHE;
 SELECT count(*) FROM system.query_cache;
+
+-- More tests for DROP with tags:
+
+SELECT 'Cache query result with different or no tag in query cache';
+SELECT 1 SETTINGS use_query_cache = true;
+SELECT 1 SETTINGS use_query_cache = true, query_cache_tag = 'abc';
+SELECT 1 SETTINGS use_query_cache = true, query_cache_tag = 'def';
+SELECT 2 SETTINGS use_query_cache = true;
+SELECT count(*) FROM system.query_cache;
+
+SELECT 'DROP entries with certain tags';
+SYSTEM DROP QUERY CACHE TAG '';
+SELECT count(*) FROM system.query_cache;
+SYSTEM DROP QUERY CACHE TAG 'def';
+SELECT count(*) FROM system.query_cache;
+SYSTEM DROP QUERY CACHE TAG 'abc';
+SELECT count(*) FROM system.query_cache;