Merge pull request ClickHouse#79417 from amosbird/projection-index-2

Add `_part_starting_offset` virtual column and key condition support for offset-based querying

Author: KochetovNicolai

docs/en/engines/table-engines/mergetree-family/mergetree.md

@@ -966,8 +966,11 @@ ClickHouse versions 22.3 through 22.7 use a different cache configuration, see [
 
 - `_part` — Name of a part.
 - `_part_index` — Sequential index of the part in the query result.
+- `_part_starting_offset` — Cumulative starting row of the part in the query result.
+- `_part_offset` — Number of row in the part.
 - `_partition_id` — Name of a partition.
 - `_part_uuid` — Unique part identifier (if enabled MergeTree setting `assign_part_uuids`).
+- `_part_data_version` — Data version of part (either min block number or mutation version).
 - `_partition_value` — Values (a tuple) of a `partition by` expression.
 - `_sample_factor` — Sample factor (from the query).
 - `_block_number` — Block number of the row, it is persisted on merges when `allow_experimental_block_number_column` is set to true.

src/Interpreters/MutationsInterpreter.cpp

@@ -1283,10 +1283,9 @@ void MutationsInterpreter::Source::read(
             plan,
             *data,
             storage_snapshot,
-            part,
+            RangesInDataPart(part),
             alter_conversions,
             nullptr,
-            0,
             required_columns,
             nullptr,
             apply_deleted_mask_,

src/Processors/QueryPlan/Optimizations/actionsDAGUtils.cpp

@@ -434,4 +434,66 @@ void applyActionsToSortDescription(
     description.resize(prefix_size);
 }
 
+std::optional<std::unordered_map<const ActionsDAG::Node *, const ActionsDAG::Node *>> resolveMatchedInputs(
+    const MatchedTrees::Matches & matches,
+    const std::unordered_set<const ActionsDAG::Node *> & allowed_inputs,
+    const ActionsDAG::NodeRawConstPtrs & nodes)
+{
+    struct Frame
+    {
+        const ActionsDAG::Node * node;
+        size_t next_child_to_visit = 0;
+    };
+
+    std::stack<Frame> stack;
+    std::unordered_set<const ActionsDAG::Node *> visited;
+    std::unordered_map<const ActionsDAG::Node *, const ActionsDAG::Node *> new_inputs;
+
+    for (const auto * node : nodes)
+    {
+        if (visited.contains(node))
+            continue;
+
+        stack.push({.node = node});
+
+        while (!stack.empty())
+        {
+            auto & frame = stack.top();
+
+            if (frame.next_child_to_visit == 0)
+            {
+                auto jt = matches.find(frame.node);
+                if (jt != matches.end())
+                {
+                    const auto & match = jt->second;
+                    if (match.node && !match.monotonicity && allowed_inputs.contains(match.node))
+                    {
+                        visited.insert(frame.node);
+                        new_inputs[frame.node] = match.node;
+                        stack.pop();
+                        continue;
+                    }
+                }
+            }
+
+            if (frame.next_child_to_visit < frame.node->children.size())
+            {
+                stack.push({.node = frame.node->children[frame.next_child_to_visit]});
+                ++frame.next_child_to_visit;
+                continue;
+            }
+
+            /// Not a match and there is no matched child.
+            if (frame.node->type == ActionsDAG::ActionType::INPUT)
+                return std::nullopt;
+
+            /// Not a match, but all children matched.
+            visited.insert(frame.node);
+            stack.pop();
+        }
+    }
+
+    return new_inputs;
+}
+
 }

src/Processors/QueryPlan/Optimizations/actionsDAGUtils.h

@@ -63,4 +63,20 @@ void applyActionsToSortDescription(
     const ActionsDAG & dag,
     /// Ignore this one output of DAG. Used for FilterStep where filter column is removed.
     const ActionsDAG::Node * output_to_skip = nullptr);
+
+/// Traverses a set of root nodes (`nodes`) and attempts to replace them (and their subtrees)
+/// with matched input nodes from the `matches` map, under the condition that:
+/// - The match is non-monotonic (`!match.monotonicity`)
+/// - The matched node is allowed (present in `allowed_inputs`)
+///
+/// This function performs a depth-first traversal to ensure all nodes and their dependencies
+/// can be substituted by valid inputs. If any node ultimately resolves to an INPUT node that
+/// has no valid match in `matches`, the function returns `std::nullopt` to indicate that
+/// a consistent input substitution map cannot be constructed.
+///
+/// The primary use case is to construct the input substitution map required by `ActionsDAG::foldActionsByProjection`.
+std::optional<std::unordered_map<const ActionsDAG::Node *, const ActionsDAG::Node *>> resolveMatchedInputs(
+    const MatchedTrees::Matches & matches,
+    const std::unordered_set<const ActionsDAG::Node *> & allowed_inputs,
+    const ActionsDAG::NodeRawConstPtrs & nodes);
 }

src/Processors/QueryPlan/Optimizations/optimizeUseAggregateProjection.cpp

@@ -335,67 +335,11 @@ std::optional<ActionsDAG> analyzeAggregateProjection(
 
     /// Here we want to match query keys with projection keys.
     /// Query key can be any expression depending on projection keys.
+    auto new_inputs = resolveMatchedInputs(matches, proj_key_nodes, query_key_nodes);
+    if (!new_inputs)
+        return {};
 
-    struct Frame
-    {
-        const ActionsDAG::Node * node;
-        size_t next_child_to_visit = 0;
-    };
-
-    std::stack<Frame> stack;
-    std::unordered_set<const ActionsDAG::Node *> visited;
-    std::unordered_map<const ActionsDAG::Node *, const ActionsDAG::Node *> new_inputs;
-
-    for (const auto * key_node : query_key_nodes)
-    {
-        if (visited.contains(key_node))
-            continue;
-
-        stack.push({.node = key_node});
-
-        while (!stack.empty())
-        {
-            auto & frame = stack.top();
-
-            if (frame.next_child_to_visit == 0)
-            {
-                auto jt = matches.find(frame.node);
-                if (jt != matches.end())
-                {
-                    auto & match = jt->second;
-                    if (match.node && !match.monotonicity && proj_key_nodes.contains(match.node))
-                    {
-                        visited.insert(frame.node);
-                        new_inputs[frame.node] = match.node;
-                        stack.pop();
-                        continue;
-                    }
-                }
-            }
-
-            if (frame.next_child_to_visit < frame.node->children.size())
-            {
-                stack.push({.node = frame.node->children[frame.next_child_to_visit]});
-                ++frame.next_child_to_visit;
-                continue;
-            }
-
-            /// Not a match and there is no matched child.
-            if (frame.node->type == ActionsDAG::ActionType::INPUT)
-            {
-                // LOG_TRACE(getLogger("optimizeUseProjections"), "Cannot find match for {}", frame.node->result_name);
-                return {};
-            }
-
-            /// Not a match, but all children matched.
-            visited.insert(frame.node);
-            stack.pop();
-        }
-    }
-
-    // LOG_TRACE(getLogger("optimizeUseProjections"), "Folding actions by projection");
-
-    auto proj_dag = ActionsDAG::foldActionsByProjection(new_inputs, query_key_nodes);
+    auto proj_dag = ActionsDAG::foldActionsByProjection(*new_inputs, query_key_nodes);
     appendAggregateFunctions(proj_dag, aggregates, *matched_aggregates);
     return proj_dag;
 }

src/Processors/QueryPlan/Optimizations/projectionsCommon.cpp

@@ -229,17 +229,22 @@ bool analyzeProjectionCandidate(
     const PartitionIdToMaxBlockPtr & max_added_blocks,
     const ActionsDAG * dag)
 {
-    MergeTreeData::DataPartsVector projection_parts;
-    MergeTreeData::DataPartsVector normal_parts;
+    RangesInDataParts projection_parts;
+    RangesInDataParts normal_parts;
 
     for (const auto & part_with_ranges : parts_with_ranges)
     {
         const auto & created_projections = part_with_ranges.data_part->getProjectionParts();
         auto it = created_projections.find(candidate.projection->name);
         if (it != created_projections.end() && !it->second->is_broken)
-            projection_parts.push_back(it->second);
+        {
+            projection_parts.push_back(
+                RangesInDataPart(it->second, part_with_ranges.part_index_in_query, part_with_ranges.part_starting_offset_in_query));
+        }
         else
-            normal_parts.push_back(part_with_ranges.data_part);
+        {
+            normal_parts.push_back(part_with_ranges);
+        }
     }
 
     if (projection_parts.empty())

src/Processors/QueryPlan/PartsSplitter.cpp

@@ -251,6 +251,7 @@ class RangesInDataPartsBuilder
             ranges_in_data_parts.emplace_back(
                 initial_ranges_in_data_parts[part_index].data_part,
                 initial_ranges_in_data_parts[part_index].part_index_in_query,
+                initial_ranges_in_data_parts[part_index].part_starting_offset_in_query,
                 MarkRanges{mark_range});
             part_index_to_initial_ranges_in_data_parts_index[it->second] = part_index;
             return;

src/Processors/QueryPlan/ReadFromMergeTree.cpp

@@ -90,12 +90,6 @@ size_t countPartitions(const RangesInDataParts & parts_with_ranges)
     return countPartitions(parts_with_ranges, get_partition_id);
 }
 
-size_t countPartitions(const MergeTreeData::DataPartsVector & prepared_parts)
-{
-    auto get_partition_id = [](const MergeTreeData::DataPartPtr data_part) { return data_part->info.getPartitionId(); };
-    return countPartitions(prepared_parts, get_partition_id);
-}
-
 bool restoreDAGInputs(ActionsDAG & dag, const NameSet & inputs)
 {
     std::unordered_set<const ActionsDAG::Node *> outputs(dag.getOutputs().begin(), dag.getOutputs().end());
@@ -336,7 +330,7 @@ void ReadFromMergeTree::AnalysisResult::checkLimits(const Settings & settings, c
 }
 
 ReadFromMergeTree::ReadFromMergeTree(
-    MergeTreeData::DataPartsVector parts_,
+    RangesInDataParts parts_,
     MergeTreeData::MutationsSnapshotPtr mutations_,
     Names all_column_names_,
     const MergeTreeData & data_,
@@ -1245,7 +1239,8 @@ Pipe ReadFromMergeTree::spreadMarkRangesAmongStreamsWithOrder(
                 }
 
                 ranges_to_get_from_part = split_ranges(ranges_to_get_from_part, input_order_info->direction);
-                new_parts.emplace_back(part.data_part, part.part_index_in_query, std::move(ranges_to_get_from_part));
+                new_parts.emplace_back(
+                    part.data_part, part.part_index_in_query, part.part_starting_offset_in_query, std::move(ranges_to_get_from_part));
             }
 
             split_parts_and_ranges.emplace_back(std::move(new_parts));
@@ -1488,7 +1483,8 @@ Pipe ReadFromMergeTree::spreadMarkRangesAmongStreamsFinal(
             RangesInDataParts new_parts;
 
             for (auto part_it = parts_to_merge_ranges[range_index]; part_it != parts_to_merge_ranges[range_index + 1]; ++part_it)
-                new_parts.emplace_back(part_it->data_part, part_it->part_index_in_query, part_it->ranges);
+                new_parts.emplace_back(
+                    part_it->data_part, part_it->part_index_in_query, part_it->part_starting_offset_in_query, part_it->ranges);
 
             if (new_parts.empty())
                 continue;
@@ -1631,7 +1627,7 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(bool
     return selectRangesToRead(prepared_parts, find_exact_ranges);
 }
 
-ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(MergeTreeData::DataPartsVector parts, bool find_exact_ranges) const
+ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(RangesInDataParts parts, bool find_exact_ranges) const
 {
     return selectRangesToRead(
         std::move(parts),
@@ -1653,7 +1649,7 @@ static void buildIndexes(
     std::optional<ReadFromMergeTree::Indexes> & indexes,
     const ActionsDAG * filter_actions_dag,
     const MergeTreeData & data,
-    const MergeTreeData::DataPartsVector & parts,
+    const RangesInDataParts & parts,
     const MergeTreeData::MutationsSnapshotPtr & mutations_snapshot,
     [[maybe_unused]] const std::optional<VectorSearchParameters> & vector_search_parameters,
     const ContextPtr & context,
@@ -1684,8 +1680,15 @@ static void buildIndexes(
         indexes->partition_pruner.emplace(metadata_snapshot, filter_dag, context, false /* strict */);
     }
 
-    indexes->part_values = MergeTreeDataSelectExecutor::filterPartsByVirtualColumns(metadata_snapshot, data, parts, filter_dag.predicate, context);
-    MergeTreeDataSelectExecutor::buildKeyConditionFromPartOffset(indexes->part_offset_condition, filter_dag.predicate, context);
+    indexes->part_values
+        = MergeTreeDataSelectExecutor::filterPartsByVirtualColumns(metadata_snapshot, data, parts, filter_dag.predicate, context);
+
+    /// Perform virtual column key analysis only when no corresponding physical columns exist.
+    const auto & columns = metadata_snapshot->getColumns();
+    if (!columns.has("_part_offset") && !columns.has("_part"))
+        MergeTreeDataSelectExecutor::buildKeyConditionFromPartOffset(indexes->part_offset_condition, filter_dag.predicate, context);
+    if (!columns.has("_part_offset") && !columns.has("_part_starting_offset"))
+        MergeTreeDataSelectExecutor::buildKeyConditionFromTotalOffset(indexes->total_offset_condition, filter_dag.predicate, context);
 
     indexes->use_skip_indexes = settings[Setting::use_skip_indexes];
     if (query_info.isFinal() && !settings[Setting::use_skip_indexes_if_final])
@@ -1834,7 +1837,7 @@ void ReadFromMergeTree::applyFilters(ActionDAGNodes added_filter_nodes)
 }
 
 ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
-    MergeTreeData::DataPartsVector parts,
+    RangesInDataParts parts,
     MergeTreeData::MutationsSnapshotPtr mutations_snapshot,
     const std::optional<VectorSearchParameters> & vector_search_parameters,
     const StorageMetadataPtr & metadata_snapshot,
@@ -1867,7 +1870,17 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
     const Names & primary_key_column_names = primary_key.column_names;
 
     if (!indexes)
-        buildIndexes(indexes, query_info_.filter_actions_dag.get(), data, parts, mutations_snapshot, vector_search_parameters, context_, query_info_, metadata_snapshot, log);
+        buildIndexes(
+            indexes,
+            query_info_.filter_actions_dag.get(),
+            data,
+            parts,
+            mutations_snapshot,
+            vector_search_parameters,
+            context_,
+            query_info_,
+            metadata_snapshot,
+            log);
 
     if (indexes->part_values && indexes->part_values->empty())
         return std::make_shared<AnalysisResult>(std::move(result));
@@ -1890,6 +1903,9 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
     if (indexes->part_offset_condition)
         LOG_DEBUG(log, "Part offset condition: {}", indexes->part_offset_condition->toString());
 
+    if (indexes->total_offset_condition)
+        LOG_DEBUG(log, "Total offset condition: {}", indexes->total_offset_condition->toString());
+
     if (indexes->key_condition.alwaysFalse())
         return std::make_shared<AnalysisResult>(std::move(result));
 
@@ -1924,7 +1940,7 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
             return std::make_shared<AnalysisResult>(std::move(result));
 
         for (const auto & part : parts)
-            total_marks_pk += part->index_granularity->getMarksCountWithoutFinal();
+            total_marks_pk += part.data_part->index_granularity->getMarksCountWithoutFinal();
         parts_before_pk = parts.size();
 
         auto reader_settings = getMergeTreeReaderSettings(context_, query_info_);
@@ -1934,6 +1950,7 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
             context_,
             indexes->key_condition,
             indexes->part_offset_condition,
+            indexes->total_offset_condition,
             indexes->skip_indexes,
             reader_settings,
             log,
@@ -1945,12 +1962,11 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
 
         MergeTreeDataSelectExecutor::filterPartsByQueryConditionCache(result.parts_with_ranges, query_info_, context_, log);
 
-        if (indexes->use_skip_indexes && !indexes->skip_indexes.useful_indices.empty()  && query_info_.isFinal() && settings[Setting::use_skip_indexes_if_final_exact_mode])
+        if (indexes->use_skip_indexes && !indexes->skip_indexes.useful_indices.empty() && query_info_.isFinal()
+            && settings[Setting::use_skip_indexes_if_final_exact_mode])
         {
-            result.parts_with_ranges = findPKRangesForFinalAfterSkipIndex(primary_key,
-                                                                          metadata_snapshot->getSortingKey(),
-                                                                          result.parts_with_ranges,
-                                                                          log);
+            result.parts_with_ranges
+                = findPKRangesForFinalAfterSkipIndex(primary_key, metadata_snapshot->getSortingKey(), result.parts_with_ranges, log);
             add_index_stat_row_for_pk_expand = true;
         }
     }
@@ -2135,7 +2151,7 @@ bool ReadFromMergeTree::requestOutputEachPartitionThroughSeparatePort()
     {
         std::unordered_map<String, size_t> partition_rows;
         for (const auto & part : prepared_parts)
-            partition_rows[part->info.getPartitionId()] += part->rows_count;
+            partition_rows[part.data_part->info.getPartitionId()] += part.data_part->rows_count;
         size_t sum_rows = 0;
         size_t max_rows = 0;
         for (const auto & [_, rows] : partition_rows)