[fix](ann-index) Fix ANN range search state leakage and incorrect slot index tracking. (#63666)

Issue Number: close #xxx

Related PR: #xxx

Problem Summary:

### Release note

ANN range search execution state was stored on shared VExpr roots.
VExprContext clones share the root expression, so a segment that
executed ANN range search could leak that state into another segment
without an ANN index and incorrectly remove the common expression. ANN
range search also mixed schema column indexes with storage column ids
when updating common expression index status, so remapped schemas failed
to mark the source slot expression as evaluated. This patch returns ANN
execution state through the current evaluation call, stores ANN root
bitmap in the current segment IndexContext, and updates slot index
status by source column index.

### Check List (For Author)

- Test <!-- At least one of them must be included. -->
    - [ ] Regression test
    - [ ] Unit Test
    - [ ] Manual test (add detailed scripts or steps below)
    - [ ] No need to test or manual test. Explain why:
- [ ] This is a refactor/code format and no logic has been changed.
        - [ ] Previous test can cover this change.
        - [ ] No code files have been changed.
        - [ ] Other reason <!-- Add your reason?  -->

- Behavior changed:
    - [ ] No.
    - [ ] Yes. <!-- Explain the behavior change -->

- Does this need documentation?
    - [ ] No.
- [ ] Yes. <!-- Add document PR link here. eg:
apache/doris-website#1214 -->

### Check List (For Reviewer who merge this PR)

- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label <!-- Add branch pick label that this PR
should merge into -->

Author: kaka11chen

be/src/exprs/vectorized_fn_call.cpp

@@ -557,7 +557,8 @@ Status VectorizedFnCall::evaluate_ann_range_search(
         const std::vector<ColumnId>& idx_to_cid,
         const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
         roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-        bool enable_result_cache) {
+        bool enable_result_cache, AnnRangeSearchEvaluationResult& evaluation_result) {
+    evaluation_result = {};
     if (range_search_runtime.is_ann_range_search == false) {
         return Status::OK();
     }
@@ -566,8 +567,8 @@ Status VectorizedFnCall::evaluate_ann_range_search(
                               range_search_runtime.to_string());
     size_t origin_num = row_bitmap.cardinality();
 
-    int idx_in_block = static_cast<int>(range_search_runtime.src_col_idx);
-    DCHECK(idx_in_block < idx_to_cid.size())
+    const auto idx_in_block = range_search_runtime.src_col_idx;
+    DCHECK_LT(idx_in_block, idx_to_cid.size())
             << "idx_in_block: " << idx_in_block << ", idx_to_cid.size(): " << idx_to_cid.size();
 
     ColumnId src_col_cid = idx_to_cid[idx_in_block];
@@ -649,6 +650,7 @@ Status VectorizedFnCall::evaluate_ann_range_search(
     row_bitmap = *result.roaring;
 
     // Process virtual column
+    bool dist_fulfilled = false;
     if (range_search_runtime.dst_col_idx >= 0) {
         // Prepare materialization if we can use result from index.
         // Typical situation: range search and operator is LE or LT.
@@ -672,7 +674,7 @@ Status VectorizedFnCall::evaluate_ann_range_search(
             }
             virtual_column_iterator->prepare_materialization(std::move(distance_col),
                                                              std::move(result.row_ids));
-            _virtual_column_is_fulfilled = true;
+            dist_fulfilled = true;
         } else {
             // Whether the ANN index should have produced distance depends on metric and operator:
             //  - L2: distance is produced for LE/LT; not produced for GE/GT
@@ -686,17 +688,17 @@ Status VectorizedFnCall::evaluate_ann_range_search(
             // If we expected distance but didn't get it, assert in debug to catch logic errors.
             DCHECK(!should_have_distance) << "Expected distance from ANN index but got none";
 #endif
-            _virtual_column_is_fulfilled = false;
         }
     } else {
         // Dest is not virtual column.
-        _virtual_column_is_fulfilled = true;
+        dist_fulfilled = true;
     }
 
-    _has_been_executed = true;
+    evaluation_result.executed = true;
+    evaluation_result.dist_fulfilled = dist_fulfilled;
     VLOG_DEBUG << fmt::format(
             "Ann range search filtered {} rows, origin {} rows, virtual column is full-filled: {}",
-            origin_num - row_bitmap.cardinality(), origin_num, _virtual_column_is_fulfilled);
+            origin_num - row_bitmap.cardinality(), origin_num, dist_fulfilled);
 
     ann_index_stats = *stats;
     return Status::OK();

be/src/exprs/vectorized_fn_call.h

@@ -92,7 +92,7 @@ class VectorizedFnCall : public VExpr {
             const std::vector<ColumnId>& idx_to_cid,
             const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
             roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-            bool enable_result_cache) override;
+            bool enable_result_cache, AnnRangeSearchEvaluationResult& result) override;
 
     void prepare_ann_range_search(const doris::VectorSearchUserParams& params,
                                   segment_v2::AnnRangeSearchRuntime& runtime,

be/src/exprs/vexpr.cpp

@@ -1038,7 +1038,9 @@ Status VExpr::evaluate_ann_range_search(
         const std::vector<std::unique_ptr<segment_v2::IndexIterator>>& index_iterators,
         const std::vector<ColumnId>& idx_to_cid,
         const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
-        roaring::Roaring& row_bitmap, AnnIndexStats& ann_index_stats, bool enable_result_cache) {
+        roaring::Roaring& row_bitmap, AnnIndexStats& ann_index_stats, bool enable_result_cache,
+        AnnRangeSearchEvaluationResult& result) {
+    result = {};
     return Status::OK();
 }
 
@@ -1056,14 +1058,6 @@ void VExpr::prepare_ann_range_search(const doris::VectorSearchUserParams& params
     }
 }
 
-bool VExpr::ann_range_search_executedd() {
-    return _has_been_executed;
-}
-
-bool VExpr::ann_dist_is_fulfilled() const {
-    return _virtual_column_is_fulfilled;
-}
-
 Status VExpr::execute_filter(VExprContext* context, const Block* block,
                              uint8_t* __restrict result_filter_data, size_t rows, bool accept_null,
                              bool* can_filter_all) const {

be/src/exprs/vexpr.h

@@ -79,6 +79,15 @@ struct AnnRangeSearchRuntime;
 
 using Selector = IColumn::Selector;
 
+struct AnnRangeSearchEvaluationResult {
+    // Indicates whether the expr row_bitmap has been updated.
+    bool executed = false;
+    // Indicates whether the virtual column is fulfilled.
+    // NOTE, if there is no virtual column in the expr tree, and expr
+    // is evaluated by ann index, this flag is still true.
+    bool dist_fulfilled = false;
+};
+
 class VExpr {
 public:
     // resize inserted param column to make sure column size equal to block.rows() and return param column index
@@ -346,7 +355,7 @@ class VExpr {
             const std::vector<ColumnId>& idx_to_cid,
             const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
             roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-            bool enable_result_cache);
+            bool enable_result_cache, AnnRangeSearchEvaluationResult& result);
 
     // Prepare the runtime for ANN range search.
     // AnnRangeSearchRuntime is used to store the runtime information of ann range search.
@@ -356,10 +365,6 @@ class VExpr {
                                           segment_v2::AnnRangeSearchRuntime& range_search_runtime,
                                           bool& suitable_for_ann_index);
 
-    bool ann_range_search_executedd();
-
-    bool ann_dist_is_fulfilled() const;
-
     virtual uint64_t get_digest(uint64_t seed) const;
 
 protected:
@@ -442,13 +447,6 @@ class VExpr {
     // ensuring uniqueness during index traversal
     uint32_t _index_unique_id = 0;
     bool _enable_inverted_index_query = true;
-
-    // Indicates whether the expr row_bitmap has been updated.
-    bool _has_been_executed = false;
-    // Indicates whether the virtual column is fulfilled.
-    // NOTE, if there is no virtual column in the expr tree, and expr
-    // is evaluated by ann index, this flag is still true.
-    bool _virtual_column_is_fulfilled = false;
 };
 
 // NOLINTBEGIN(readability-function-size)

be/src/exprs/vexpr_context.cpp

@@ -29,6 +29,7 @@
 #include "core/block/columns_with_type_and_name.h"
 #include "core/column/column.h"
 #include "core/column/column_const.h"
+#include "exec/common/util.hpp"
 #include "exprs/function_context.h"
 #include "exprs/vexpr.h"
 #include "runtime/runtime_state.h"
@@ -440,41 +441,57 @@ Status VExprContext::evaluate_ann_range_search(
         const std::unordered_map<VExprContext*, std::unordered_map<ColumnId, VExpr*>>&
                 common_expr_to_slotref_map,
         roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-        bool enable_result_cache) {
+        bool enable_result_cache, bool* ann_range_search_executed) {
+    if (ann_range_search_executed != nullptr) {
+        *ann_range_search_executed = false;
+    }
     if (_root == nullptr) {
         return Status::OK();
     }
 
+    AnnRangeSearchEvaluationResult evaluation_result;
     RETURN_IF_ERROR(_root->evaluate_ann_range_search(
             _ann_range_search_runtime, cid_to_index_iterators, idx_to_cid, column_iterators,
-            row_bitmap, ann_index_stats, enable_result_cache));
+            row_bitmap, ann_index_stats, enable_result_cache, evaluation_result));
 
-    if (!_root->ann_range_search_executedd()) {
+    if (!evaluation_result.executed) {
         return Status::OK();
     }
+    if (ann_range_search_executed != nullptr) {
+        *ann_range_search_executed = true;
+    }
+
+    DCHECK(_index_context != nullptr);
+    _index_context->set_index_result_for_expr(
+            _root.get(),
+            segment_v2::InvertedIndexResultBitmap(std::make_shared<roaring::Roaring>(row_bitmap),
+                                                  std::make_shared<roaring::Roaring>()));
 
-    if (!_root->ann_dist_is_fulfilled()) {
+    if (!evaluation_result.dist_fulfilled) {
         // Do not perform index scan in this case.
         return Status::OK();
     }
 
-    auto src_col_idx = _ann_range_search_runtime.src_col_idx;
+    DCHECK_LT(_ann_range_search_runtime.src_col_idx, idx_to_cid.size());
+    const auto src_col_idx = cast_set<int>(_ann_range_search_runtime.src_col_idx);
+    const auto src_col_key = cast_set<ColumnId>(_ann_range_search_runtime.src_col_idx);
     auto slot_ref_map_it = common_expr_to_slotref_map.find(this);
     if (slot_ref_map_it == common_expr_to_slotref_map.end()) {
         return Status::OK();
     }
     auto& slot_ref_map = slot_ref_map_it->second;
-    ColumnId cid = idx_to_cid[src_col_idx];
-    if (slot_ref_map.find(cid) == slot_ref_map.end()) {
+    auto slot_ref_it = slot_ref_map.find(src_col_key);
+    if (slot_ref_it == slot_ref_map.end()) {
         return Status::OK();
     }
-    const VExpr* slot_ref_expr_addr = slot_ref_map.find(cid)->second;
-    _index_context->set_true_for_index_status(slot_ref_expr_addr, idx_to_cid[cid]);
+    const VExpr* slot_ref_expr_addr = slot_ref_it->second;
+    _index_context->set_true_for_index_status(slot_ref_expr_addr, src_col_idx);
 
     VLOG_DEBUG << fmt::format(
             "Evaluate ann range search for expr {}, src_col_idx {}, cid {}, row_bitmap "
             "cardinality {}",
-            _root->debug_string(), src_col_idx, cid, row_bitmap.cardinality());
+            _root->debug_string(), src_col_idx, idx_to_cid[_ann_range_search_runtime.src_col_idx],
+            row_bitmap.cardinality());
     return Status::OK();
 }
 

be/src/exprs/vexpr_context.h

@@ -389,7 +389,7 @@ class VExprContext {
             const std::unordered_map<VExprContext*, std::unordered_map<ColumnId, VExpr*>>&
                     common_expr_to_slotref_map,
             roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-            bool enable_result_cache);
+            bool enable_result_cache, bool* ann_range_search_executed);
 
     uint64_t get_digest(uint64_t seed) const;
 

be/src/exprs/virtual_slot_ref.cpp

@@ -235,11 +235,9 @@ Status VirtualSlotRef::evaluate_ann_range_search(
         const std::vector<ColumnId>& idx_to_cid,
         const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
         roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-        bool enable_result_cache) {
+        bool enable_result_cache, AnnRangeSearchEvaluationResult& result) {
     return _virtual_column_expr->evaluate_ann_range_search(
             range_search_runtime, cid_to_index_iterators, idx_to_cid, column_iterators, row_bitmap,
-            ann_index_stats, enable_result_cache);
-
-    return Status::OK();
+            ann_index_stats, enable_result_cache, result);
 }
 } // namespace doris

be/src/exprs/virtual_slot_ref.h

@@ -108,7 +108,7 @@ class VirtualSlotRef MOCK_REMOVE(final) : public VExpr {
             const std::vector<ColumnId>& idx_to_cid,
             const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,
             roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,
-            bool enable_result_cache) override;
+            bool enable_result_cache, AnnRangeSearchEvaluationResult& result) override;
 
 #ifdef BE_TEST
     // Test-only setter methods for unit testing

be/src/storage/segment/segment_iterator.cpp

@@ -1366,10 +1366,14 @@ Status SegmentIterator::_apply_index_expr() {
     for (const auto& expr_ctx : _common_expr_ctxs_push_down) {
         segment_v2::AnnIndexStats ann_index_stats;
         size_t origin_rows = _row_bitmap.cardinality();
+        bool ann_range_search_executed = false;
         RETURN_IF_ERROR(expr_ctx->evaluate_ann_range_search(
                 _index_iterators, _schema->column_ids(), _column_iterators,
                 _common_expr_to_slotref_map, _row_bitmap, ann_index_stats,
-                enable_ann_index_result_cache));
+                enable_ann_index_result_cache, &ann_range_search_executed));
+        if (ann_range_search_executed) {
+            _opts.stats->ann_index_range_search_cnt++;
+        }
         _opts.stats->rows_ann_index_range_filtered += (origin_rows - _row_bitmap.cardinality());
         _opts.stats->ann_index_load_ns += ann_index_stats.load_index_costs_ns.value();
         _opts.stats->ann_index_range_search_ns += ann_index_stats.search_costs_ns.value();
@@ -1386,14 +1390,6 @@ Status SegmentIterator::_apply_index_expr() {
         _opts.stats->ann_index_range_cache_hits += ann_index_stats.range_cache_hits.value();
     }
 
-    for (auto it = _common_expr_ctxs_push_down.begin(); it != _common_expr_ctxs_push_down.end();) {
-        if ((*it)->root()->ann_range_search_executedd()) {
-            _opts.stats->ann_index_range_search_cnt++;
-            it = _common_expr_ctxs_push_down.erase(it);
-        } else {
-            ++it;
-        }
-    }
     return Status::OK();
 }