Defer IVF list scan setup until needed

faiss/IndexIVF.cpp

@@ -17,6 +17,7 @@
 #include <cinttypes>
 #include <cstdio>
 #include <limits>
+#include <optional>
 
 #include <faiss/utils/utils.h>
 
@@ -437,6 +438,10 @@ void IndexIVF::search_preassigned(
     const bool ensure_topk_full = params ? params->ensure_topk_full : false;
 
     IDSelector* sel = params ? params->sel : nullptr;
+    FAISS_THROW_IF_NOT_MSG(
+            !(sel && store_pairs),
+            "selector and store_pairs cannot be combined");
+
     const IDSelectorRange* selr = dynamic_cast<const IDSelectorRange*>(sel);
     if (selr) {
         if (selr->assume_sorted) {
@@ -446,10 +451,6 @@ void IndexIVF::search_preassigned(
         }
     }
 
-    FAISS_THROW_IF_NOT_MSG(
-            !(sel && store_pairs),
-            "selector and store_pairs cannot be combined");
-
     FAISS_THROW_IF_NOT_MSG(
             !invlists->use_iterator ||
                     (cur_max_codes == 0 && store_pairs == false),
@@ -568,10 +569,9 @@ void IndexIVF::search_preassigned(
                     return (size_t)0;
                 }
 
-                scanner->set_list(key, coarse_dis_i);
-
                 nlistv++;
                 if (invlists->use_iterator) {
+                    scanner->set_list(key, coarse_dis_i);
                     size_t list_size = 0;
                     std::unique_ptr<InvertedListsIterator> it(
                             invlists->get_iterator(key, inverted_list_context));
@@ -586,31 +586,30 @@ void IndexIVF::search_preassigned(
                         list_size = static_cast<size_t>(list_size_max);
                     }
 
-                    InvertedLists::ScopedCodes scodes(invlists, key);
-                    const uint8_t* codes = scodes.get();
-
-                    std::unique_ptr<InvertedLists::ScopedIds> sids;
+                    std::optional<InvertedLists::ScopedIds> sids;
                     const idx_t* ids = nullptr;
-
                     if (!store_pairs) {
-                        sids = std::make_unique<InvertedLists::ScopedIds>(
-                                invlists, key);
+                        sids.emplace(invlists, key);
                         ids = sids->get();
                     }
 
+                    size_t jmin = 0;
                     if (selr) { // IDSelectorRange
                         // restrict search to a section of the inverted list
-                        size_t jmin, jmax;
+                        size_t jmax;
                         selr->find_sorted_ids_bounds(
                                 list_size, ids, &jmin, &jmax);
                         list_size = jmax - jmin;
                         if (list_size == 0) {
                             return (size_t)0;
                         }
-                        codes += jmin * code_size;
                         ids += jmin;
                     }
 
+                    scanner->set_list(key, coarse_dis_i);
+                    InvertedLists::ScopedCodes scodes(invlists, key);
+                    const uint8_t* codes = scodes.get() + jmin * code_size;
+
                     size_t old_scan_cnt = 0;
                     size_t old_heap_updates = 0;
                     if (is_similarity_metric(metric_type)) {

tests/test_ivf_index.cpp

@@ -19,6 +19,7 @@
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 #include <faiss/impl/ResultHandler.h>
 
 namespace {
@@ -256,6 +257,49 @@ TEST(IVF, list_context) {
     }
 }
 
+TEST(IVF, sorted_range_selector_rejects_store_pairs) {
+    constexpr size_t d = 4;
+    constexpr size_t nb = 100;
+    constexpr faiss::idx_t k = 50;
+
+    faiss::IndexFlatL2 quantizer(d);
+    const std::vector<float> centroid(d, 0.0f);
+    quantizer.add(1, centroid.data());
+    faiss::IndexIVFFlat index(&quantizer, d, 1);
+    index.is_trained = true;
+
+    std::vector<float> database(nb * d);
+    std::vector<faiss::idx_t> ids(nb);
+    for (size_t i = 0; i < nb; ++i) {
+        ids[i] = static_cast<faiss::idx_t>(1000 + i);
+        for (size_t j = 0; j < d; ++j) {
+            database[i * d + j] = static_cast<float>(i + j);
+        }
+    }
+    index.add_with_ids(nb, database.data(), ids.data());
+
+    faiss::IDSelectorRange selector(1020, 1040, true);
+    faiss::SearchParametersIVF params;
+    params.nprobe = 1;
+    params.sel = &selector;
+    const faiss::idx_t key = 0;
+    const float coarse_distance = 0.0f;
+    std::vector<float> distances(k);
+    std::vector<faiss::idx_t> labels(k);
+    EXPECT_THROW(
+            index.search_preassigned(
+                    1,
+                    centroid.data(),
+                    k,
+                    &key,
+                    &coarse_distance,
+                    distances.data(),
+                    labels.data(),
+                    true,
+                    &params),
+            faiss::FaissException);
+}
+
 TEST(IVF, jaccard_search_returns_most_similar_vector) {
     constexpr int d = 3;
     constexpr int nb = 3;