fix

Author: yingfeng

src/storage/invertedindex/search/blockmax_wand_iterator_impl.cpp

@@ -82,18 +82,27 @@ BlockMaxWandIterator::BlockMaxWandIterator(std::vector<std::unique_ptr<DocIterat
     UpdateScoreUpperBoundPrefixSums();
 }
 
-// Highly optimized pivot calculation with multiple strategies
+// Highly optimized pivot calculation with multiple strategies and improved safety
 size_t BlockMaxWandIterator::FindPivotOptimized(float threshold) {
     const size_t num_iterators = sorted_iterators_.size();
 
+    // Safety check for empty iterator list
+    if (num_iterators == 0) {
+        return 0;
+    }
+    
     // For very large sets, try fast estimation first
     if (num_iterators > FAST_PIVOT_THRESHOLD) {
         size_t estimated_pivot;
         if (TryFastPivotEstimation(threshold, estimated_pivot)) {
-            return estimated_pivot;
+            // Double-check the estimated pivot for safety
+            if (estimated_pivot < num_iterators) {
+                return estimated_pivot;
+            }
         }
     }
 
+    // Ensure prefix sums are valid
     if (!prefix_sums_valid_) {
         UpdateScoreUpperBoundPrefixSums();
     }
@@ -108,7 +117,9 @@ size_t BlockMaxWandIterator::FindPivotOptimized(float threshold) {
             left = mid + 1;
         }
     }
-    return left;
+    
+    // Safety check - ensure we don't return an out-of-bounds pivot
+    return std::min(left, num_iterators);
 }
 
 void BlockMaxWandIterator::UpdateScoreUpperBoundPrefixSums() {
@@ -160,25 +171,64 @@ void BlockMaxWandIterator::OptimizedPartialSort(size_t limit) {
     }
 }
 
-// Fast pivot estimation for very large keyword sets
+// Aggressively optimized pivot estimation with safety guarantees
 bool BlockMaxWandIterator::TryFastPivotEstimation(float threshold, size_t& estimated_pivot) {
     const size_t num_iterators = sorted_iterators_.size();
-    
-    // Sample-based estimation: check every Nth iterator
-    const size_t sample_step = std::max(1UL, num_iterators / 20); // Sample ~5% of iterators
-    float accumulated_score = 0.0f;
-    
-    for (size_t i = 0; i < num_iterators; i += sample_step) {
-        accumulated_score += sorted_iterators_[i]->BM25ScoreUpperBound();
-        if (accumulated_score > threshold) {
-            // Estimate the actual pivot position
-            estimated_pivot = std::min(i + sample_step, num_iterators - 1);
+    if (num_iterators <= FAST_PIVOT_THRESHOLD) {
+        return false;
+    }
+
+    // Level 1: Top 5 high-score terms check (hot path)
+    constexpr size_t TOP_K = 5;
+    float sum = 0.0f;
+    for (size_t i = 0; i < std::min(TOP_K, num_iterators); ++i) {
+        sum += sorted_iterators_[i]->BM25ScoreUpperBound();
+        if (sum > threshold) {
+            estimated_pivot = i;
             return true;
         }
     }
-    
-    estimated_pivot = num_iterators;
-    return false;
+
+    // Level 2: Strided sampling with early termination
+    const size_t stride = std::max<size_t>(1, num_iterators / 20); // 5% sampling
+    float prev_sum = sum;
+    size_t prev_i = TOP_K;
+
+    for (size_t i = TOP_K; i < num_iterators; i += stride) {
+        prev_sum = sum;
+        prev_i = i;
+        sum += sorted_iterators_[i]->BM25ScoreUpperBound();
+        
+        if (sum > threshold) {
+            // Linear search backward to find exact pivot
+            for (size_t j = prev_i; j < i; ++j) {
+                prev_sum += sorted_iterators_[j]->BM25ScoreUpperBound();
+                if (prev_sum > threshold) {
+                    estimated_pivot = j;
+                    return true;
+                }
+            }
+            estimated_pivot = i;
+            return true;
+        }
+        
+        // Early termination if remaining terms cannot reach threshold
+        float max_remaining = (num_iterators - i) * sorted_iterators_[i]->BM25ScoreUpperBound();
+        if (sum + max_remaining <= threshold) {
+            break;
+        }
+    }
+
+    // Fallback: Use prefix sums if available
+    if (prefix_sums_valid_) {
+        estimated_pivot = std::lower_bound(score_ub_prefix_sums_.begin(), 
+                                          score_ub_prefix_sums_.end(),
+                                          threshold) - score_ub_prefix_sums_.begin();
+        estimated_pivot = std::min(estimated_pivot, num_iterators - 1);
+        return true;
+    }
+
+    return false; // Fall back to standard method
 }
 
 void BlockMaxWandIterator::UpdateScoreThreshold(const float threshold) {