bool is_remote_function,

ContextPtr context,

const Settings & settings,

const StorageID & main_table,

ASTPtr additional_filter_ast,

LoggerPtr log,

const DistributedSettings * distributed_settings)

{

ClientInfo new_client_info = context->getClientInfo();

Settings new_settings {settings};

new_settings[Setting::queue_max_wait_ms] = Cluster::saturate(new_settings[Setting::queue_max_wait_ms], settings[Setting::max_execution_time]);

/// In case of interserver mode we should reset initial_user for remote() function to use passed user from the query.

if (is_remote_function)

{

const auto & address = cluster.getShardsAddresses().front().front();

new_client_info.initial_user = address.user;

}

/// If "secret" (in remote_servers) is not in use,

/// user on the shard is not the same as the user on the initiator,

/// hence per-user limits should not be applied.

const bool interserver_mode = !cluster.getSecret().empty();

if (!interserver_mode)

{

/// Does not matter on remote servers, because queries are sent under different user.

new_settings[Setting::max_concurrent_queries_for_user] = 0;

new_settings[Setting::max_memory_usage_for_user] = 0;

/// Set as unchanged to avoid sending to remote server.

new_settings[Setting::max_concurrent_queries_for_user].changed = false;

new_settings[Setting::max_memory_usage_for_user].changed = false;

}

if (settings[Setting::force_optimize_skip_unused_shards_nesting] && settings[Setting::force_optimize_skip_unused_shards])

{

if (new_settings[Setting::force_optimize_skip_unused_shards_nesting] == 1)

{

new_settings[Setting::force_optimize_skip_unused_shards] = false;

new_settings[Setting::force_optimize_skip_unused_shards].changed = false;

if (log)

LOG_TRACE(log, "Disabling force_optimize_skip_unused_shards for nested queries (force_optimize_skip_unused_shards_nesting exceeded)");

}

else

{

--new_settings[Setting::force_optimize_skip_unused_shards_nesting].value;

new_settings[Setting::force_optimize_skip_unused_shards_nesting].changed = true;

if (log)

LOG_TRACE(

log, "force_optimize_skip_unused_shards_nesting is now {}", new_settings[Setting::force_optimize_skip_unused_shards_nesting].value);

}

}

if (settings[Setting::optimize_skip_unused_shards_nesting] && settings[Setting::optimize_skip_unused_shards])

{

if (new_settings[Setting::optimize_skip_unused_shards_nesting] == 1)

{

new_settings[Setting::optimize_skip_unused_shards] = false;

new_settings[Setting::optimize_skip_unused_shards].changed = false;

if (log)

LOG_TRACE(log, "Disabling optimize_skip_unused_shards for nested queries (optimize_skip_unused_shards_nesting exceeded)");

}

else

{

--new_settings[Setting::optimize_skip_unused_shards_nesting].value;

new_settings[Setting::optimize_skip_unused_shards_nesting].changed = true;

if (log)

LOG_TRACE(log, "optimize_skip_unused_shards_nesting is now {}", new_settings[Setting::optimize_skip_unused_shards_nesting].value);

}

}

if (!settings[Setting::skip_unavailable_shards].changed && distributed_settings)

{

new_settings[Setting::skip_unavailable_shards] = (*distributed_settings)[DistributedSetting::skip_unavailable_shards].value;

new_settings[Setting::skip_unavailable_shards].changed = true;

}

if (!settings[Setting::skip_unavailable_shards_mode].changed && distributed_settings)

{

new_settings[Setting::skip_unavailable_shards_mode] = (*distributed_settings)[DistributedSetting::skip_unavailable_shards_mode].value;

new_settings[Setting::skip_unavailable_shards_mode].changed = true;

}

if (settings[Setting::offset])

{

new_settings[Setting::offset] = 0;

new_settings[Setting::offset].changed = false;

}

if (settings[Setting::limit])

{

new_settings[Setting::limit] = 0;

new_settings[Setting::limit].changed = false;

}

/// Setting additional_table_filters may be applied to Distributed table.

/// In case if query is executed up to WithMergableState on remote shard, it is impossible to filter on initiator.

/// We need to propagate the setting, but change the table name from distributed to source.

///

/// Here we don't try to analyze setting again. In case if query_info->additional_filter_ast is not empty, some filter was applied.

/// It's just easier to add this filter for a source table.

if (additional_filter_ast)

{

Tuple tuple;

tuple.push_back(main_table.getShortName());

tuple.push_back(additional_filter_ast->formatWithSecretsOneLine());

new_settings[Setting::additional_table_filters].value.push_back(std::move(tuple));

}

/// disable parallel replicas if cluster contains only shards with 1 replica

if (context->canUseTaskBasedParallelReplicas())

{

bool disable_parallel_replicas = false;

if (is_remote_function)

{

if (cluster.getName().empty()) // disable parallel replicas with remote() table functions w/o configured cluster

disable_parallel_replicas = true;

else

new_settings[Setting::cluster_for_parallel_replicas] = cluster.getName();

}

if (!disable_parallel_replicas)

{

disable_parallel_replicas = true;

for (const auto & shard : cluster.getShardsInfo())

{

if (shard.getAllNodeCount() > 1)

{

disable_parallel_replicas = false;

break;

}

}

}

if (disable_parallel_replicas)

new_settings[Setting::allow_experimental_parallel_reading_from_replicas] = 0;

}

if (settings[Setting::max_execution_time_leaf].totalMicroseconds() > 0)

{

/// Replace 'max_execution_time' of this sub-query with 'max_execution_time_leaf' and 'timeout_overflow_mode'

/// with 'timeout_overflow_mode_leaf'

new_settings[Setting::max_execution_time] = settings[Setting::max_execution_time_leaf];

new_settings[Setting::timeout_overflow_mode] = settings[Setting::timeout_overflow_mode_leaf];

}

/// in case of parallel replicas custom key use round robing load balancing

/// so custom key partitions will be spread over nodes in round-robin fashion

if (context->canUseParallelReplicasCustomKeyForCluster(cluster) && !settings[Setting::load_balancing].changed)

{

new_settings[Setting::load_balancing] = LoadBalancing::ROUND_ROBIN;

}

/// disable plan serialization for sample and custom key modes

/// until filter generation for these modes are done on query plan level

if (context->canUseOffsetParallelReplicas())

new_settings[Setting::serialize_query_plan] = false;

auto new_context = Context::createCopy(context);

new_context->setSettings(new_settings);

new_context->setClientInfo(new_client_info);

if (context->canUseParallelReplicasCustomKeyForCluster(cluster))

new_context->disableOffsetParallelReplicas();

return new_context;

{

return updateSettingsAndClientInfoForCluster(cluster,

/* is_remote_function= */ false,

context,

settings,

main_table,

/* additional_filter_ast= */ {},

/* log= */ {},

/* distributed_settings= */ {});

{

const Settings & settings = context->getSettingsRef();

ThrottlerPtr user_level_throttler;

if (auto process_list_element = context->getProcessListElement())

user_level_throttler = process_list_element->getUserNetworkThrottler();

/// Network bandwidth limit, if needed.

ThrottlerPtr throttler;

if (settings[Setting::max_network_bandwidth] || settings[Setting::max_network_bytes])

{

throttler = std::make_shared<Throttler>(

settings[Setting::max_network_bandwidth],

settings[Setting::max_network_bytes],

"Limit for bytes to send or receive over network exceeded.",

user_level_throttler);

}

else

throttler = user_level_throttler;

return throttler;

{

if (!context->canUseParallelReplicasCustomKeyForCluster(cluster))

return {};

const auto & settings = context->getSettingsRef();

auto custom_key_ast = parseCustomKeyForTable(settings[Setting::parallel_replicas_custom_key], *context);

if (custom_key_ast == nullptr)

return {};

return [my_custom_key_ast = std::move(custom_key_ast),

column_description = columns,

custom_key_type = settings[Setting::parallel_replicas_mode].value,

custom_key_range_lower = settings[Setting::parallel_replicas_custom_key_range_lower].value,

custom_key_range_upper = settings[Setting::parallel_replicas_custom_key_range_upper].value,

query_context = context,

replica_count = cluster.getShardsInfo().front().per_replica_pools.size()](uint64_t replica_num) -> ASTPtr

{

return getCustomKeyFilterForParallelReplica(

replica_count,

replica_num - 1,

my_custom_key_ast,

{custom_key_type, custom_key_range_lower, custom_key_range_upper},

column_description,

query_context);

};

QueryPlan & query_plan,

SharedHeader header,

QueryProcessingStage::Enum processed_stage,

const StorageID & main_table,

const ASTPtr & table_func_ptr,

SelectStreamFactory & stream_factory,

LoggerPtr log,

ContextPtr context,

const SelectQueryInfo & query_info,

const ExpressionActionsPtr & sharding_key_expr,

const std::string & sharding_key_column_name,

const DistributedSettings & distributed_settings,

AdditionalShardFilterGenerator shard_filter_generator,

bool is_remote_function)

{

const Settings & settings = context->getSettingsRef();

if (settings[Setting::max_distributed_depth] && context->getClientInfo().distributed_depth >= settings[Setting::max_distributed_depth])

throw Exception(ErrorCodes::TOO_LARGE_DISTRIBUTED_DEPTH, "Maximum distributed depth exceeded");

const ClusterPtr & not_optimized_cluster = query_info.cluster;

std::vector<QueryPlanPtr> plans;

SelectStreamFactory::Shards remote_shards;

auto cluster = query_info.getCluster();

auto new_context = updateSettingsAndClientInfoForCluster(*cluster, is_remote_function, context,

settings, main_table, query_info.additional_filter_ast, log, &distributed_settings);

if (context->getSettingsRef()[Setting::allow_experimental_parallel_reading_from_replicas].value

&& context->getSettingsRef()[Setting::allow_experimental_parallel_reading_from_replicas].value

!= new_context->getSettingsRef()[Setting::allow_experimental_parallel_reading_from_replicas].value)

{

LOG_TRACE(

log,

"Parallel reading from replicas is disabled for cluster. There are no shards with more than 1 replica: cluster={}",

cluster->getName());

}

new_context->increaseDistributedDepth();

const size_t shards = cluster->getShardCount();

ProfileEvents::increment(ProfileEvents::Shards, shards);

/// Tracker is shared between local-missing-table skip path in SelectStreamFactory and

/// remote unavailable-shard skip path in ReadFromRemote so max_skip_unavailable_shards_num

/// and max_skip_unavailable_shards_ratio are enforced uniformly across both paths.

UnavailableShardTrackerPtr unavailable_shard_tracker;

{

const auto & new_settings_ref = new_context->getSettingsRef();

if (new_settings_ref[Setting::skip_unavailable_shards])

{

size_t max_num = new_settings_ref[Setting::max_skip_unavailable_shards_num];

Float64 max_ratio = static_cast<double>(new_settings_ref[Setting::max_skip_unavailable_shards_ratio]);

if (max_num > 0 || max_ratio > 0)

unavailable_shard_tracker = std::make_shared<UnavailableShardTracker>(shards, max_num, max_ratio);

}

}

if (context->getSettingsRef()[Setting::allow_experimental_analyzer])

{

for (size_t i = 0, s = cluster->getShardsInfo().size(); i < s; ++i)

{

const auto & shard_info = cluster->getShardsInfo()[i];

auto query_for_shard = query_info.query_tree->clone();

if (sharding_key_expr && query_info.optimized_cluster && settings[Setting::optimize_skip_unused_shards_rewrite_in] && shards > 1 &&

/// TODO: support composite sharding key

sharding_key_expr->getRequiredColumns().size() == 1)

{

OptimizeShardingKeyRewriteInVisitor::Data visitor_data{

sharding_key_expr,

sharding_key_column_name,

shard_info,

not_optimized_cluster->getSlotToShard(),

};

optimizeShardingKeyRewriteIn(query_for_shard, std::move(visitor_data), new_context);

}

// decide for each shard if parallel reading from replicas should be enabled

// according to settings and number of replicas declared per shard

const auto & addresses = cluster->getShardsAddresses().at(i);

const bool parallel_replicas_enabled = addresses.size() > 1 && new_context->canUseTaskBasedParallelReplicas();

stream_factory.createForShard(

shard_info,

query_for_shard,

main_table,

table_func_ptr,

new_context,

plans,

remote_shards,

static_cast<UInt32>(shards),

parallel_replicas_enabled,

shard_filter_generator,

unavailable_shard_tracker);

}

}

else

{

for (size_t i = 0, s = cluster->getShardsInfo().size(); i < s; ++i)

{

const auto & shard_info = cluster->getShardsInfo()[i];

ASTPtr query_ast_for_shard = query_info.query->clone();

if (sharding_key_expr && query_info.optimized_cluster && settings[Setting::optimize_skip_unused_shards_rewrite_in] && shards > 1 &&

/// TODO: support composite sharding key

sharding_key_expr->getRequiredColumns().size() == 1)

{

OptimizeShardingKeyRewriteInVisitor::Data visitor_data{

sharding_key_expr,

sharding_key_column_name,

shard_info,

not_optimized_cluster->getSlotToShard(),

};

OptimizeShardingKeyRewriteInVisitor visitor(visitor_data);

visitor.visit(query_ast_for_shard);

}

// decide for each shard if parallel reading from replicas should be enabled

// according to settings and number of replicas declared per shard

const auto & addresses = cluster->getShardsAddresses().at(i);

bool parallel_replicas_enabled = addresses.size() > 1 && context->canUseTaskBasedParallelReplicas();

stream_factory.createForShard(

shard_info,

query_ast_for_shard,

main_table,

table_func_ptr,

new_context,

plans,

remote_shards,

static_cast<UInt32>(shards),

parallel_replicas_enabled,

shard_filter_generator,

unavailable_shard_tracker);

}

}

if (!remote_shards.empty())

{

Scalars scalars = context->hasQueryContext() ? context->getQueryContext()->getScalars() : Scalars{};

scalars.emplace(

"_shard_count", Block{{DataTypeUInt32().createColumnConst(1, shards), std::make_shared<DataTypeUInt32>(), "_shard_count"}});

auto external_tables = context->getExternalTables();

auto plan = std::make_unique<QueryPlan>();

auto read_from_remote = std::make_unique<ReadFromRemote>(

std::move(remote_shards),

header,

processed_stage,

main_table,

table_func_ptr,

new_context,

getThrottler(context),

std::move(scalars),

std::move(external_tables),

log,

shards,

query_info.storage_limits,

not_optimized_cluster->getName(),

std::move(unavailable_shard_tracker));

read_from_remote->setStepDescription("Read from remote replica");

plan->addStep(std::move(read_from_remote));

plan->addInterpreterContext(new_context);

plans.emplace_back(std::move(plan));

}

if (plans.empty())

return;

if (plans.size() == 1)

{

query_plan = std::move(*plans.front());

return;

}

SharedHeaders input_headers;

input_headers.reserve(plans.size());

for (auto & plan : plans)

input_headers.emplace_back(plan->getCurrentHeader());

auto union_step = std::make_unique<UnionStep>(std::move(input_headers));

query_plan.unitePlans(std::move(union_step), std::move(plans));

{

const auto & settings = context->getSettingsRef();

auto context_mutable = Context::createCopy(context);

/// check hedged connections setting

if (settings[Setting::use_hedged_requests].value)

{

LOG_INFO(

logger,

"Disabling 'use_hedged_requests' in favor of 'enable_parallel_replicas'. "

"Hedged connections are not used for parallel reading from replicas");

/// disable hedged connections -> parallel replicas uses own logic to choose replicas

context_mutable->setSetting("use_hedged_requests", Field{false});

}

/// If parallel replicas executed over distributed table i.e. in scope of a shard,

/// currently, local plan for parallel replicas is not created.

/// Having local plan is prerequisite to use projection with parallel replicas.

/// So, currently, we disable projection support with parallel replicas when reading over distributed table with parallel replicas

/// Otherwise, it can lead to incorrect results, in particular with use of implicit projection min_max_count

if (shard_num > 0 && settings[Setting::parallel_replicas_support_projection].value)

{

LOG_TRACE(

logger,

"Disabling 'parallel_replicas_support_projection'. Currently, it's not supported for queries with parallel replicas over distributed tables");

context_mutable->setSetting("parallel_replicas_support_projection", Field{false});

}

return context_mutable;

{

/// check cluster for parallel replicas

auto not_optimized_cluster = context->getClusterForParallelReplicas();

auto scalars = context->hasQueryContext() ? context->getQueryContext()->getScalars() : Scalars{};

UInt64 shard_num = 0; /// shard_num is 1-based, so 0 - no shard specified

const auto it = scalars.find("_shard_num");

if (it != scalars.end())

{

const Block & block = it->second;

const auto & column = block.safeGetByPosition(0).column;

shard_num = column->getUInt(0);

}

ClusterPtr new_cluster = not_optimized_cluster;

/// if got valid shard_num from query initiator, then parallel replicas scope is the specified shard

/// shards are numbered in order of appearance in the cluster config

if (shard_num > 0)

{

const auto shard_count = not_optimized_cluster->getShardCount();

if (shard_num > shard_count)

throw Exception(

ErrorCodes::LOGICAL_ERROR,

"Shard number is greater than shard count: shard_num={} shard_count={} cluster={}",

shard_num,

shard_count,

not_optimized_cluster->getName());

chassert(shard_count == not_optimized_cluster->getShardsAddresses().size());

LOG_DEBUG(logger, "Parallel replicas query in shard scope: shard_num={} cluster={}", shard_num, not_optimized_cluster->getName());

// get cluster for shard specified by shard_num

// shard_num is 1-based, but getClusterWithSingleShard expects 0-based index

new_cluster = not_optimized_cluster->getClusterWithSingleShard(shard_num - 1);

}

else

{

if (not_optimized_cluster->getShardCount() > 1)

throw DB::Exception(

ErrorCodes::UNEXPECTED_CLUSTER,

"`cluster_for_parallel_replicas` setting refers to cluster with several shards. Expected a cluster with one shard");

}

return {new_cluster, shard_num};

{

const String cluster_name = context->getSettingsRef()[Setting::cluster_for_parallel_replicas];

ReplicasInfo replicas_info;

#if CLICKHOUSE_CLOUD

/// The shared catalog cluster is exposed in `system.clusters` both under its plain name and under the

/// `all_groups.` prefix, and both report the same `is_active` data - accept either spelling here.

if (SharedDatabaseCatalog::initialized())

{

const String & catalog_cluster_name = SharedDatabaseCatalog::instance().getClusterName();

if (cluster_name == catalog_cluster_name

|| cluster_name == SharedDatabaseCatalog::ALL_GROUPS_CLUSTER_PREFIX + catalog_cluster_name)

replicas_info = SharedDatabaseCatalog::instance().getClusterWithReplicasInfo(cluster_name).second;

}

#endif

if (replicas_info.replicas.empty())

{

/// A `Replicated` database is exposed in `system.clusters` both as `<db>` and as `all_groups.<db>`;

/// the latter resolves to the same database after stripping the prefix (see `tryGetReplicatedDatabaseCluster`).

/// Strip it here too, otherwise an `all_groups.<db>` cluster gets no liveness data and we fall back to

/// counting inactive replicas again.

String database_name = cluster_name;

bool all_groups = false;

static constexpr std::string_view all_groups_prefix = DatabaseReplicated::ALL_GROUPS_CLUSTER_PREFIX;

if (database_name.starts_with(all_groups_prefix))

{

database_name = database_name.substr(all_groups_prefix.size());

all_groups = true;

}

if (auto database = DatabaseCatalog::instance().tryGetDatabase(database_name))

if (const auto * replicated = typeid_cast<const DatabaseReplicated *>(database.get()))

{

/// `cluster_for_parallel_replicas` is resolved by `Context::getCluster`, which prefers a

/// configured or discovered cluster over a `Replicated` database of the same name. When such a

/// cluster shadows the database, the resolved `cluster` is not the database's own cluster: its

/// replica names do not match the database's ZooKeeper nodes, so `tryGetReplicasInfo` would

/// report every replica inactive and the coordinator would collapse onto the local replica.

/// Only trust the database liveness when the resolved cluster really is this database's own

/// cluster (`system.clusters` likewise reports such configured clusters with unknown `is_active`).

const ClusterPtr database_cluster = all_groups ? replicated->tryGetAllGroupsCluster() : replicated->tryGetCluster();

if (database_cluster == cluster)

replicas_info = replicated->tryGetReplicasInfo(cluster);

}

}

std::vector<bool> is_active;

is_active.reserve(replicas_info.replicas.size());

for (const auto & replica : replicas_info.replicas)

is_active.push_back(replica.is_active);

return is_active;

{

const auto & settings = context->getSettingsRef();

const auto & shard = cluster->getShardsInfo().at(0);

/// Exclude replicas that are known to be inactive (e.g. stale registrations left over after autoscaling).

/// The reading coordinator distributes mark segments by hashing over the number of replicas, so counting

/// replicas that never participate leaves "phantom" segments that only the source replica picks up, which

/// produces a severe work-distribution skew. See `is_active` in `system.clusters`.

std::vector<bool> is_active = getActiveReplicasForParallelReplicas(context, cluster);

if (!is_active.empty() && is_active.size() != shard.getAllNodeCount())

is_active.clear(); /// Liveness does not match the cluster definition; fall back to using all replicas.

if (!is_active.empty())

{

/// Identify the local replica the same way `findLocalReplicaIndexAndUpdatePools` does (host name + port).

const auto & addresses = cluster->getShardsAddresses().at(0);

std::optional<size_t> local_replica_index;

for (size_t i = 0; i < is_active.size() && i < addresses.size(); ++i)

{

const auto & address = addresses[i];

const bool is_local_replica = std::any_of(

shard.local_addresses.begin(),

shard.local_addresses.end(),

[&](const Cluster::Address & local_addr)

{ return local_addr.host_name == address.host_name && local_addr.port == address.port; });

if (is_local_replica)

{

local_replica_index = i;

break;

}

}

/// Test-only: simulate a transient window where the initiator's own `active` znode is momentarily

/// missing, so liveness reports the local replica as inactive. The forcing below must still keep it;

/// otherwise the local replica is filtered out and `findLocalReplicaIndexAndUpdatePools` throws

/// INCONSISTENT_CLUSTER_DEFINITION, turning a query that used to run into an error.

fiu_do_on(FailPoints::parallel_replicas_force_local_replica_inactive,

{

if (local_replica_index)

is_active[*local_replica_index] = false;

});

/// The local replica is the initiator - it is running this query, so it is online by definition even

/// if its `active` znode is transiently missing. Force it active so liveness never filters it out.

if (local_replica_index)

is_active[*local_replica_index] = true;

/// Test-only: simulate liveness drifting between the two passes of an INSERT SELECT - drop one active

/// non-local replica from this (first) snapshot so the coordinator is sized smaller than the cluster's

/// current active set. The remote-pool pass must reuse this snapshot's pools; if it recomputed liveness

/// instead, it would see the dropped replica active again and assign it a replica number that is out of

/// range for the already-sized coordinator. ONCE, so only the first (coordinator-building) call is hit.

/// Skip the discarded suitability-probe plan (see `in_insert_select_suitability_probe`), otherwise it

/// would consume the ONCE failpoint before the executed coordinator is built.

if (!in_insert_select_suitability_probe)

{

fiu_do_on(FailPoints::parallel_replicas_insert_select_drop_active_replica,

{

for (size_t i = 0; i < is_active.size(); ++i)

{

if (is_active[i] && (!local_replica_index || i != *local_replica_index))

{

is_active[i] = false;

break;

}

}

});

}

}

size_t available_replicas = shard.getAllNodeCount();

if (!is_active.empty())

{

available_replicas = std::count(is_active.begin(), is_active.end(), true);

/// Safety net: if liveness reports no active replicas (it should not, since this query is running),

/// ignore it rather than ending up with an empty replica set.

if (available_replicas == 0)

{

is_active.clear();

available_replicas = shard.getAllNodeCount();

}

}

size_t max_replicas_to_use = settings[Setting::max_parallel_replicas];

if (max_replicas_to_use > available_replicas)

{

LOG_TRACE(

logger,

"The number of replicas requested ({}) is bigger than the real number available in the cluster ({}). "

"Will use the latter number to execute the query.",

settings[Setting::max_parallel_replicas].value,

available_replicas);

max_replicas_to_use = available_replicas;

}

std::vector<ConnectionPoolWithFailover::Base::ShuffledPool> shuffled_pool;

if (max_replicas_to_use < available_replicas)

{

// will be shuffled according to `load_balancing` setting

shuffled_pool = shard.pool->getShuffledPools(settings);

}

else

{

/// If all (active) replicas in cluster are used for query execution,

/// try to preserve replicas order as in cluster definition.

/// It's important for data locality during query execution

/// independently of the query initiator

auto priority_func = [](size_t i) { return Priority{static_cast<Int64>(i)}; };

shuffled_pool = shard.pool->getShuffledPools(settings, priority_func);

}

std::vector<ConnectionPoolPtr> pools_to_use;

pools_to_use.reserve(shuffled_pool.size());

for (auto & pool : shuffled_pool)

{

/// Skip inactive replicas so they do not occupy a slot in the reading coordinator.

if (!is_active.empty() && !is_active[pool.index])

continue;

pools_to_use.emplace_back(std::move(pool.pool));

}

return {pools_to_use, max_replicas_to_use};

{

const auto & shard = cluster->getShardsInfo().at(0);

/// find local replica index in pool

std::optional<size_t> local_replica_index;

for (size_t i = 0, s = pools.size(); i < s; ++i)

{

const auto & hostname = pools[i]->getHost();

const auto & port = pools[i]->getPort();

const auto found = std::find_if(

begin(shard.local_addresses),

end(shard.local_addresses),

[&hostname, &port](const Cluster::Address & local_addr)

{

return hostname == local_addr.host_name && port == local_addr.port;

});

if (found != shard.local_addresses.end())

{

local_replica_index = i;

break;

}

}

if (!local_replica_index)

throw Exception(

ErrorCodes::INCONSISTENT_CLUSTER_DEFINITION,

"Local replica is not found in '{}' cluster definition, see 'cluster_for_parallel_replicas' setting",

cluster->getName());

// resize the pool but keep local replicas in it (and update its index)

chassert(max_replicas_to_use <= pools.size());

if (local_replica_index >= max_replicas_to_use)

{

std::swap(pools[max_replicas_to_use - 1], pools[local_replica_index.value()]);

local_replica_index = max_replicas_to_use - 1;

}

pools.resize(max_replicas_to_use);

return local_replica_index;

QueryPlan & query_plan,

const StorageID & storage_id,

SharedHeader header,

QueryProcessingStage::Enum processed_stage,

const ASTPtr & query_ast,

QueryTreeNodePtr query_tree,

PlannerContextPtr planner_context,

ContextPtr context,

std::shared_ptr<const StorageLimitsList> storage_limits,

QueryPlanStepPtr analyzed_read_from_merge_tree)

{

auto logger = getLogger("executeQueryWithParallelReplicas");

LOG_DEBUG(logger, "Executing read from {}, header {}, query ({}), stage {} with parallel replicas",

storage_id.getNameForLogs(), header->dumpStructure(), query_ast->formatForLogging(), processed_stage);

auto [cluster, shard_num] = prepareClusterForParallelReplicas(logger, context);

auto new_context = updateContextForParallelReplicas(logger, context, shard_num);

auto [connection_pools, max_replicas_to_use] = prepareConnectionPoolsForParallelReplicas(logger, new_context, cluster);

auto external_tables = new_context->getExternalTables();

auto coordinator = std::make_shared<ParallelReplicasReadingCoordinator>(max_replicas_to_use);

auto scalars = new_context->hasQueryContext() ? new_context->getQueryContext()->getScalars() : Scalars{};

const auto & shard = cluster->getShardsInfo().at(0);

/// do not build local plan for distributed queries for now (address it later);

/// when `parallel_replicas_prefer_local_replica` is false, skip local plan to allow the

/// load balancer to pick any replica.

if (canUseLocalPlanForParallelReplicas(new_context))

{

auto local_replica_index = findLocalReplicaIndexAndUpdatePools(connection_pools, max_replicas_to_use, cluster);

auto [local_plan, with_parallel_replicas] = createLocalPlanForParallelReplicas(

query_tree,

*header,

new_context,

processed_stage,

coordinator,

std::move(analyzed_read_from_merge_tree),

local_replica_index.value());

if (!with_parallel_replicas || connection_pools.size() == 1)

{

query_plan = std::move(*local_plan);

return;

}

std::shared_ptr<const QueryPlan> remote_query_plan;

if (new_context->getSettingsRef()[Setting::serialize_query_plan])

{

remote_query_plan = createRemotePlanForParallelReplicas(query_tree, *header, new_context, processed_stage);

remote_query_plan->ensureSerialized(DBMS_QUERY_PLAN_SERIALIZATION_VERSION);

}

/// The subquery carries its own SETTINGS (shipped to remote replicas via the AST). Pass its

/// context down so the local plan is optimized with the same read-in-order settings as the

/// replicas, and the initiator does not end up with a different coordination mode.

ContextPtr local_context = new_context;

if (const auto * query_node = query_tree->as<QueryNode>())

local_context = query_node->getContext();

else if (const auto * union_node = query_tree->as<UnionNode>())

local_context = union_node->getContext();

auto read_from_local = std::make_unique<ReadFromLocalParallelReplicaStep>(std::move(local_plan), std::move(local_context));

auto stub_local_plan = std::make_unique<QueryPlan>();

stub_local_plan->addStep(std::move(read_from_local));

LOG_DEBUG(logger, "Local replica got replica number {}", local_replica_index.value());

auto read_from_remote = std::make_unique<ReadFromParallelRemoteReplicasStep>(

query_ast,

query_tree,

planner_context,

cluster,

storage_id,

coordinator,

header,

processed_stage,

new_context,

getThrottler(new_context),

std::move(scalars),

std::move(external_tables),

getLogger("ReadFromParallelRemoteReplicasStep"),

std::move(storage_limits),

std::move(connection_pools),

local_replica_index,

shard.pool,

std::move(remote_query_plan));

auto remote_plan = std::make_unique<QueryPlan>();

remote_plan->addStep(std::move(read_from_remote));

SharedHeaders input_headers;

input_headers.reserve(2);

input_headers.emplace_back(stub_local_plan->getCurrentHeader());

input_headers.emplace_back(remote_plan->getCurrentHeader());

std::vector<QueryPlanPtr> plans;

plans.emplace_back(std::move(stub_local_plan));

plans.emplace_back(std::move(remote_plan));

auto union_step = std::make_unique<UnionStep>(std::move(input_headers));

query_plan.unitePlans(std::move(union_step), std::move(plans));

}

else

{

chassert(max_replicas_to_use <= connection_pools.size());

connection_pools.resize(max_replicas_to_use);

auto read_from_remote = std::make_unique<ReadFromParallelRemoteReplicasStep>(

query_ast,

query_tree,

planner_context,

cluster,

storage_id,

std::move(coordinator),

header,

processed_stage,

new_context,

getThrottler(new_context),

std::move(scalars),

std::move(external_tables),

getLogger("ReadFromParallelRemoteReplicasStep"),

std::move(storage_limits),

std::move(connection_pools),

std::nullopt,

shard.pool);

query_plan.addStep(std::move(read_from_remote));

}

QueryPlan & query_plan,

const StorageID & storage_id,

QueryProcessingStage::Enum processed_stage,

const QueryTreeNodePtr & query_tree,