#include <Processors/QueryPlan/ReadFromMergeTree.h>

#include <Analyzer/QueryNode.h>

#include <Core/Settings.h>

#include <Functions/IFunction.h>

#include <IO/Operators.h>

#include <Interpreters/Cluster.h>

#include <Interpreters/Context.h>

#include <Interpreters/ExpressionAnalyzer.h>

#include <Interpreters/ExpressionActions.h>

#include <Interpreters/InterpreterSelectQuery.h>

#include <Interpreters/TreeRewriter.h>

#include <Interpreters/Cache/QueryConditionCache.h>

#include <Interpreters/ClusterProxy/distributedIndexAnalysis.h>

#include <Parsers/ASTFunction.h>

#include <Parsers/ASTIdentifier.h>

#include <Parsers/ASTLiteral.h>

#include <Parsers/ASTSelectQuery.h>

#include <Parsers/parseIdentifierOrStringLiteral.h>

#include <Processors/ConcatProcessor.h>

#include <Processors/Merges/AggregatingSortedTransform.h>

#include <Processors/Merges/CoalescingSortedTransform.h>

#include <Processors/Merges/CollapsingSortedTransform.h>

#include <Processors/Merges/GraphiteRollupSortedTransform.h>

#include <Processors/Merges/MergingSortedTransform.h>

#include <Processors/Merges/ReplacingSortedTransform.h>

#include <Processors/Merges/SummingSortedTransform.h>

#include <Processors/Merges/VersionedCollapsingTransform.h>

#include <Processors/QueryPlan/IQueryPlanStep.h>

#include <Processors/QueryPlan/PartsSplitter.h>

#include <Processors/QueryPlan/LazilyReadFromMergeTree.h>

#include <Processors/QueryPlan/QueryIdHolder.h>

#include <Processors/Sources/NullSource.h>

#include <Processors/Transforms/ExpressionTransform.h>

#include <Processors/Transforms/FilterTransform.h>

#include <Processors/Transforms/ReverseTransform.h>

#include <Processors/Transforms/SelectByIndicesTransform.h>

#include <Processors/Transforms/VirtualRowTransform.h>

#include <QueryPipeline/QueryPipelineBuilder.h>

#include <Storages/MergeTree/MergeTreeDataSelectExecutor.h>

#include <Storages/MergeTree/MergeTreeIndexMinMax.h>

#include <Storages/MergeTree/MergeTreeIndexText.h>

#include <Storages/MergeTree/MergeTreeIndexVectorSimilarity.h>

#include <Storages/MergeTree/MergeTreePrefetchedReadPool.h>

#include <Storages/MergeTree/MergeTreeReadPool.h>

#include <Storages/MergeTree/MergeTreeReadPoolInOrder.h>

#include <Storages/MergeTree/MergeTreeReadPoolParallelReplicas.h>

#include <Storages/MergeTree/MergeTreeReadPoolParallelReplicasInOrder.h>

#include <Storages/MergeTree/MergeTreeIndexReadResultPool.h>

#include <Storages/MergeTree/MergeTreeReadPoolProjectionIndex.h>

#include <Storages/MergeTree/MergeTreeSettings.h>

#include <Storages/MergeTree/MergeTreeSource.h>

#include <Storages/MergeTree/RangesInDataPart.h>

#include <Storages/MergeTree/RequestResponse.h>

#include <Storages/Statistics/ConditionSelectivityEstimator.h>

#include <Storages/VirtualColumnUtils.h>

#include <Common/JSONBuilder.h>

#include <Common/logger_useful.h>

#include <Common/thread_local_rng.h>

#include <algorithm>

#include <iterator>

#include <memory>

#include <string_view>

#include <unordered_map>

#include <unordered_set>

#include <fmt/ranges.h>

#include "config.h"

using namespace DB;

namespace

{

template <typename Container, typename Getter>

size_t countPartitions(const Container & parts, Getter get_partition_id)

{

if (parts.empty())

return 0;

String cur_partition_id = get_partition_id(parts[0]);

size_t unique_partitions = 1;

for (size_t i = 1; i < parts.size(); ++i)

{

if (get_partition_id(parts[i]) != cur_partition_id)

{

++unique_partitions;

cur_partition_id = get_partition_id(parts[i]);

}

}

return unique_partitions;

}

size_t countPartitions(const RangesInDataParts & parts_with_ranges)

{

auto get_partition_id = [](const RangesInDataPart & rng) { return rng.data_part->info.getPartitionId(); };

return countPartitions(parts_with_ranges, get_partition_id);

}

/// check if a DAG node only depends on sorting key columns (ActionsDAG version of isExpressionOverSortingKey)

bool isNodeOverSortingKey(const ActionsDAG::Node * node, const NameSet & sorting_key_set)

{

if (sorting_key_set.contains(node->result_name))

return true;

if (node->type == ActionsDAG::ActionType::COLUMN)

return true; // constants are fine

if (node->type == ActionsDAG::ActionType::INPUT || node->type == ActionsDAG::ActionType::PLACEHOLDER)

return false; // already checked result_name

for (const auto * child : node->children)

if (!isNodeOverSortingKey(child, sorting_key_set))

return false;

return true;

}

bool restoreDAGInputs(ActionsDAG & dag, const NameSet & inputs)

{

std::unordered_set<const ActionsDAG::Node *> outputs(dag.getOutputs().begin(), dag.getOutputs().end());

bool added = false;

for (const auto * input : dag.getInputs())

{

if (inputs.contains(input->result_name) && !outputs.contains(input))

{

dag.getOutputs().push_back(input);

added = true;

}

}

return added;

}

bool restorePrewhereInputs(FilterDAGInfo * row_level_filter, PrewhereInfo * info, const NameSet & inputs)

{

bool added = false;

if (row_level_filter)

added = added || restoreDAGInputs(row_level_filter->actions, inputs);

if (info)

added = added || restoreDAGInputs(info->prewhere_actions, inputs);

return added;

}

}

namespace ProfileEvents

{

extern const Event IndexAnalysisRounds;

extern const Event SelectedParts;

extern const Event SelectedPartsTotal;

extern const Event SelectedRanges;

extern const Event SelectedMarks;

extern const Event SelectedMarksTotal;

extern const Event SelectQueriesWithPrimaryKeyUsage;

}

namespace DB

{

namespace Setting

{

extern const SettingsBool allow_asynchronous_read_from_io_pool_for_merge_tree;

extern const SettingsBool allow_prefetched_read_pool_for_local_filesystem;

extern const SettingsBool allow_prefetched_read_pool_for_remote_filesystem;

extern const SettingsBool compile_sort_description;

extern const SettingsBool do_not_merge_across_partitions_select_final;

extern const SettingsBool enable_automatic_decision_for_merging_across_partitions_for_final;

extern const SettingsBool enable_vertical_final;

extern const SettingsBool force_aggregate_partitions_independently;

extern const SettingsBool force_primary_key;

extern const SettingsString ignore_data_skipping_indices;

extern const SettingsUInt64 max_number_of_partitions_for_independent_aggregation;

extern const SettingsInt64 max_partitions_to_read;

extern const SettingsUInt64 max_rows_to_read;

extern const SettingsUInt64 max_rows_to_read_leaf;

extern const SettingsMaxThreads max_final_threads;

extern const SettingsUInt64 max_parser_backtracks;

extern const SettingsUInt64 max_parser_depth;

extern const SettingsUInt64 max_query_size;

extern const SettingsUInt64 max_streams_for_merge_tree_reading;

extern const SettingsMaxThreads max_threads;

extern const SettingsUInt64 merge_tree_max_bytes_to_use_cache;

extern const SettingsUInt64 merge_tree_max_rows_to_use_cache;

extern const SettingsUInt64 merge_tree_min_bytes_for_concurrent_read_for_remote_filesystem;

extern const SettingsUInt64 merge_tree_min_rows_for_concurrent_read_for_remote_filesystem;

extern const SettingsUInt64 merge_tree_min_bytes_for_concurrent_read;

extern const SettingsUInt64 merge_tree_min_rows_for_concurrent_read;

extern const SettingsFloat merge_tree_read_split_ranges_into_intersecting_and_non_intersecting_injection_probability;

extern const SettingsBool merge_tree_use_const_size_tasks_for_remote_reading;

extern const SettingsUInt64 min_count_to_compile_sort_description;

extern const SettingsOverflowMode read_overflow_mode;

extern const SettingsOverflowMode read_overflow_mode_leaf;

extern const SettingsUInt64 parallel_replicas_count;

extern const SettingsBool parallel_replicas_local_plan;

extern const SettingsBool parallel_replicas_index_analysis_only_on_coordinator;

extern const SettingsBool parallel_replicas_support_projection;

extern const SettingsBool distributed_index_analysis;

extern const SettingsBool distributed_index_analysis_for_non_shared_merge_tree;

extern const SettingsUInt64 preferred_block_size_bytes;

extern const SettingsUInt64 preferred_max_column_in_block_size_bytes;

extern const SettingsUInt64 read_in_order_two_level_merge_threshold;

extern const SettingsBool split_parts_ranges_into_intersecting_and_non_intersecting_final;

extern const SettingsBool split_intersecting_parts_ranges_into_layers_final;

extern const SettingsBool use_primary_key;

extern const SettingsBool use_partition_pruning;

extern const SettingsBool use_skip_indexes;

extern const SettingsBool use_skip_indexes_if_final;

extern const SettingsBool use_skip_indexes_for_disjunctions;

extern const SettingsBool use_uncompressed_cache;

extern const SettingsNonZeroUInt64 merge_tree_min_read_task_size;

extern const SettingsBool read_in_order_use_virtual_row;

extern const SettingsBool use_skip_indexes_if_final_exact_mode;

extern const SettingsBool use_skip_indexes_on_data_read;

extern const SettingsBool use_skip_indexes_for_top_k;

extern const SettingsBool use_top_k_dynamic_filtering;

extern const SettingsBool use_query_condition_cache;

extern const SettingsNonZeroUInt64 max_parallel_replicas;

extern const SettingsBool enable_shared_storage_snapshot_in_query;

extern const SettingsUInt64 query_plan_max_step_description_length;

extern const SettingsBool apply_row_policy_after_final;

extern const SettingsBool apply_prewhere_after_final;

extern const SettingsBool distributed_index_analysis_only_on_coordinator;

}

namespace MergeTreeSetting

{

extern const MergeTreeSettingsUInt64 index_granularity;

extern const MergeTreeSettingsUInt64 index_granularity_bytes;

extern const MergeTreeSettingsUInt64 max_concurrent_queries;

extern const MergeTreeSettingsInt64 max_partitions_to_read;

extern const MergeTreeSettingsUInt64 min_marks_to_honor_max_concurrent_queries;

extern const MergeTreeSettingsUInt64 distributed_index_analysis_min_parts_to_activate;

extern const MergeTreeSettingsUInt64 distributed_index_analysis_min_indexes_bytes_to_activate;

}

namespace ErrorCodes

{

extern const int INDEX_NOT_USED;

extern const int LOGICAL_ERROR;

extern const int TOO_MANY_PARTITIONS;

}

static bool checkAllPartsOnRemoteFS(const RangesInDataParts & parts)

{

for (const auto & part : parts)

{

if (!part.data_part->isStoredOnRemoteDisk())

return false;

}

return true;

}

/// build sort description for output stream

static SortDescription getSortDescriptionForOutputHeader(

const SharedHeader & output_header,

const Names & sorting_key_columns,

const std::vector<bool> & reverse_flags,

const int sort_direction,

InputOrderInfoPtr input_order_info,

const FilterDAGInfoPtr & row_level_filter,

const PrewhereInfoPtr & prewhere_info,

bool enable_vertical_final)

{

/// Updating sort description can be done after PREWHERE actions are applied to the header.

/// After PREWHERE actions are applied, column names in header can differ from storage column names due to aliases

/// To mitigate it, we're trying to build original header and use it to deduce sorting description

/// TODO: this approach is fragile, it'd be more robust to update sorting description for the whole plan during plan optimization

Block original_header = output_header->cloneEmpty();

if (prewhere_info)

{

{

FindOriginalNodeForOutputName original_column_finder(prewhere_info->prewhere_actions);

for (auto & column : original_header)

{

const auto * original_node = original_column_finder.find(column.name);

if (original_node)

column.name = original_node->result_name;

}

}

}

if (row_level_filter)

{

FindOriginalNodeForOutputName original_column_finder(row_level_filter->actions);

for (auto & column : original_header)

{

const auto * original_node = original_column_finder.find(column.name);

if (original_node)

column.name = original_node->result_name;

}

}

SortDescription sort_description;

const Block & header = *output_header;

size_t sort_columns_size = sorting_key_columns.size();

sort_description.reserve(sort_columns_size);

for (size_t i = 0; i < sort_columns_size; ++i)

{

const auto & sorting_key = sorting_key_columns[i];

const auto it = std::find_if(

original_header.begin(), original_header.end(), [&sorting_key](const auto & column) { return column.name == sorting_key; });

if (it == original_header.end())

break;

const size_t column_pos = std::distance(original_header.begin(), it);

if (!reverse_flags.empty() && reverse_flags[i])

sort_description.emplace_back((header.begin() + column_pos)->name, sort_direction * -1);

else

sort_description.emplace_back((header.begin() + column_pos)->name, sort_direction);

}

if (input_order_info && !enable_vertical_final)

{

// output_stream.sort_scope = DataStream::SortScope::Stream;

const size_t used_prefix_of_sorting_key_size = input_order_info->used_prefix_of_sorting_key_size;

if (sort_description.size() > used_prefix_of_sorting_key_size)

sort_description.resize(used_prefix_of_sorting_key_size);

return sort_description;

}

return {};

}

std::shared_ptr<QueryIdHolder> ReadFromMergeTree::AnalysisResult::checkLimits(

const Context & context_, const MergeTreeData & data_, const MergeTreeSettings & data_settings_) const

{

const Settings & settings = context_.getSettingsRef();

auto max_partitions_to_read = settings[Setting::max_partitions_to_read].changed

? settings[Setting::max_partitions_to_read].value

: data_settings_[MergeTreeSetting::max_partitions_to_read].value;

if (max_partitions_to_read > 0)

{

std::set<String> partitions;

for (const auto & part_with_ranges : parts_with_ranges)

partitions.insert(part_with_ranges.data_part->info.getPartitionId());

if (partitions.size() > static_cast<size_t>(max_partitions_to_read))

{

throw Exception(

ErrorCodes::TOO_MANY_PARTITIONS,

"Too many partitions to read. Current {}, max {}",

partitions.size(),

max_partitions_to_read);

}

}

if (data_settings_[MergeTreeSetting::max_concurrent_queries] > 0

&& data_settings_[MergeTreeSetting::min_marks_to_honor_max_concurrent_queries] > 0

&& selected_marks >= data_settings_[MergeTreeSetting::min_marks_to_honor_max_concurrent_queries])

{

auto query_id = context_.getCurrentQueryId();

if (!query_id.empty())

return data_.getQueryIdHolder(query_id, data_settings_[MergeTreeSetting::max_concurrent_queries]);

}

return nullptr;

}

ReadFromMergeTree::ReadFromMergeTree(

RangesInDataPartsPtr parts_,

MergeTreeData::MutationsSnapshotPtr mutations_,

Names all_column_names_,

const MergeTreeData & data_,

MergeTreeSettingsPtr data_settings_,

const SelectQueryInfo & query_info_,

const StorageSnapshotPtr & storage_snapshot_,

const ContextPtr & context_,

size_t max_block_size_,

size_t num_streams_,

PartitionIdToMaxBlockPtr max_block_numbers_to_read_,

LoggerPtr log_,

AnalysisResultPtr analyzed_result_ptr_,

bool enable_parallel_reading_,

std::optional<MergeTreeAllRangesCallback> all_ranges_callback_,

std::optional<MergeTreeReadTaskCallback> read_task_callback_,

std::optional<size_t> number_of_current_replica_)

: SourceStepWithFilter(std::make_shared<const Block>(MergeTreeSelectProcessor::transformHeader(

storage_snapshot_->getSampleBlockForColumns(all_column_names_),

query_info_.row_level_filter,

query_info_.prewhere_info)), all_column_names_, query_info_, storage_snapshot_, context_)

, data_settings(std::move(data_settings_))

, reader_settings(MergeTreeReaderSettings::createForQuery(context_, *data_settings, query_info_))

, prepared_parts(std::move(parts_))

, mutations_snapshot(std::move(mutations_))

, all_column_names(std::move(all_column_names_))

, data(data_)

, actions_settings(ExpressionActionsSettings(context_))

, block_size{

.max_block_size_rows = max_block_size_,

.preferred_block_size_bytes = context->getSettingsRef()[Setting::preferred_block_size_bytes],

.preferred_max_column_in_block_size_bytes = context->getSettingsRef()[Setting::preferred_max_column_in_block_size_bytes]}

, requested_num_streams(num_streams_)

, max_block_numbers_to_read(std::move(max_block_numbers_to_read_))

, log(std::move(log_))

, analyzed_result_ptr(analyzed_result_ptr_)

, is_parallel_reading_from_replicas(enable_parallel_reading_)

, enable_remove_parts_from_snapshot_optimization(!context->getSettingsRef()[Setting::enable_shared_storage_snapshot_in_query] && query_info_.merge_tree_enable_remove_parts_from_snapshot_optimization)

, number_of_current_replica(number_of_current_replica_)

{

if (is_parallel_reading_from_replicas)

{

if (all_ranges_callback_.has_value())

all_ranges_callback = all_ranges_callback_.value();

else

all_ranges_callback = context->getMergeTreeAllRangesCallback();

if (read_task_callback_.has_value())

read_task_callback = read_task_callback_.value();

else

read_task_callback = context->getMergeTreeReadTaskCallback();

}

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

if (settings[Setting::max_streams_for_merge_tree_reading])

{

if (settings[Setting::allow_asynchronous_read_from_io_pool_for_merge_tree])

{

/// When async reading is enabled, allow to read using more streams.

/// Will add resize to output_streams_limit to reduce memory usage.

output_streams_limit = std::min<size_t>(requested_num_streams, settings[Setting::max_streams_for_merge_tree_reading]);

/// We intentionally set `max_streams` to 1 in InterpreterSelectQuery in case of small limit.

/// Changing it here to `max_streams_for_merge_tree_reading` proven itself as a threat for performance.

if (requested_num_streams != 1)

requested_num_streams = std::max<size_t>(requested_num_streams, settings[Setting::max_streams_for_merge_tree_reading]);

}

else

/// Just limit requested_num_streams otherwise.

requested_num_streams = std::min<size_t>(requested_num_streams, settings[Setting::max_streams_for_merge_tree_reading]);

}

/// Add explicit description.

std::string description = data.getStorageID().getFullNameNotQuoted();

setStepDescription(description, context->getSettingsRef()[Setting::query_plan_max_step_description_length]);

enable_vertical_final = query_info.isFinal() && context->getSettingsRef()[Setting::enable_vertical_final]

&& data.merging_params.mode == MergeTreeData::MergingParams::Replacing;

}

std::unique_ptr<ReadFromMergeTree> ReadFromMergeTree::createLocalParallelReplicasReadingStep(

ContextPtr & context_,

AnalysisResultPtr analyzed_result_ptr_,

MergeTreeAllRangesCallback all_ranges_callback_,

MergeTreeReadTaskCallback read_task_callback_,

size_t replica_number)

{

const bool enable_parallel_reading = true;

return std::make_unique<ReadFromMergeTree>(

/// Optimized version of getParts() to avoid extra copy

analyzed_result_ptr ? std::make_shared<RangesInDataParts>(analyzed_result_ptr->parts_with_ranges) : prepared_parts,

mutations_snapshot,

all_column_names,

data,

data_settings,

getQueryInfo(),

getStorageSnapshot(),

context_,

block_size.max_block_size_rows,

requested_num_streams,

max_block_numbers_to_read,

log,

std::move(analyzed_result_ptr_),

enable_parallel_reading,

all_ranges_callback_,

read_task_callback_,

replica_number);

}

Pipe ReadFromMergeTree::readFromPoolParallelReplicas(

RangesInDataParts parts_with_range,

const MergeTreeIndexBuildContextPtr & index_build_context,

Names required_columns,

PoolSettings pool_settings)

{

const auto & client_info = context->getClientInfo();

auto extension = ParallelReadingExtension{

all_ranges_callback.value(),

read_task_callback.value(),

number_of_current_replica.value_or(client_info.number_of_current_replica),

context->getClusterForParallelReplicas()->getShardsInfo().at(0).getAllNodeCount()};

auto pool = std::make_shared<MergeTreeReadPoolParallelReplicas>(

std::move(extension),

std::move(parts_with_range),

mutations_snapshot,

shared_virtual_fields,

index_read_tasks,

storage_snapshot,

query_info.row_level_filter,

query_info.prewhere_info,

actions_settings,

reader_settings,

required_columns,

pool_settings,

block_size,

context);

Pipes pipes;

for (size_t i = 0; i < pool_settings.threads; ++i)

{

auto algorithm = std::make_unique<MergeTreeThreadSelectAlgorithm>(i);

auto processor = std::make_unique<MergeTreeSelectProcessor>(

pool,

std::move(algorithm),

query_info.row_level_filter,

query_info.prewhere_info,

index_read_tasks,

actions_settings,

reader_settings,

index_build_context);

auto source = std::make_shared<MergeTreeSource>(std::move(processor), data.getLogName());

pipes.emplace_back(std::move(source));

}

return Pipe::unitePipes(std::move(pipes));

}

Pipe ReadFromMergeTree::readFromPool(

RangesInDataParts parts_with_range,

const MergeTreeIndexBuildContextPtr & index_build_context,

Names required_columns,

PoolSettings pool_settings)

{

size_t total_rows = parts_with_range.getRowsCountAllParts();

if (query_info.trivial_limit > 0 && query_info.trivial_limit < total_rows)

total_rows = query_info.trivial_limit;

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

/// round min_marks_to_read up to nearest multiple of block_size expressed in marks

/// If granularity is adaptive it doesn't make sense

/// Maybe it will make sense to add settings `max_block_size_bytes`

if (block_size.max_block_size_rows && !data.canUseAdaptiveGranularity())

{

size_t fixed_index_granularity = (*data_settings)[MergeTreeSetting::index_granularity];

pool_settings.min_marks_for_concurrent_read

= (pool_settings.min_marks_for_concurrent_read * fixed_index_granularity + block_size.max_block_size_rows - 1)

/ block_size.max_block_size_rows * block_size.max_block_size_rows / fixed_index_granularity;

}

bool all_parts_are_remote = true;

bool all_parts_are_local = true;

for (const auto & part : parts_with_range)

{

const bool is_remote = part.data_part->isStoredOnRemoteDisk();

all_parts_are_local &= !is_remote;

all_parts_are_remote &= is_remote;

}

MergeTreeReadPoolPtr pool;

bool allow_prefetched_remote = all_parts_are_remote && settings[Setting::allow_prefetched_read_pool_for_remote_filesystem]

&& MergeTreePrefetchedReadPool::checkReadMethodAllowed(reader_settings.read_settings.remote_fs_method);

bool allow_prefetched_local = all_parts_are_local && settings[Setting::allow_prefetched_read_pool_for_local_filesystem]

&& MergeTreePrefetchedReadPool::checkReadMethodAllowed(reader_settings.read_settings.local_fs_method);

/** Do not use prefetched read pool if query is trivial limit query.

* Because time spend during filling per thread tasks can be greater than whole query

* execution for big tables with small limit.

*/

bool use_prefetched_read_pool = query_info.trivial_limit == 0 && (allow_prefetched_remote || allow_prefetched_local);

if (use_prefetched_read_pool)

{

pool = std::make_shared<MergeTreePrefetchedReadPool>(

std::move(parts_with_range),

mutations_snapshot,

shared_virtual_fields,

index_read_tasks,

storage_snapshot,

query_info.row_level_filter,

query_info.prewhere_info,

actions_settings,

reader_settings,

required_columns,

pool_settings,

block_size,

context,

dataflow_cache_updater);

}

else

{

pool = std::make_shared<MergeTreeReadPool>(

std::move(parts_with_range),

mutations_snapshot,

shared_virtual_fields,

index_read_tasks,

storage_snapshot,

query_info.row_level_filter,

query_info.prewhere_info,

actions_settings,

reader_settings,

required_columns,

pool_settings,

block_size,

context,

dataflow_cache_updater);

}

LOG_DEBUG(log, "Reading approx. {} rows with {} streams", total_rows, pool_settings.threads);

Pipes pipes;

for (size_t i = 0; i < pool_settings.threads; ++i)

{

auto algorithm = std::make_unique<MergeTreeThreadSelectAlgorithm>(i);

auto processor = std::make_unique<MergeTreeSelectProcessor>(

pool,

std::move(algorithm),

query_info.row_level_filter,

query_info.prewhere_info,

index_read_tasks,

actions_settings,

reader_settings,

index_build_context);

auto source = std::make_shared<MergeTreeSource>(std::move(processor), data.getLogName());

if (i == 0)

source->addTotalRowsApprox(total_rows);

pipes.emplace_back(std::move(source));

}

auto pipe = Pipe::unitePipes(std::move(pipes));

if (output_streams_limit && output_streams_limit < pipe.numOutputPorts())

pipe.resize(output_streams_limit);

return pipe;

}

Pipe ReadFromMergeTree::readInOrder(

RangesInDataParts parts_with_ranges,

const MergeTreeIndexBuildContextPtr & index_build_context,

Names required_columns,

PoolSettings pool_settings,

ReadType read_type,

UInt64 read_limit)

{

/// For reading in order it makes sense to read only

/// one range per task to reduce number of read rows.

bool has_limit_below_one_block = read_type != ReadType::Default && read_limit && read_limit < block_size.max_block_size_rows;

MergeTreeReadPoolPtr pool;

if (is_parallel_reading_from_replicas)

{

const auto & client_info = context->getClientInfo();

ParallelReadingExtension extension{

all_ranges_callback.value(),

read_task_callback.value(),

number_of_current_replica.value_or(client_info.number_of_current_replica),

context->getClusterForParallelReplicas()->getShardsInfo().at(0).getAllNodeCount()};

CoordinationMode mode = read_type == ReadType::InOrder

? CoordinationMode::WithOrder

: CoordinationMode::ReverseOrder;

pool = std::make_shared<MergeTreeReadPoolParallelReplicasInOrder>(

std::move(extension),

mode,

parts_with_ranges,

mutations_snapshot,

shared_virtual_fields,

index_read_tasks,

has_limit_below_one_block,

storage_snapshot,

query_info.row_level_filter,

query_info.prewhere_info,

actions_settings,

reader_settings,

required_columns,

pool_settings,

block_size,

context);

}

else

{

pool = std::make_shared<MergeTreeReadPoolInOrder>(

has_limit_below_one_block,

read_type,

parts_with_ranges,

mutations_snapshot,

shared_virtual_fields,

index_read_tasks,

storage_snapshot,

query_info.row_level_filter,

query_info.prewhere_info,

actions_settings,

reader_settings,

required_columns,

pool_settings,

block_size,

context,

dataflow_cache_updater);

}

/// If parallel replicas enabled, set total rows in progress here only on initiator with local plan

/// Otherwise rows will counted multiple times

const UInt64 in_order_limit = query_info.input_order_info ? query_info.input_order_info->limit : 0;

const bool set_total_rows_approx = !is_parallel_reading_from_replicas || isParallelReplicasLocalPlanForInitiator();

Pipes pipes;

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

{

const auto & part_with_ranges = parts_with_ranges[i];

UInt64 total_rows = part_with_ranges.getRowsCount();

if (query_info.trivial_limit > 0 && query_info.trivial_limit < total_rows)

total_rows = query_info.trivial_limit;

else if (in_order_limit > 0 && in_order_limit < total_rows)

total_rows = in_order_limit;

LOG_TRACE(log, "Reading {} ranges in{}order from part {}, approx. {} rows starting from {}",

part_with_ranges.ranges.size(),

read_type == ReadType::InReverseOrder ? " reverse " : " ",

part_with_ranges.data_part->name, total_rows,

part_with_ranges.data_part->index_granularity->getMarkStartingRow(part_with_ranges.ranges.front().begin));

MergeTreeSelectAlgorithmPtr algorithm;

if (read_type == ReadType::InReverseOrder)

algorithm = std::make_unique<MergeTreeInReverseOrderSelectAlgorithm>(i);

else

algorithm = std::make_unique<MergeTreeInOrderSelectAlgorithm>(i);

auto processor = std::make_unique<MergeTreeSelectProcessor>(

pool,

std::move(algorithm),

query_info.row_level_filter,

query_info.prewhere_info,

index_read_tasks,

actions_settings,

reader_settings,

index_build_context);

processor->addPartLevelToChunk(isQueryWithFinal());

auto source = std::make_shared<MergeTreeSource>(std::move(processor), data.getLogName());

if (set_total_rows_approx)

source->addTotalRowsApprox(total_rows);

Pipe pipe(source);

if (virtual_row_conversion && (read_type == ReadType::InOrder || read_type == ReadType::InReverseOrder))

{

const auto & index = part_with_ranges.data_part->getIndex();

const auto & primary_key = storage_snapshot->metadata->primary_key;

bool has_final_mark = part_with_ranges.data_part->index_granularity->hasFinalMark();

bool read_in_direct_order = read_type == ReadType::InOrder;

size_t mark_range_pos = read_in_direct_order ? part_with_ranges.ranges.front().begin : part_with_ranges.ranges.back().end;

bool has_pk_value = (read_in_direct_order || has_final_mark) && std::ranges::all_of(*index, [&](const auto & col) { return col->size() > mark_range_pos; });

/// The index may have fewer columns than the primary key if suffix columns were

/// removed by optimizeIndexColumns (controlled by primary_key_ratio_of_unique_prefix_values_to_skip_suffix_columns).

/// In that case, we cannot apply virtual row optimization because we don't have all required columns.

size_t num_pk_columns_required = virtual_row_conversion->getRequiredColumnsWithTypes().size();

if (index->size() >= num_pk_columns_required && has_pk_value)

{

ColumnsWithTypeAndName pk_columns;

pk_columns.reserve(num_pk_columns_required);

for (size_t j = 0; j < num_pk_columns_required; ++j)

{

auto column = primary_key.data_types[j]->createColumn()->cloneEmpty();

column->insert((*(*index)[j])[mark_range_pos]);

pk_columns.push_back({std::move(column), primary_key.data_types[j], primary_key.column_names[j]});

}

Block pk_block(std::move(pk_columns));

pipe.addSimpleTransform([&](const SharedHeader & header)

{

return std::make_shared<VirtualRowTransform>(header, pk_block, virtual_row_conversion);

});

}

}

pipes.emplace_back(std::move(pipe));

}

auto pipe = Pipe::unitePipes(std::move(pipes));

if (read_type == ReadType::InReverseOrder)

{

pipe.addSimpleTransform([&](const SharedHeader & header)

{

return std::make_shared<ReverseTransform>(header);

});

}

return pipe;

}

Pipe ReadFromMergeTree::read(

RangesInDataParts parts_with_range,

const MergeTreeIndexBuildContextPtr & index_build_context,

Names required_columns,

ReadType read_type,

size_t max_streams,

size_t min_marks_for_concurrent_read,

bool use_uncompressed_cache)

{

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

size_t sum_marks = parts_with_range.getMarksCountAllParts();

const size_t total_query_nodes = is_parallel_reading_from_replicas

? std::min<size_t>(

context->getClusterForParallelReplicas()->getShardsInfo().at(0).getAllNodeCount(),

context->getSettingsRef()[Setting::max_parallel_replicas])

: 1;

PoolSettings pool_settings{

.threads = max_streams,

.sum_marks = sum_marks,

.min_marks_for_concurrent_read = min_marks_for_concurrent_read,

.preferred_block_size_bytes = settings[Setting::preferred_block_size_bytes],

.use_uncompressed_cache = use_uncompressed_cache,

.use_const_size_tasks_for_remote_reading = settings[Setting::merge_tree_use_const_size_tasks_for_remote_reading],

.total_query_nodes = total_query_nodes,

};

if (read_type == ReadType::ParallelReplicas)

return readFromPoolParallelReplicas(

std::move(parts_with_range), index_build_context, std::move(required_columns), std::move(pool_settings));

/// Reading from default thread pool is beneficial for remote storage because of new prefetches.

if (read_type == ReadType::Default && (max_streams > 1 || checkAllPartsOnRemoteFS(parts_with_range)))

return readFromPool(

std::move(parts_with_range), index_build_context, std::move(required_columns), std::move(pool_settings));

auto pipe = readInOrder(parts_with_range, index_build_context, required_columns, pool_settings, read_type, /*limit=*/0);

/// Use ConcatProcessor to concat sources together.

/// It is needed to read in parts order (and so in PK order) if single thread is used.

if (read_type == ReadType::Default && pipe.numOutputPorts() > 1)

pipe.addTransform(std::make_shared<ConcatProcessor>(pipe.getSharedHeader(), pipe.numOutputPorts()));

return pipe;

}

namespace

{

struct PartRangesReadInfo

{

std::vector<size_t> sum_marks_in_parts;

size_t sum_marks = 0;

size_t total_rows = 0;

size_t adaptive_parts = 0;

size_t index_granularity_bytes = 0;

size_t max_marks_to_use_cache = 0;

size_t min_marks_for_concurrent_read = 0;

bool use_uncompressed_cache = false;

PartRangesReadInfo(

const RangesInDataParts & parts,

const Settings & settings,

const MergeTreeSettings & data_settings)

{

/// Count marks for each part.

sum_marks_in_parts.resize(parts.size());

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

{

total_rows += parts[i].getRowsCount();

sum_marks_in_parts[i] = parts[i].getMarksCount();

sum_marks += sum_marks_in_parts[i];

if (parts[i].data_part->index_granularity_info.mark_type.adaptive)

++adaptive_parts;

}

if (adaptive_parts > parts.size() / 2)

index_granularity_bytes = data_settings[MergeTreeSetting::index_granularity_bytes];

max_marks_to_use_cache = MergeTreeDataSelectExecutor::roundRowsOrBytesToMarks(

settings[Setting::merge_tree_max_rows_to_use_cache],

settings[Setting::merge_tree_max_bytes_to_use_cache],

data_settings[MergeTreeSetting::index_granularity],

index_granularity_bytes);

auto all_parts_on_remote_disk = checkAllPartsOnRemoteFS(parts);

size_t min_rows_for_concurrent_read;

size_t min_bytes_for_concurrent_read;

if (all_parts_on_remote_disk)

{

min_rows_for_concurrent_read = settings[Setting::merge_tree_min_rows_for_concurrent_read_for_remote_filesystem];

min_bytes_for_concurrent_read = settings[Setting::merge_tree_min_bytes_for_concurrent_read_for_remote_filesystem];

}

else

{

min_rows_for_concurrent_read = settings[Setting::merge_tree_min_rows_for_concurrent_read];

min_bytes_for_concurrent_read = settings[Setting::merge_tree_min_bytes_for_concurrent_read];

}

min_marks_for_concurrent_read = MergeTreeDataSelectExecutor::minMarksForConcurrentRead(

min_rows_for_concurrent_read, min_bytes_for_concurrent_read,

data_settings[MergeTreeSetting::index_granularity], index_granularity_bytes, settings[Setting::merge_tree_min_read_task_size], sum_marks);

use_uncompressed_cache = settings[Setting::use_uncompressed_cache];

if (sum_marks > max_marks_to_use_cache)

use_uncompressed_cache = false;

}

};

}

Pipe ReadFromMergeTree::readByLayers(

const RangesInDataParts & parts_with_ranges,

SplitPartsByRanges split_parts,

const MergeTreeIndexBuildContextPtr & index_build_context,

const Names & column_names,

const InputOrderInfoPtr & input_order_info)

{

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

LOG_TRACE(log, "Spreading mark ranges among streams (reading by layers)");

PartRangesReadInfo info(parts_with_ranges, settings, *data_settings);

if (0 == info.sum_marks)

return {};

ReadingInOrderStepGetter reading_step_getter;

Names in_order_column_names_to_read;

SortDescription sort_description;

if (reader_settings.read_in_order)

{

NameSet column_names_set(column_names.begin(), column_names.end());

in_order_column_names_to_read = column_names;

/// Add columns needed to calculate the sorting expression

for (const auto & column_name : storage_snapshot->metadata->getColumnsRequiredForSortingKey())

{

if (column_names_set.contains(column_name))

continue;

in_order_column_names_to_read.push_back(column_name);

column_names_set.insert(column_name);

}

auto sorting_expr = storage_snapshot->metadata->getSortingKey().expression;

const auto & sorting_columns = storage_snapshot->metadata->getSortingKey().column_names;

std::vector<bool> reverse_flags = storage_snapshot->metadata->getSortingKeyReverseFlags();

sort_description.compile_sort_description = settings[Setting::compile_sort_description];

sort_description.min_count_to_compile_sort_description = settings[Setting::min_count_to_compile_sort_description];

sort_description.reserve(input_order_info->used_prefix_of_sorting_key_size);

for (size_t i = 0; i < input_order_info->used_prefix_of_sorting_key_size; ++i)

{

if (!reverse_flags.empty() && reverse_flags[i])

sort_description.emplace_back(sorting_columns[i], input_order_info->direction * -1);

else

sort_description.emplace_back(sorting_columns[i], input_order_info->direction);

}

reading_step_getter

= [this, &index_build_context, &in_order_column_names_to_read, &info, sorting_expr, &sort_description](auto parts)

{

auto pipe = this->read(

std::move(parts),

index_build_context,

in_order_column_names_to_read,

ReadType::InOrder,

1 /* num_streams */,

0 /* min_marks_for_concurrent_read */,

info.use_uncompressed_cache);

if (pipe.empty())

{

auto header = std::make_shared<const Block>(MergeTreeSelectProcessor::transformHeader(

storage_snapshot->getSampleBlockForColumns(in_order_column_names_to_read),

query_info.row_level_filter,

query_info.prewhere_info));

pipe = Pipe(std::make_shared<NullSource>(header));

}

pipe.addSimpleTransform([sorting_expr](const SharedHeader & header)

{

return std::make_shared<ExpressionTransform>(header, sorting_expr);

});

if (pipe.numOutputPorts() != 1)

{

auto transform = std::make_shared<MergingSortedTransform>(

pipe.getSharedHeader(),

pipe.numOutputPorts(),

sort_description,

block_size.max_block_size_rows,

/*max_block_size_bytes=*/ 0,

/*max_dynamic_subcolumns*/ std::nullopt,

SortingQueueStrategy::Batch,

/*limit=*/ 0,

/*always_read_till_end=*/ false,