Reuse IVF iterators after empty checks

faiss/IndexIVF.cpp

@@ -563,8 +563,17 @@ void IndexIVF::search_preassigned(
                         key,
                         nlist);
 
-                // don't waste time on empty lists
-                if (invlists->is_empty(key, inverted_list_context)) {
+                // Iterator backends may need to seek or perform remote setup to
+                // create an iterator. Keep the iterator used for the empty
+                // check so that a non-empty list is opened only once.
+                std::unique_ptr<InvertedListsIterator> iterator;
+                if (invlists->use_iterator) {
+                    iterator.reset(
+                            invlists->get_iterator(key, inverted_list_context));
+                    if (!iterator->is_available()) {
+                        return (size_t)0;
+                    }
+                } else if (invlists->is_empty(key, inverted_list_context)) {
                     return (size_t)0;
                 }
 
@@ -573,11 +582,8 @@ void IndexIVF::search_preassigned(
                 nlistv++;
                 if (invlists->use_iterator) {
                     size_t list_size = 0;
-                    std::unique_ptr<InvertedListsIterator> it(
-                            invlists->get_iterator(key, inverted_list_context));
-
                     nheap += scanner->iterate_codes(
-                            it.get(), simi, idxi, k, list_size);
+                            iterator.get(), simi, idxi, k, list_size);
 
                     return list_size;
                 } else {
@@ -897,19 +903,27 @@ void IndexIVF::range_search_preassigned(
                         ik,
                         nlist);
 
-                if (invlists->is_empty(key, inverted_list_context)) {
+                // Reuse the iterator that performed the empty check. The base
+                // InvertedLists::is_empty implementation creates an iterator,
+                // so calling it here and get_iterator below would open every
+                // non-empty list twice.
+                std::unique_ptr<InvertedListsIterator> iterator;
+                if (invlists->use_iterator) {
+                    iterator.reset(
+                            invlists->get_iterator(key, inverted_list_context));
+                    if (!iterator->is_available()) {
+                        return;
+                    }
+                } else if (invlists->is_empty(key, inverted_list_context)) {
                     return;
                 }
 
                 scanner->set_list(key, coarse_dis[i * cur_nprobe + ik]);
                 const size_t scan_cnt0 = qres.stats.scan_cnt;
                 if (invlists->use_iterator) {
                     size_t list_size = 0;
-                    std::unique_ptr<InvertedListsIterator> it(
-                            invlists->get_iterator(key, inverted_list_context));
-
                     scanner->iterate_codes_range(
-                            it.get(), radius, qres, list_size);
+                            iterator.get(), radius, qres, list_size);
                     qres.stats.scan_cnt += list_size;
                 } else {
                     InvertedLists::ScopedCodes scodes(invlists, key);

tests/test_ivf_index.cpp

@@ -44,6 +44,7 @@ class TestContext {
     std::unordered_map<faiss::idx_t, size_t> list_nos;
     faiss::idx_t id = 0;
     std::set<size_t> lists_probed;
+    size_t iterator_creations = 0;
 };
 
 // the iterator that iterates over the codes stored in context object
@@ -103,6 +104,7 @@ class TestInvertedLists : public faiss::InvertedLists {
             const override {
         auto testContext = (TestContext*)context;
         testContext->lists_probed.insert(list_no);
+        testContext->iterator_creations++;
         return new TestInvertedListIterator(list_no, testContext);
     }
 
@@ -238,6 +240,8 @@ TEST(IVF, list_context) {
                 &params);
         EXPECT_EQ(nprobe, context.lists_probed.size())
                 << "should probe nprobe lists";
+        EXPECT_EQ(nprobe, context.iterator_creations)
+                << "each non-empty iterator list should be opened once";
 
         // check the result contains the query vector, the probablity of
         // this fail should be low
@@ -254,6 +258,23 @@ TEST(IVF, list_context) {
                 labels.cend())
                 << "should return the query vector";
     }
+    {
+        constexpr size_t nprobe = 10;
+        faiss::SearchParametersIVF params;
+        params.inverted_list_context = &context;
+        params.nprobe = nprobe;
+
+        context.iterator_creations = 0;
+        faiss::RangeSearchResult result(1);
+        index.range_search(
+                1,
+                query_vector.data(),
+                std::numeric_limits<float>::max(),
+                &result,
+                &params);
+        EXPECT_EQ(nprobe, context.iterator_creations)
+                << "range search should open each iterator list once";
+    }
 }
 
 TEST(IVF, jaccard_search_returns_most_similar_vector) {