Flexible drop cache

src/Disks/IO/CachedOnDiskReadBufferFromFile.cpp

@@ -74,19 +74,22 @@ CachedOnDiskReadBufferFromFile::CachedOnDiskReadBufferFromFile(
 }
 
 void CachedOnDiskReadBufferFromFile::appendFilesystemCacheLog(
-    const FileSegment::Range & file_segment_range, CachedOnDiskReadBufferFromFile::ReadType type)
+    const FileSegment & file_segment, CachedOnDiskReadBufferFromFile::ReadType type)
 {
     if (!cache_log)
         return;
 
+    const auto range = file_segment.range();
     FilesystemCacheLogElement elem
     {
         .event_time = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now()),
         .query_id = query_id,
         .source_file_path = source_file_path,
-        .file_segment_range = { file_segment_range.left, file_segment_range.right },
+        .file_segment_range = { range.left, range.right },
         .requested_range = { first_offset, read_until_position },
-        .file_segment_size = file_segment_range.size(),
+        .file_segment_key = file_segment.key().toString(),
+        .file_segment_offset = file_segment.offset(),
+        .file_segment_size = range.size(),
         .read_from_cache_attempted = true,
         .read_buffer_id = current_buffer_id,
         .profile_counters = std::make_shared<ProfileEvents::Counters::Snapshot>(
@@ -495,7 +498,7 @@ bool CachedOnDiskReadBufferFromFile::completeFileSegmentAndGetNext()
     auto completed_range = current_file_segment->range();
 
     if (cache_log)
-        appendFilesystemCacheLog(completed_range, read_type);
+        appendFilesystemCacheLog(*current_file_segment, read_type);
 
     chassert(file_offset_of_buffer_end > completed_range.right);
 
@@ -521,7 +524,7 @@ CachedOnDiskReadBufferFromFile::~CachedOnDiskReadBufferFromFile()
 {
     if (cache_log && file_segments && !file_segments->empty())
     {
-        appendFilesystemCacheLog(file_segments->front().range(), read_type);
+        appendFilesystemCacheLog(file_segments->front(), read_type);
     }
 }
 

src/Disks/IO/CachedOnDiskReadBufferFromFile.h

@@ -90,7 +90,7 @@ private:
 
     bool completeFileSegmentAndGetNext();
 
-    void appendFilesystemCacheLog(const FileSegment::Range & file_segment_range, ReadType read_type);
+    void appendFilesystemCacheLog(const FileSegment & file_segment, ReadType read_type);
 
     bool writeCache(char * data, size_t size, size_t offset, FileSegment & file_segment);
 

src/Disks/IO/ReadBufferFromRemoteFSGather.cpp

@@ -88,6 +88,8 @@ void ReadBufferFromRemoteFSGather::appendUncachedReadInfo()
         .source_file_path = current_object.remote_path,
         .file_segment_range = { 0, current_object.bytes_size },
         .cache_type = FilesystemCacheLogElement::CacheType::READ_FROM_FS_BYPASSING_CACHE,
+        .file_segment_key = {},
+        .file_segment_offset = {},
         .file_segment_size = current_object.bytes_size,
         .read_from_cache_attempted = false,
     };

src/Disks/getDiskConfigurationFromAST.cpp

@@ -8,6 +8,7 @@
 #include <Parsers/ASTFunction.h>
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTLiteral.h>
+#include <Parsers/formatAST.h>
 #include <Parsers/ASTFunctionWithKeyValueArguments.h>
 #include <Storages/checkAndGetLiteralArgument.h>
 #include <Poco/DOM/Document.h>

src/Disks/getOrCreateDiskFromAST.cpp

@@ -2,6 +2,7 @@
 #include <Common/logger_useful.h>
 #include <Common/assert_cast.h>
 #include <Common/filesystemHelpers.h>
+#include <Parsers/isDiskFunction.h>
 #include <Disks/getDiskConfigurationFromAST.h>
 #include <Disks/DiskSelector.h>
 #include <Parsers/formatAST.h>
@@ -9,7 +10,6 @@
 #include <Parsers/ASTLiteral.h>
 #include <Parsers/ASTIdentifier.h>
 #include <Parsers/ASTFunction.h>
-#include <Parsers/isDiskFunction.h>
 #include <Interpreters/Context.h>
 #include <Parsers/IAST.h>
 #include <Interpreters/InDepthNodeVisitor.h>
@@ -26,12 +26,20 @@ namespace
 {
     std::string getOrCreateDiskFromDiskAST(const ASTFunction & function, ContextPtr context)
     {
-        /// We need a unique name for a created custom disk, but it needs to be the same
-        /// after table is reattached or server is restarted, so take a hash of the disk
-        /// configuration serialized ast as a disk name suffix.
-        auto disk_setting_string = serializeAST(function, true);
-        auto disk_name = DiskSelector::TMP_INTERNAL_DISK_PREFIX
-            + toString(sipHash128(disk_setting_string.data(), disk_setting_string.size()));
+        std::string disk_name;
+        if (function.name == "disk")
+        {
+            /// We need a unique name for a created custom disk, but it needs to be the same
+            /// after table is reattached or server is restarted, so take a hash of the disk
+            /// configuration serialized ast as a disk name suffix.
+            auto disk_setting_string = serializeAST(function, true);
+            disk_name = DiskSelector::TMP_INTERNAL_DISK_PREFIX
+                + toString(sipHash128(disk_setting_string.data(), disk_setting_string.size()));
+        }
+        else
+        {
+            disk_name = function.name.substr(std::strlen("disk_"));
+        }
 
         auto result_disk = context->getOrCreateDisk(disk_name, [&](const DisksMap & disks_map) -> DiskPtr {
             const auto * function_args_expr = assert_cast<const ASTExpressionList *>(function.arguments.get());
@@ -43,6 +51,9 @@ namespace
             return disk;
         });
 
+        if (!result_disk->isCustomDisk())
+            throw Exception(ErrorCodes::BAD_ARGUMENTS, "Disk with name `{}` already exist", disk_name);
+
         if (!result_disk->isRemote())
         {
             static constexpr auto custom_disks_base_dir_in_config = "custom_local_disks_base_directory";
@@ -91,8 +102,8 @@ namespace
 
 std::string getOrCreateDiskFromDiskAST(const ASTPtr & disk_function, ContextPtr context)
 {
-    if (!isDiskFunction(disk_function))
-        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected a disk function");
+    if (!disk_function->as<ASTFunction>())
+        throw Exception(ErrorCodes::BAD_ARGUMENTS, "Expected a function");
 
     auto ast = disk_function->clone();
 

src/Interpreters/Cache/FileCache.cpp

@@ -807,6 +807,17 @@ bool FileCache::tryReserve(FileSegment & file_segment, const size_t size)
     return true;
 }
 
+void FileCache::removeKey(const Key & key)
+{
+    assertInitialized();
+
+    auto locked_key = metadata.lockKeyMetadata(key, CacheMetadata::KeyNotFoundPolicy::THROW);
+    if (!locked_key)
+        return;
+
+    locked_key->removeAllReleasable();
+}
+
 void FileCache::removeKeyIfExists(const Key & key)
 {
     assertInitialized();
@@ -822,6 +833,17 @@ void FileCache::removeKeyIfExists(const Key & key)
     locked_key->removeAllReleasable();
 }
 
+void FileCache::removeFileSegment(const Key & key, size_t offset)
+{
+    assertInitialized();
+
+    auto locked_key = metadata.lockKeyMetadata(key, CacheMetadata::KeyNotFoundPolicy::RETURN_NULL);
+    if (!locked_key)
+        return;
+
+    locked_key->removeFileSegment(offset);
+}
+
 void FileCache::removePathIfExists(const String & path)
 {
     removeKeyIfExists(createKeyForPath(path));
@@ -916,7 +938,7 @@ void FileCache::loadMetadata()
                 continue;
             }
 
-            const auto key = Key(unhexUInt<UInt128>(key_directory.filename().string().data()));
+            const auto key = Key::fromKeyString(key_directory.filename().string());
             auto locked_key = metadata.lockKeyMetadata(key, CacheMetadata::KeyNotFoundPolicy::CREATE_EMPTY, /* is_initial_load */true);
 
             for (fs::directory_iterator offset_it{key_directory}; offset_it != fs::directory_iterator(); ++offset_it)

src/Interpreters/Cache/FileCache.h

@@ -83,13 +83,19 @@ public:
 
     FileSegmentsHolderPtr set(const Key & key, size_t offset, size_t size, const CreateFileSegmentSettings & settings);
 
-    /// Remove files by `key`. Removes files which might be used at the moment.
+    /// Remove file segment by `key` and `offset`. Throws if file segment does not exist.
+    void removeFileSegment(const Key & key, size_t offset);
+
+    /// Remove files by `key`. Throws if key does not exist.
+    void removeKey(const Key & key);
+
+    /// Remove files by `key`.
     void removeKeyIfExists(const Key & key);
 
-    /// Removes files by `path`. Removes files which might be used at the moment.
+    /// Removes files by `path`.
     void removePathIfExists(const String & path);
 
-    /// Remove files by `key`. Will not remove files which are used at the moment.
+    /// Remove files by `key`.
     void removeAllReleasable();
 
     std::vector<String> tryGetCachePaths(const Key & key);

src/Interpreters/Cache/FileCacheKey.cpp

@@ -28,4 +28,9 @@ FileCacheKey FileCacheKey::random()
     return FileCacheKey(UUIDHelpers::generateV4().toUnderType());
 }
 
+FileCacheKey FileCacheKey::fromKeyString(const std::string & key_str)
+{
+    return FileCacheKey(unhexUInt<UInt128>(key_str.data()));
+}
+
 }

src/Interpreters/Cache/FileCacheKey.h

@@ -21,6 +21,8 @@ struct FileCacheKey
     static FileCacheKey random();
 
     bool operator==(const FileCacheKey & other) const { return key == other.key; }
+
+    static FileCacheKey fromKeyString(const std::string & key_str);
 };
 
 using FileCacheKeyAndOffset = std::pair<FileCacheKey, size_t>;

src/Interpreters/Cache/Metadata.cpp

@@ -381,17 +381,32 @@ void LockedKey::removeAllReleasable()
     }
 }
 
-KeyMetadata::iterator LockedKey::removeFileSegment(size_t offset, const FileSegmentGuard::Lock & segment_lock)
+KeyMetadata::iterator LockedKey::removeFileSegment(size_t offset)
 {
     auto it = key_metadata->find(offset);
     if (it == key_metadata->end())
         throw Exception(ErrorCodes::LOGICAL_ERROR, "There is no offset {}", offset);
 
     auto file_segment = it->second->file_segment;
+    return removeFileSegmentImpl(it, file_segment->lock());
+}
+
+KeyMetadata::iterator LockedKey::removeFileSegment(size_t offset, const FileSegmentGuard::Lock & segment_lock)
+{
+    auto it = key_metadata->find(offset);
+    if (it == key_metadata->end())
+        throw Exception(ErrorCodes::LOGICAL_ERROR, "There is no offset {}", offset);
+
+    return removeFileSegmentImpl(it, segment_lock);
+}
+
+KeyMetadata::iterator LockedKey::removeFileSegmentImpl(KeyMetadata::iterator it, const FileSegmentGuard::Lock & segment_lock)
+{
+    auto file_segment = it->second->file_segment;
 
     LOG_DEBUG(
         key_metadata->log, "Remove from cache. Key: {}, offset: {}, size: {}",
-        getKey(), offset, file_segment->reserved_size);
+        getKey(), file_segment->offset(), file_segment->reserved_size);
 
     chassert(file_segment->assertCorrectnessUnlocked(segment_lock));
 

src/Interpreters/Cache/Metadata.h

@@ -159,6 +159,7 @@ struct LockedKey : private boost::noncopyable
     void removeAllReleasable();
 
     KeyMetadata::iterator removeFileSegment(size_t offset, const FileSegmentGuard::Lock &);
+    KeyMetadata::iterator removeFileSegment(size_t offset);
 
     void shrinkFileSegmentToDownloadedSize(size_t offset, const FileSegmentGuard::Lock &);
 
@@ -173,6 +174,8 @@ struct LockedKey : private boost::noncopyable
     std::string toString() const;
 
 private:
+    KeyMetadata::iterator removeFileSegmentImpl(KeyMetadata::iterator it, const FileSegmentGuard::Lock &);
+
     const std::shared_ptr<KeyMetadata> key_metadata;
     KeyGuard::Lock lock; /// `lock` must be destructed before `key_metadata`.
 };

src/Interpreters/FilesystemCacheLog.cpp

@@ -40,6 +40,8 @@ NamesAndTypesList FilesystemCacheLogElement::getNamesAndTypes()
         {"source_file_path", std::make_shared<DataTypeString>()},
         {"file_segment_range", std::make_shared<DataTypeTuple>(types)},
         {"total_requested_range", std::make_shared<DataTypeTuple>(types)},
+        {"key", std::make_shared<DataTypeString>()},
+        {"offset", std::make_shared<DataTypeUInt64>()},
         {"size", std::make_shared<DataTypeUInt64>()},
         {"read_type", std::make_shared<DataTypeString>()},
         {"read_from_cache_attempted", std::make_shared<DataTypeUInt8>()},
@@ -60,6 +62,8 @@ void FilesystemCacheLogElement::appendToBlock(MutableColumns & columns) const
     columns[i++]->insert(source_file_path);
     columns[i++]->insert(Tuple{file_segment_range.first, file_segment_range.second});
     columns[i++]->insert(Tuple{requested_range.first, requested_range.second});
+    columns[i++]->insert(file_segment_key);
+    columns[i++]->insert(file_segment_offset);
     columns[i++]->insert(file_segment_size);
     columns[i++]->insert(typeToString(cache_type));
     columns[i++]->insert(read_from_cache_attempted);

src/Interpreters/FilesystemCacheLog.h

@@ -39,6 +39,8 @@ struct FilesystemCacheLogElement
     std::pair<size_t, size_t> file_segment_range{};
     std::pair<size_t, size_t> requested_range{};
     CacheType cache_type{};
+    std::string file_segment_key;
+    size_t file_segment_offset;
     size_t file_segment_size;
     bool read_from_cache_attempted;
     String read_buffer_id;

src/Interpreters/InterpreterSystemQuery.cpp

@@ -370,7 +370,18 @@ BlockIO InterpreterSystemQuery::execute()
             else
             {
                 auto cache = FileCacheFactory::instance().getByName(query.filesystem_cache_name).cache;
-                cache->removeAllReleasable();
+                if (query.delete_key.empty())
+                {
+                    cache->removeAllReleasable();
+                }
+                else
+                {
+                    auto key = FileCacheKey::fromKeyString(query.delete_key);
+                    if (query.delete_offset.has_value())
+                        cache->removeFileSegment(key, query.delete_offset.value());
+                    else
+                        cache->removeKey(key);
+                }
             }
             break;
         }

src/Parsers/ASTSystemQuery.cpp

@@ -210,7 +210,15 @@ void ASTSystemQuery::formatImpl(const FormatSettings & settings, FormatState &,
     else if (type == Type::DROP_FILESYSTEM_CACHE)
     {
         if (!filesystem_cache_name.empty())
+        {
             settings.ostr << (settings.hilite ? hilite_none : "") << " " << filesystem_cache_name;
+            if (!delete_key.empty())
+            {
+                settings.ostr << (settings.hilite ? hilite_none : "") << " KEY " << delete_key;
+                if (delete_offset.has_value())
+                    settings.ostr << (settings.hilite ? hilite_none : "") << " OFFSET " << delete_offset.value();
+            }
+        }
     }
     else if (type == Type::UNFREEZE)
     {

src/Parsers/ASTSystemQuery.h

@@ -107,6 +107,8 @@ public:
     UInt64 seconds{};
 
     String filesystem_cache_name;
+    std::string delete_key;
+    std::optional<size_t> delete_offset;
 
     String backup_name;
 

src/Parsers/ParserSetQuery.cpp

@@ -215,7 +215,7 @@ bool ParserSetQuery::parseNameValuePair(SettingChange & change, IParser::Pos & p
     else if (ParserKeyword("FALSE").ignore(pos, expected))
         value = std::make_shared<ASTLiteral>(Field(static_cast<UInt64>(0)));
     /// for SETTINGS disk=disk(type='s3', path='', ...)
-    else if (function_p.parse(pos, function_ast, expected) && function_ast->as<ASTFunction>()->name == "disk")
+    else if (function_p.parse(pos, function_ast, expected) && function_ast->as<ASTFunction>()->name.starts_with("disk"))
     {
         tryGetIdentifierNameInto(name, change.name);
         change.value = createFieldFromAST(function_ast);
@@ -280,7 +280,7 @@ bool ParserSetQuery::parseNameValuePairWithParameterOrDefault(
         node = std::make_shared<ASTLiteral>(Field(static_cast<UInt64>(1)));
     else if (ParserKeyword("FALSE").ignore(pos, expected))
         node = std::make_shared<ASTLiteral>(Field(static_cast<UInt64>(0)));
-    else if (function_p.parse(pos, function_ast, expected) && function_ast->as<ASTFunction>()->name == "disk")
+    else if (function_p.parse(pos, function_ast, expected) && function_ast->as<ASTFunction>()->name.starts_with("disk"))
     {
         change.name = name;
         change.value = createFieldFromAST(function_ast);

src/Parsers/ParserSystemQuery.cpp

@@ -405,7 +405,15 @@ bool ParserSystemQuery::parseImpl(IParser::Pos & pos, ASTPtr & node, Expected &
             ParserLiteral path_parser;
             ASTPtr ast;
             if (path_parser.parse(pos, ast, expected))
+            {
                 res->filesystem_cache_name = ast->as<ASTLiteral>()->value.safeGet<String>();
+                if (ParserKeyword{"KEY"}.ignore(pos, expected) && ParserIdentifier().parse(pos, ast, expected))
+                {
+                    res->delete_key = ast->as<ASTIdentifier>()->name();
+                    if (ParserKeyword{"OFFSET"}.ignore(pos, expected) && ParserLiteral().parse(pos, ast, expected))
+                        res->delete_offset = ast->as<ASTLiteral>()->value.safeGet<UInt64>();
+                }
+            }
             if (!parseQueryWithOnCluster(res, pos, expected))
                 return false;
             break;

src/Parsers/isDiskFunction.cpp

@@ -10,7 +10,7 @@ bool isDiskFunction(ASTPtr ast)
         return false;
 
     const auto * function = ast->as<ASTFunction>();
-    return function && function->name == "disk" && function->arguments->as<ASTExpressionList>();
+    return function && function->name.starts_with("disk") && function->arguments->as<ASTExpressionList>();
 }
 
 }

tests/queries/0_stateless/02808_filesystem_cache_drop_query.sh

@@ -0,0 +1,66 @@
+# Tags: no-fasttest, no-parallel, no-s3-storage, no-random-settings
+
+# set -x
+
+CUR_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+# shellcheck source=../shell_config.sh
+. "$CUR_DIR"/../shell_config.sh
+
+
+disk_name="${CLICKHOUSE_TEST_UNIQUE_NAME}"
+$CLICKHOUSE_CLIENT -nm --query """
+DROP TABLE IF EXISTS test;
+CREATE TABLE test (a Int32, b String)
+ENGINE = MergeTree() ORDER BY tuple()
+SETTINGS disk = disk_$disk_name(type = cache, max_size = '100Ki', path = ${CLICKHOUSE_TEST_UNIQUE_NAME}, disk = s3disk);
+
+INSERT INTO test SELECT 1, 'test';
+"""
+
+query_id=$RANDOM
+
+$CLICKHOUSE_CLIENT --query_id "$query_id" --query "SELECT * FROM test FORMAT Null SETTINGS enable_filesystem_cache_log = 1"
+
+${CLICKHOUSE_CLIENT} -q " system flush logs"
+
+key=$($CLICKHOUSE_CLIENT -nm --query """
+SELECT key FROM system.filesystem_cache_log WHERE query_id = '$query_id' ORDER BY size DESC LIMIT 1;
+""")
+
+offset=$($CLICKHOUSE_CLIENT -nm --query """
+SELECT offset FROM system.filesystem_cache_log WHERE query_id = '$query_id' ORDER BY size DESC LIMIT 1;
+""")
+
+$CLICKHOUSE_CLIENT -nm --query """
+SELECT count() FROM system.filesystem_cache WHERE key = '$key' AND file_segment_range_begin = $offset;
+"""
+
+$CLICKHOUSE_CLIENT -nm --query """
+SYSTEM DROP FILESYSTEM CACHE '$disk_name' KEY $key OFFSET $offset;
+"""
+
+$CLICKHOUSE_CLIENT -nm --query """
+SELECT count() FROM system.filesystem_cache WHERE key = '$key' AND file_segment_range_begin = $offset;
+"""
+
+query_id=$RANDOM$RANDOM
+
+$CLICKHOUSE_CLIENT --query_id "$query_id" --query "SELECT * FROM test FORMAT Null SETTINGS enable_filesystem_cache_log = 1"
+
+${CLICKHOUSE_CLIENT} -q " system flush logs"
+
+key=$($CLICKHOUSE_CLIENT -nm --query """
+SELECT key FROM system.filesystem_cache_log WHERE query_id = '$query_id' ORDER BY size DESC LIMIT 1;
+""")
+
+$CLICKHOUSE_CLIENT -nm --query """
+SELECT count() FROM system.filesystem_cache WHERE key = '$key';
+"""
+
+$CLICKHOUSE_CLIENT -nm --query """
+SYSTEM DROP FILESYSTEM CACHE '$disk_name' KEY $key
+"""
+
+$CLICKHOUSE_CLIENT -nm --query """
+SELECT count() FROM system.filesystem_cache WHERE key = '$key';
+"""