Merge pull request ClickHouse#78350 from shankar-iyer/additional_primary_key_scan_for_final_with_skip_index

Additional primary key scan after using skip index for query with FINAL clause.

Author: shankar-iyer

src/Core/Settings.cpp

@@ -1410,6 +1410,16 @@ By default, this setting is disabled because skip indexes may exclude rows (gran
 
 Possible values:
 
+- 0 — Disabled.
+- 1 — Enabled.
+)", 0) \
+    DECLARE(Bool, use_skip_indexes_if_final_exact_mode, 0, R"(
+Controls whether granules returned by a skipping index are expanded in newer parts to return correct results when executing a query with the FINAL modifier.
+
+Using skip indexes may exclude rows (granules) containing the latest data which could lead to incorrect results. This setting can ensure that correct results are returned by scanning newer parts that have overlap with the ranges returned by the skip index.
+
+Possible values:
+
 - 0 — Disabled.
 - 1 — Enabled.
 )", 0) \

src/Core/SettingsChangesHistory.cpp

@@ -71,6 +71,7 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory()
             {"geotoh3_lon_lat_input_order", true, false, "A new setting for legacy behaviour to set lon and lat order"},
             {"secondary_indices_enable_bulk_filtering", false, true, "A new algorithm for filtering by data skipping indices"},
             {"implicit_table_at_top_level", "", "", "A new setting, used in clickhouse-local"},
+            {"use_skip_indexes_if_final_exact_mode", 0, 0, "This setting was introduced to help FINAL query return correct results with skip indexes"},
         });
         addSettingsChanges(settings_changes_history, "25.4",
         {

src/Processors/QueryPlan/PartsSplitter.cpp

@@ -256,7 +256,10 @@ class RangesInDataPartsBuilder
             return;
         }
 
-        ranges_in_data_parts[it->second].ranges.push_back(mark_range);
+        if (ranges_in_data_parts[it->second].ranges.back().end == mark_range.begin)
+            ranges_in_data_parts[it->second].ranges.back().end = mark_range.end;
+        else
+            ranges_in_data_parts[it->second].ranges.push_back(mark_range);
     }
 
     RangesInDataParts & getCurrentRangesInDataParts()
@@ -289,6 +292,10 @@ struct PartsRangesIterator
 
         if (event == other.event)
         {
+            if (!selected && other.selected)
+                return true;
+            if (selected && !other.selected)
+                return false;
             if (part_index == other.part_index)
             {
                 /// Within the same part we should process events in order of mark numbers,
@@ -347,6 +354,7 @@ struct PartsRangesIterator
     MarkRange range;
     size_t part_index;
     EventType event;
+    bool selected; /// Whether this range was selected or rejected in skip index filtering
 };
 
 struct PartRangeIndex
@@ -475,7 +483,8 @@ SplitPartsRangesResult splitPartsRanges(RangesInDataParts ranges_in_data_parts,
                  in_reverse_order,
                  range,
                  part_index,
-                 PartsRangesIterator::EventType::RangeStart});
+                 PartsRangesIterator::EventType::RangeStart,
+                 false});
 
             const bool value_is_defined_at_end_mark = range.end < index_granularity->getMarksCount();
             if (!value_is_defined_at_end_mark)
@@ -486,7 +495,8 @@ SplitPartsRangesResult splitPartsRanges(RangesInDataParts ranges_in_data_parts,
                  in_reverse_order,
                  range,
                  part_index,
-                 PartsRangesIterator::EventType::RangeEnd});
+                 PartsRangesIterator::EventType::RangeEnd,
+                 false});
         }
     }
 
@@ -713,7 +723,8 @@ SplitPartsByRanges splitIntersectingPartsRangesIntoLayers(
                 in_reverse_order,
                 range,
                 part_index,
-                PartsRangesIterator::EventType::RangeStart};
+                PartsRangesIterator::EventType::RangeStart,
+                false};
             PartRangeIndex parts_range_start_index(parts_range_start);
             parts_ranges_queue.push({std::move(parts_range_start), std::move(parts_range_start_index)});
 
@@ -726,7 +737,8 @@ SplitPartsByRanges splitIntersectingPartsRangesIntoLayers(
                 in_reverse_order,
                 range,
                 part_index,
-                PartsRangesIterator::EventType::RangeEnd};
+                PartsRangesIterator::EventType::RangeEnd,
+                false};
             PartRangeIndex parts_range_end_index(parts_range_end);
             parts_ranges_queue.push({std::move(parts_range_end), std::move(parts_range_end_index)});
         }
@@ -911,6 +923,130 @@ static ASTs buildFilters(const KeyDescription & primary_key, const std::vector<V
     return filters;
 }
 
+RangesInDataParts findPKRangesForFinalAfterSkipIndexImpl(RangesInDataParts & ranges_in_data_parts, bool cannot_sort_primary_key, const LoggerPtr & logger)
+{
+    IndexAccess index_access(ranges_in_data_parts);
+    std::vector<PartsRangesIterator> selected_ranges;
+    std::vector<PartsRangesIterator> rejected_ranges;
+
+    RangesInDataPartsBuilder result(ranges_in_data_parts);
+
+    auto skip_and_return_all_part_ranges = [&]()
+    {
+        RangesInDataParts all_part_ranges(std::move(ranges_in_data_parts));
+        for (auto & all_part_range : all_part_ranges)
+        {
+            const auto & index_granularity = all_part_range.data_part->index_granularity;
+            all_part_range.ranges = MarkRanges{{MarkRange{0, index_granularity->getMarksCountWithoutFinal()}}};
+        }
+        return all_part_ranges;
+    };
+
+    if (cannot_sort_primary_key) /// just expand to all parts + ranges
+    {
+        return skip_and_return_all_part_ranges();
+    }
+
+    for (size_t part_index = 0; part_index < ranges_in_data_parts.size(); ++part_index)
+    {
+        const auto & index_granularity = ranges_in_data_parts[part_index].data_part->index_granularity;
+        std::vector<bool> is_selected_range(index_granularity->getMarksCountWithoutFinal(), false);
+        for (const auto & range : ranges_in_data_parts[part_index].ranges)
+        {
+            const bool value_is_defined_at_end_mark = range.end < index_granularity->getMarksCount();
+            if (!value_is_defined_at_end_mark)
+            {
+                return skip_and_return_all_part_ranges();
+            }
+
+            selected_ranges.push_back(
+                {index_access.getValue(part_index, range.begin), false, range, part_index, PartsRangesIterator::EventType::RangeStart, true});
+            for (auto i = range.begin; i < range.end;i++)
+               is_selected_range[i] = true;
+        }
+
+        for (size_t range_begin = 0; range_begin < is_selected_range.size(); range_begin++)
+        {
+            const bool value_is_defined_at_end_mark = ((range_begin + 1) < index_granularity->getMarksCount());
+            if (!value_is_defined_at_end_mark)
+            {
+                return skip_and_return_all_part_ranges();
+            }
+
+            if (is_selected_range[range_begin])
+                continue;
+            MarkRange rejected_range(range_begin, range_begin + 1);
+            rejected_ranges.push_back(
+                {index_access.getValue(part_index, rejected_range.begin), false, rejected_range, part_index, PartsRangesIterator::EventType::RangeStart, false});
+        }
+    }
+
+    ::sort(selected_ranges.begin(), selected_ranges.end());
+
+    ::sort(rejected_ranges.begin(), rejected_ranges.end());
+
+    LOG_TRACE(logger, "findPKRangesForFinalAfterSkipIndex : sorting phase complete");
+
+    std::vector<PartsRangesIterator>::iterator selected_ranges_iter = selected_ranges.begin();
+    std::vector<PartsRangesIterator>::iterator rejected_ranges_iter = rejected_ranges.begin();
+    size_t more_ranges_added = 0;
+
+    while (selected_ranges_iter != selected_ranges.end() && rejected_ranges_iter != rejected_ranges.end())
+    {
+        auto selected_range_start = selected_ranges_iter->value;
+        auto selected_range_end = index_access.getValue(selected_ranges_iter->part_index, selected_ranges_iter->range.end);
+        auto rejected_range_start = rejected_ranges_iter->value;
+
+        int result1 = compareValues(rejected_range_start, selected_range_start, false);
+        int result2 = compareValues(rejected_range_start, selected_range_end, false);
+
+        if (result1 == 0 || result2 == 0 || (result1 > 0 && result2 < 0)) /// rejected_range_start inside [selected_range]
+        {
+            result.addRange(rejected_ranges_iter->part_index, rejected_ranges_iter->range);
+            rejected_ranges_iter++;
+            more_ranges_added++;
+        }
+        else if (result1 > 0) /// rejected_range_start beyond [selected_range]
+        {
+            result.addRange(selected_ranges_iter->part_index, selected_ranges_iter->range);
+            selected_ranges_iter++;
+        }
+        else
+        {
+            auto rejected_range_end = index_access.getValue(rejected_ranges_iter->part_index, rejected_ranges_iter->range.end);
+            int result3 = compareValues(rejected_range_end, selected_range_start, false);
+            int result4 = compareValues(rejected_range_end, selected_range_end, false);
+            /// rejected_range_end inside [selected range] OR [rejected range] encompasses [selected range]
+            if (result3 == 0 || result4 == 0 || (result3 > 0 && result4 < 0) || (result1 < 0 && result4 > 0))
+            {
+                result.addRange(rejected_ranges_iter->part_index, rejected_ranges_iter->range);
+                more_ranges_added++;
+            }
+            rejected_ranges_iter++;
+        }
+    }
+
+    while (selected_ranges_iter != selected_ranges.end())
+    {
+        result.addRange(selected_ranges_iter->part_index, selected_ranges_iter->range);
+        selected_ranges_iter++;
+    }
+
+    auto result_final_ranges = result.getCurrentRangesInDataParts();
+    std::stable_sort(
+        result_final_ranges.begin(),
+        result_final_ranges.end(),
+        [](const auto & lhs, const auto & rhs) { return lhs.part_index_in_query < rhs.part_index_in_query; });
+    for (auto & result_final_range : result_final_ranges)
+    {
+        std::sort(result_final_range.ranges.begin(), result_final_range.ranges.end());
+    }
+
+    LOG_TRACE(logger, "findPKRangesForFinalAfterSkipIndex : processed {} parts, initially selected {} ranges & rejected {}, more {} ranges added", ranges_in_data_parts.size(), selected_ranges.size(), rejected_ranges.size(), more_ranges_added);
+
+    return result_final_ranges;
+}
+
 static void reorderColumns(ActionsDAG & dag, const Block & header, const std::string & filter_column)
 {
     std::unordered_map<std::string_view, const ActionsDAG::Node *> inputs_map;
@@ -1056,4 +1192,18 @@ Pipes readByLayers(
     return merging_pipes;
 }
 
+RangesInDataParts findPKRangesForFinalAfterSkipIndex(
+    const KeyDescription & primary_key,
+    const KeyDescription & sorting_key,
+    RangesInDataParts & ranges_in_data_parts,
+    const LoggerPtr & logger)
+{
+    bool cannot_sort_primary_key = false;
+    if (!isSafePrimaryKey(primary_key) || !sorting_key.reverse_flags.empty())
+    {
+        LOG_TRACE(logger, "Primary key is not sortable, expanding PK range to entire due to exact_mode.");
+        cannot_sort_primary_key = true;
+    }
+    return findPKRangesForFinalAfterSkipIndexImpl(ranges_in_data_parts, cannot_sort_primary_key, logger);
+}
 }

src/Processors/QueryPlan/PartsSplitter.h

@@ -38,6 +38,18 @@ SplitPartsWithRangesByPrimaryKeyResult splitPartsWithRangesByPrimaryKey(
     bool split_parts_ranges_into_intersecting_and_non_intersecting,
     bool split_intersecting_parts_ranges_into_layers);
 
+/**
+  *
+  * If setting use_skip_indexes_if_final_exact_mode=1, then we need to expand
+  * the initial set of granules returned from the skip index by
+  * looking for that initial set of PK ranges across all other newer parts.
+  */
+RangesInDataParts findPKRangesForFinalAfterSkipIndex(
+    const KeyDescription & primary_key,
+    const KeyDescription & sorting_key,
+    RangesInDataParts & ranges_in_data_parts,
+    const LoggerPtr & logger);
+
 struct SplitPartsByRanges
 {
     using Values = std::vector<Field>;

src/Processors/QueryPlan/ReadFromMergeTree.cpp

@@ -186,6 +186,7 @@ namespace Setting
     extern const SettingsBool query_plan_merge_filters;
     extern const SettingsUInt64 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_query_condition_cache;
     extern const SettingsBool query_condition_cache_store_conditions_as_plaintext;
     extern const SettingsBool allow_experimental_analyzer;
@@ -1891,6 +1892,7 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
 
     size_t total_marks_pk = 0;
     size_t parts_before_pk = 0;
+    bool add_index_stat_row_for_pk_expand = false;
 
     {
         MergeTreeDataSelectExecutor::filterPartsByPartition(
@@ -1935,9 +1937,19 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
             num_streams,
             result.index_stats,
             indexes->use_skip_indexes,
-            find_exact_ranges);
+            find_exact_ranges,
+            query_info_.isFinal());
 
         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])
+        {
+            result.parts_with_ranges = findPKRangesForFinalAfterSkipIndex(primary_key,
+                                                                          metadata_snapshot->getSortingKey(),
+                                                                          result.parts_with_ranges,
+                                                                          log);
+            add_index_stat_row_for_pk_expand = true;
+        }
     }
 
     size_t sum_marks_pk = total_marks_pk;
@@ -1956,6 +1968,15 @@ ReadFromMergeTree::AnalysisResultPtr ReadFromMergeTree::selectRangesToRead(
         sum_rows += part.getRowsCount();
     }
 
+    if (add_index_stat_row_for_pk_expand)
+    {
+        result.index_stats.emplace_back(ReadFromMergeTree::IndexStat{
+            .type = ReadFromMergeTree::IndexType::PrimaryKeyExpand,
+            .description = "Selects all granules that intersect by PK values with the previous skip indexes selection",
+            .num_parts_after = result.parts_with_ranges.size(),
+            .num_granules_after = sum_marks});
+    }
+
     result.total_parts = total_parts;
     result.parts_before_pk = parts_before_pk;
     result.selected_parts = result.parts_with_ranges.size();
@@ -2424,6 +2445,8 @@ static const char * indexTypeToString(ReadFromMergeTree::IndexType type)
             return "PrimaryKey";
         case ReadFromMergeTree::IndexType::Skip:
             return "Skip";
+        case ReadFromMergeTree::IndexType::PrimaryKeyExpand:
+            return "PrimaryKeyExpand";
     }
 }
 

src/Processors/QueryPlan/ReadFromMergeTree.h

@@ -71,6 +71,7 @@ class ReadFromMergeTree final : public SourceStepWithFilter
         Partition,
         PrimaryKey,
         Skip,
+        PrimaryKeyExpand,
     };
 
     /// This is a struct with information about applied indexes.

src/Storages/MergeTree/MergeTreeDataSelectExecutor.cpp

@@ -83,6 +83,7 @@ namespace Setting
     extern const SettingsUInt64 parallel_replica_offset;
     extern const SettingsUInt64 parallel_replicas_count;
     extern const SettingsParallelReplicasMode parallel_replicas_mode;
+    extern const SettingsBool use_skip_indexes_if_final_exact_mode;
     extern const SettingsBool use_query_condition_cache;
     extern const SettingsBool allow_experimental_analyzer;
     extern const SettingsBool parallel_replicas_local_plan;
@@ -619,7 +620,8 @@ RangesInDataParts MergeTreeDataSelectExecutor::filterPartsByPrimaryKeyAndSkipInd
     size_t num_streams,
     ReadFromMergeTree::IndexStats & index_stats,
     bool use_skip_indexes,
-    bool find_exact_ranges)
+    bool find_exact_ranges,
+    bool is_final_query)
 {
     const Settings & settings = context->getSettingsRef();
 
@@ -677,6 +679,7 @@ RangesInDataParts MergeTreeDataSelectExecutor::filterPartsByPrimaryKeyAndSkipInd
     std::atomic<size_t> sum_parts_pk = 0;
 
     RangesInDataParts parts_with_ranges(parts.size());
+    std::vector<size_t> skip_index_used_in_part(parts.size(), 0);
 
     /// Let's find what range to read from each part.
     {
@@ -746,6 +749,7 @@ RangesInDataParts MergeTreeDataSelectExecutor::filterPartsByPrimaryKeyAndSkipInd
                 stat.granules_dropped.fetch_add(total_granules - ranges.ranges.getNumberOfMarks(), std::memory_order_relaxed);
                 if (ranges.ranges.empty())
                     stat.parts_dropped.fetch_add(1, std::memory_order_relaxed);
+                skip_index_used_in_part[part_index] = 1; /// thread-safe
             }
 
             for (size_t idx = 0; idx < skip_indexes.merged_indices.size(); ++idx)
@@ -771,8 +775,7 @@ RangesInDataParts MergeTreeDataSelectExecutor::filterPartsByPrimaryKeyAndSkipInd
                     stat.parts_dropped.fetch_add(1, std::memory_order_relaxed);
             }
 
-            if (!ranges.ranges.empty())
-                parts_with_ranges[part_index] = std::move(ranges);
+            parts_with_ranges[part_index] = std::move(ranges);
         };
 
         size_t num_threads = std::min<size_t>(num_streams, parts.size());
@@ -825,7 +828,12 @@ RangesInDataParts MergeTreeDataSelectExecutor::filterPartsByPrimaryKeyAndSkipInd
         for (size_t part_index = 0; part_index < parts.size(); ++part_index)
         {
             auto & part = parts_with_ranges[part_index];
-            if (!part.data_part)
+            if (is_final_query && settings[Setting::use_skip_indexes_if_final_exact_mode] && skip_index_used_in_part[part_index])
+            {
+                ++next_part; /// retain this part even if empty due to FINAL
+                continue;
+            }
+            if (!part.data_part || part.ranges.empty())
                 continue;
 
             if (next_part != part_index)