fix out-range issue

microsoft · Aug 14, 2024 · 98b25c4 · 98b25c4
1 parent 25d9428
commit 98b25c4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp
@@ -1700,7 +1700,9 @@ void PQFlashIndex<T, LabelT>::cached_beam_search(const T *query1, const uint64_t
     }
 
     // copy k_search values
-    for (uint64_t i = 0; i < k_search; i++)
+    uint64_t ret_count = full_retset.size() > k_search
+        ? k_search : full_retset.size();
+    for (uint64_t i = 0; i < ret_count; i++)
     {
         indices[i] = full_retset[i].id;
         auto key = (uint32_t)indices[i];