Merge pull request ClickHouse#70659 from ClickHouse/auto-cluster-engines

Use cluster table functions automatically if parallel replicas are enabled

Author: alexey-milovidov

src/Analyzer/Resolve/QueryAnalyzer.cpp

@@ -66,6 +66,7 @@
 #include <Analyzer/Resolve/QueryExpressionsAliasVisitor.h>
 #include <Analyzer/Resolve/ReplaceColumnsVisitor.h>
 #include <Analyzer/Resolve/TableExpressionsAliasVisitor.h>
+#include <Analyzer/Resolve/TableFunctionsWithClusterAlternativesVisitor.h>
 #include <Analyzer/Resolve/TypoCorrection.h>
 
 #include <Planner/PlannerActionsVisitor.h>
@@ -5564,6 +5565,11 @@ void QueryAnalyzer::resolveQuery(const QueryTreeNodePtr & query_node, Identifier
     TableExpressionsAliasVisitor table_expressions_visitor(scope);
     table_expressions_visitor.visit(query_node_typed.getJoinTree());
 
+    TableFunctionsWithClusterAlternativesVisitor table_function_visitor;
+    table_function_visitor.visit(query_node);
+    if (!table_function_visitor.shouldReplaceWithClusterAlternatives() && scope.context->hasQueryContext())
+        scope.context->getQueryContext()->setSetting("parallel_replicas_for_cluster_engines", false);
+
     initializeQueryJoinTreeNode(query_node_typed.getJoinTree(), scope);
     scope.aliases.alias_name_to_table_expression_node = std::move(transitive_aliases);
 

src/Analyzer/Resolve/TableFunctionsWithClusterAlternativesVisitor.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <Analyzer/InDepthQueryTreeVisitor.h>
+#include <Analyzer/TableFunctionNode.h>
+#include <TableFunctions/TableFunctionFactory.h>
+#include <TableFunctions/TableFunctionFile.h>
+
+namespace DB
+{
+
+class TableFunctionsWithClusterAlternativesVisitor : public InDepthQueryTreeVisitor<TableFunctionsWithClusterAlternativesVisitor, /*const_visitor=*/true>
+{
+public:
+    void visitImpl(const QueryTreeNodePtr & node)
+    {
+        if (node->getNodeType() == QueryTreeNodeType::TABLE_FUNCTION)
+            ++table_function_count;
+        else if (node->getNodeType() == QueryTreeNodeType::TABLE)
+            ++table_count;
+    }
+
+    bool needChildVisit(const QueryTreeNodePtr &, const QueryTreeNodePtr &) { return true; }
+    bool shouldReplaceWithClusterAlternatives() const { return (table_count + table_function_count) == 1; }
+
+private:
+    size_t table_count = 0;
+    size_t table_function_count = 0;
+};
+
+}

src/Core/Settings.cpp

@@ -5824,11 +5824,12 @@ Build local plan for local replica
     DECLARE(Bool, parallel_replicas_index_analysis_only_on_coordinator, true, R"(
 Index analysis done only on replica-coordinator and skipped on other replicas. Effective only with enabled parallel_replicas_local_plan
 )", BETA) \
-    \
     DECLARE(Bool, parallel_replicas_only_with_analyzer, true, R"(
 The analyzer should be enabled to use parallel replicas. With disabled analyzer query execution fallbacks to local execution, even if parallel reading from replicas is enabled. Using parallel replicas without the analyzer enabled is not supported
 )", BETA) \
-    \
+    DECLARE(Bool, parallel_replicas_for_cluster_engines, true, R"(
+Replace table function engines with their -Cluster alternatives
+)", EXPERIMENTAL) \
     DECLARE(Bool, allow_experimental_analyzer, true, R"(
 Allow new query analyzer.
 )", IMPORTANT) ALIAS(enable_analyzer) \

src/Core/SettingsChangesHistory.cpp

@@ -70,6 +70,7 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
         {
             {"use_page_cache_with_distributed_cache", false, false, "New setting"},
             {"use_query_condition_cache", false, false, "New setting."},
+            {"parallel_replicas_for_cluster_engines", false, true, "New setting."},
         });
         addSettingsChanges(settings_changes_history, "25.2",
         {

src/Databases/DatabaseS3.cpp

@@ -138,7 +138,7 @@ StoragePtr DatabaseS3::getTableImpl(const String & name, ContextPtr context_) co
         return nullptr;
 
     /// TableFunctionS3 throws exceptions, if table cannot be created.
-    auto table_storage = table_function->execute(function, context_, name);
+    auto table_storage = table_function->execute(function, context_, name, /*cached_columns_=*/{}, /*use_global_context=*/false, /*is_insert_query=*/true);
     if (table_storage)
         addTable(name, table_storage);
 

src/Interpreters/InterpreterCreateQuery.cpp

@@ -1924,7 +1924,7 @@ bool InterpreterCreateQuery::doCreateTable(ASTCreateQuery & create,
         /// In case of CREATE AS table_function() query we should use global context
         /// in storage creation because there will be no query context on server startup
         /// and because storage lifetime is bigger than query context lifetime.
-        res = table_function->execute(table_function_ast, getContext(), create.getTable(), properties.columns, /*use_global_context=*/true);
+        res = table_function->execute(table_function_ast, getContext(), create.getTable(), properties.columns, /*use_global_context=*/true, /*is_insert_query=*/true);
         res->renameInMemory({create.getDatabase(), create.getTable(), create.uuid});
     }
     else

src/Storages/IStorageCluster.cpp

@@ -24,9 +24,11 @@
 #include <Storages/SelectQueryInfo.h>
 #include <Storages/StorageDictionary.h>
 
+#include <algorithm>
 #include <memory>
 #include <string>
 
+
 namespace DB
 {
 namespace Setting
@@ -35,6 +37,9 @@ namespace Setting
     extern const SettingsBool async_query_sending_for_remote;
     extern const SettingsBool async_socket_for_remote;
     extern const SettingsBool skip_unavailable_shards;
+    extern const SettingsBool parallel_replicas_local_plan;
+    extern const SettingsString cluster_for_parallel_replicas;
+    extern const SettingsNonZeroUInt64 max_parallel_replicas;
 }
 
 IStorageCluster::IStorageCluster(
@@ -183,29 +188,45 @@ void ReadFromCluster::initializePipeline(QueryPipelineBuilder & pipeline, const
     auto new_context = updateSettings(context->getSettingsRef());
     const auto & current_settings = new_context->getSettingsRef();
     auto timeouts = ConnectionTimeouts::getTCPTimeoutsWithFailover(current_settings);
+
+    auto max_replicas_to_use = static_cast<UInt64>(cluster->getShardsInfo().size());
+    if (current_settings[Setting::max_parallel_replicas] > 1)
+        max_replicas_to_use = std::min(max_replicas_to_use, current_settings[Setting::max_parallel_replicas].value);
+
     for (const auto & shard_info : cluster->getShardsInfo())
     {
-        auto try_results = shard_info.pool->getMany(timeouts, current_settings, PoolMode::GET_MANY);
-        for (auto & try_result : try_results)
-        {
-            auto remote_query_executor = std::make_shared<RemoteQueryExecutor>(
-                std::vector<IConnectionPool::Entry>{try_result},
-                queryToString(query_to_send),
-                getOutputHeader(),
-                new_context,
-                /*throttler=*/nullptr,
-                scalars,
-                Tables(),
-                processed_stage,
-                extension);
-
-            remote_query_executor->setLogger(log);
-            pipes.emplace_back(std::make_shared<RemoteSource>(
-                remote_query_executor,
-                add_agg_info,
-                current_settings[Setting::async_socket_for_remote],
-                current_settings[Setting::async_query_sending_for_remote]));
-        }
+        if (pipes.size() >= max_replicas_to_use)
+            break;
+
+        /// We're taking all replicas as shards,
+        /// so each shard will have only one address to connect to.
+        auto try_results = shard_info.pool->getMany(
+            timeouts,
+            current_settings,
+            PoolMode::GET_ONE,
+            {},
+            /*skip_unavailable_endpoints=*/true);
+
+        if (try_results.empty())
+            continue;
+
+        auto remote_query_executor = std::make_shared<RemoteQueryExecutor>(
+            std::vector<IConnectionPool::Entry>{try_results.front()},
+            queryToString(query_to_send),
+            getOutputHeader(),
+            new_context,
+            /*throttler=*/nullptr,
+            scalars,
+            Tables(),
+            processed_stage,
+            extension);
+
+        remote_query_executor->setLogger(log);
+        pipes.emplace_back(std::make_shared<RemoteSource>(
+            remote_query_executor,
+            add_agg_info,
+            current_settings[Setting::async_socket_for_remote],
+            current_settings[Setting::async_query_sending_for_remote]));
     }
 
     auto pipe = Pipe::unitePipes(std::move(pipes));

src/Storages/ObjectStorage/StorageObjectStorageCluster.cpp

@@ -1,6 +1,8 @@
 #include "Storages/ObjectStorage/StorageObjectStorageCluster.h"
 
 #include <Common/Exception.h>
+#include <Common/StringUtils.h>
+
 #include <Core/Settings.h>
 #include <Formats/FormatFactory.h>
 #include <Parsers/queryToString.h>
@@ -10,7 +12,7 @@
 #include <Storages/VirtualColumnUtils.h>
 #include <Storages/ObjectStorage/Utils.h>
 #include <Storages/ObjectStorage/StorageObjectStorageSource.h>
-#include <Storages/extractTableFunctionArgumentsFromSelectQuery.h>
+#include <Storages/extractTableFunctionFromSelectQuery.h>
 
 
 namespace DB
@@ -86,7 +88,16 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
     const DB::StorageSnapshotPtr & storage_snapshot,
     const ContextPtr & context)
 {
-    ASTExpressionList * expression_list = extractTableFunctionArgumentsFromSelectQuery(query);
+    auto * table_function = extractTableFunctionFromSelectQuery(query);
+    if (!table_function)
+    {
+        throw Exception(
+            ErrorCodes::LOGICAL_ERROR,
+            "Expected SELECT query from table function {}, got '{}'",
+            configuration->getEngineName(), queryToString(query));
+    }
+
+    auto * expression_list = table_function->arguments->as<ASTExpressionList>();
     if (!expression_list)
     {
         throw Exception(
@@ -105,25 +116,35 @@ void StorageObjectStorageCluster::updateQueryToSendIfNeeded(
             configuration->getEngineName());
     }
 
-    ASTPtr cluster_name_arg = args.front();
-    args.erase(args.begin());
-    configuration->addStructureAndFormatToArgsIfNeeded(args, structure, configuration->format, context, /*with_structure=*/true);
-    args.insert(args.begin(), cluster_name_arg);
+    if (!endsWith(table_function->name, "Cluster"))
+        configuration->addStructureAndFormatToArgsIfNeeded(args, structure, configuration->format, context, /*with_structure=*/true);
+    else
+    {
+        ASTPtr cluster_name_arg = args.front();
+        args.erase(args.begin());
+        configuration->addStructureAndFormatToArgsIfNeeded(args, structure, configuration->format, context, /*with_structure=*/true);
+        args.insert(args.begin(), cluster_name_arg);
+    }
 }
 
 RemoteQueryExecutor::Extension StorageObjectStorageCluster::getTaskIteratorExtension(
     const ActionsDAG::Node * predicate, const ContextPtr & local_context) const
 {
     auto iterator = StorageObjectStorageSource::createFileIterator(
         configuration, configuration->getQuerySettings(local_context), object_storage, /* distributed_processing */false,
-        local_context, predicate, virtual_columns, nullptr, local_context->getFileProgressCallback());
+        local_context, predicate, virtual_columns, nullptr, local_context->getFileProgressCallback(), /*ignore_archive_globs=*/true);
 
     auto callback = std::make_shared<std::function<String()>>([iterator]() mutable -> String
     {
         auto object_info = iterator->next(0);
-        if (object_info)
-            return object_info->getPath();
-        return "";
+        if (!object_info)
+            return "";
+
+        auto archive_object_info = std::dynamic_pointer_cast<StorageObjectStorageSource::ArchiveIterator::ObjectInfoInArchive>(object_info);
+        if (archive_object_info)
+            return archive_object_info->getPathToArchive();
+
+        return object_info->getPath();
     });
     return RemoteQueryExecutor::Extension{ .task_iterator = std::move(callback) };
 }

src/Storages/ObjectStorage/StorageObjectStorageSource.cpp

@@ -129,17 +129,25 @@ std::shared_ptr<IObjectIterator> StorageObjectStorageSource::createFileIterator(
     const ActionsDAG::Node * predicate,
     const NamesAndTypesList & virtual_columns,
     ObjectInfos * read_keys,
-    std::function<void(FileProgress)> file_progress_callback)
+    std::function<void(FileProgress)> file_progress_callback,
+    bool ignore_archive_globs)
 {
+    const bool is_archive = configuration->isArchive();
+
     if (distributed_processing)
-        return std::make_shared<ReadTaskIterator>(local_context->getReadTaskCallback(), local_context->getSettingsRef()[Setting::max_threads]);
+    {
+        auto distributed_iterator = std::make_unique<ReadTaskIterator>(local_context->getReadTaskCallback(), local_context->getSettingsRef()[Setting::max_threads]);
+
+        if (is_archive)
+            return std::make_shared<ArchiveIterator>(object_storage, configuration, std::move(distributed_iterator), local_context, nullptr);
+
+        return distributed_iterator;
+    }
 
     if (configuration->isNamespaceWithGlobs())
         throw Exception(ErrorCodes::BAD_ARGUMENTS,
                         "Expression can not have wildcards inside {} name", configuration->getNamespaceType());
 
-    const bool is_archive = configuration->isArchive();
-
     std::unique_ptr<IObjectIterator> iterator;
     if (configuration->isPathWithGlobs())
     {
@@ -190,7 +198,7 @@ std::shared_ptr<IObjectIterator> StorageObjectStorageSource::createFileIterator(
 
     if (is_archive)
     {
-        return std::make_shared<ArchiveIterator>(object_storage, configuration, std::move(iterator), local_context, read_keys);
+        return std::make_shared<ArchiveIterator>(object_storage, configuration, std::move(iterator), local_context, read_keys, ignore_archive_globs);
     }
 
     return iterator;
@@ -539,9 +547,8 @@ std::unique_ptr<ReadBufferFromFileBase> StorageObjectStorageSource::createReadBu
     if (!object_info.metadata)
     {
         if (!use_cache)
-        {
             return object_storage->readObject(StoredObject(object_info.getPath()), read_settings);
-        }
+
         object_info.metadata = object_storage->getObjectMetadata(object_info.getPath());
     }
 
@@ -646,6 +653,7 @@ std::unique_ptr<ReadBufferFromFileBase> StorageObjectStorageSource::createReadBu
         impl->setReadUntilEnd();
         impl->prefetch(DEFAULT_PREFETCH_PRIORITY);
     }
+
     return impl;
 }
 
@@ -828,7 +836,7 @@ StorageObjectStorage::ObjectInfoPtr StorageObjectStorageSource::KeysIterator::ne
     {
         size_t current_index = index.fetch_add(1, std::memory_order_relaxed);
         if (current_index >= keys.size())
-            return {};
+            return nullptr;
 
         auto key = keys[current_index];
 
@@ -898,16 +906,24 @@ StorageObjectStorageSource::ReadTaskIterator::ReadTaskIterator(
     for (auto & key_future : keys)
     {
         auto key = key_future.get();
-        if (!key.empty())
-            buffer.emplace_back(std::make_shared<ObjectInfo>(key, std::nullopt));
+        if (key.empty())
+            continue;
+
+        buffer.emplace_back(std::make_shared<ObjectInfo>(key, std::nullopt));
     }
 }
 
 StorageObjectStorage::ObjectInfoPtr StorageObjectStorageSource::ReadTaskIterator::next(size_t)
 {
     size_t current_index = index.fetch_add(1, std::memory_order_relaxed);
     if (current_index >= buffer.size())
-        return std::make_shared<ObjectInfo>(callback());
+    {
+        auto key = callback();
+        if (key.empty())
+            return nullptr;
+
+        return std::make_shared<ObjectInfo>(key, std::nullopt);
+    }
 
     return buffer[current_index];
 }
@@ -938,7 +954,8 @@ StorageObjectStorageSource::ArchiveIterator::ArchiveIterator(
     ConfigurationPtr configuration_,
     std::unique_ptr<IObjectIterator> archives_iterator_,
     ContextPtr context_,
-    ObjectInfos * read_keys_)
+    ObjectInfos * read_keys_,
+    bool ignore_archive_globs_)
     : WithContext(context_)
     , object_storage(object_storage_)
     , is_path_in_archive_with_globs(configuration_->isPathInArchiveWithGlobs())
@@ -947,6 +964,7 @@ StorageObjectStorageSource::ArchiveIterator::ArchiveIterator(
     , log(getLogger("ArchiveIterator"))
     , path_in_archive(is_path_in_archive_with_globs ? "" : configuration_->getPathInArchive())
     , read_keys(read_keys_)
+    , ignore_archive_globs(ignore_archive_globs_)
 {
 }
 
@@ -977,12 +995,15 @@ ObjectInfoPtr StorageObjectStorageSource::ArchiveIterator::next(size_t processor
                 if (!archive_object)
                     return {};
 
+                if (!archive_object->metadata)
+                    archive_object->metadata = object_storage->getObjectMetadata(archive_object->getPath());
+
                 archive_reader = createArchiveReader(archive_object);
                 file_enumerator = archive_reader->firstFile();
                 if (!file_enumerator)
                     continue;
             }
-            else if (!file_enumerator->nextFile())
+            else if (!file_enumerator->nextFile() || ignore_archive_globs)
             {
                 file_enumerator.reset();
                 continue;