Fix selector handling for IVF iterator scans

faiss/IndexIVF.cpp

@@ -1498,6 +1498,9 @@ size_t InvertedListScanner::iterate_codes(
     if (!keep_max) {
         for (; it->is_available(); it->next()) {
             auto id_and_codes = it->get_id_and_codes();
+            if (sel && !sel->is_member(id_and_codes.first)) {
+                continue;
+            }
             float dis = distance_to_code(id_and_codes.second);
             if (has_cb) {
                 it->on_distance_computed(id_and_codes.first, dis);
@@ -1514,6 +1517,9 @@ size_t InvertedListScanner::iterate_codes(
     } else {
         for (; it->is_available(); it->next()) {
             auto id_and_codes = it->get_id_and_codes();
+            if (sel && !sel->is_member(id_and_codes.first)) {
+                continue;
+            }
             float dis = distance_to_code(id_and_codes.second);
             if (has_cb) {
                 it->on_distance_computed(id_and_codes.first, dis);
@@ -1561,6 +1567,9 @@ void InvertedListScanner::iterate_codes_range(
     list_size = 0;
     for (; it->is_available(); it->next()) {
         auto id_and_codes = it->get_id_and_codes();
+        if (sel && !sel->is_member(id_and_codes.first)) {
+            continue;
+        }
         float dis = distance_to_code(id_and_codes.second);
         bool keep = !keep_max
                 ? dis < radius

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 {
@@ -254,6 +255,48 @@ TEST(IVF, list_context) {
                 labels.cend())
                 << "should return the query vector";
     }
+    {
+        // Iterator scans must honor the same selector contract as array-backed
+        // scans. Keep only the upper half of the IDs and verify both KNN and
+        // range search.
+        constexpr faiss::idx_t k = 100;
+        constexpr size_t nprobe = 10;
+        faiss::IDSelectorRange selector(nb / 2, nb);
+        faiss::SearchParametersIVF params;
+        params.inverted_list_context = &context;
+        params.nprobe = nprobe;
+        params.sel = &selector;
+
+        context.lists_probed.clear();
+        std::vector<float> distances(k);
+        std::vector<faiss::idx_t> labels(k);
+        index.search(
+                1,
+                query_vector.data(),
+                k,
+                distances.data(),
+                labels.data(),
+                &params);
+        const size_t rejected_knn = std::count_if(
+                labels.cbegin(), labels.cend(), [&](faiss::idx_t id) {
+                    return id != -1 && !selector.is_member(id);
+                });
+        EXPECT_EQ(0, rejected_knn);
+
+        context.lists_probed.clear();
+        faiss::RangeSearchResult result(1);
+        index.range_search(
+                1,
+                query_vector.data(),
+                std::numeric_limits<float>::max(),
+                &result,
+                &params);
+        const size_t rejected_range = std::count_if(
+                result.labels,
+                result.labels + result.lims[1],
+                [&](faiss::idx_t id) { return !selector.is_member(id); });
+        EXPECT_EQ(0, rejected_range);
+    }
 }
 
 TEST(IVF, jaccard_search_returns_most_similar_vector) {