set keep_max for InvertedListScanner child classes (facebookresearch#…

…3169) Summary: Pull Request resolved: facebookresearch#3169 InvertedListScanner.keep_max is used in iterate_codes method, this field is not set to true when metric is IP, causing problems when using InvertedListsIterator interface. Reviewed By: algoriddle Differential Revision: D52006595 fbshipit-source-id: 6adbe377c0b9f198c9fd8704869133407ef845e5
tarang-jain · Dec 11, 2023 · 79f558f · 79f558f
1 parent 9f8f238
commit 79f558f
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 1 deletion.
diff --git a/faiss/IndexIVFFlat.cpp b/faiss/IndexIVFFlat.cpp
@@ -130,7 +130,9 @@ struct IVFFlatScanner : InvertedListScanner {
     size_t d;
 
     IVFFlatScanner(size_t d, bool store_pairs, const IDSelector* sel)
-            : InvertedListScanner(store_pairs, sel), d(d) {}
+            : InvertedListScanner(store_pairs, sel), d(d) {
+        keep_max = is_similarity_metric(metric);
+    }
 
     const float* xi;
     void set_query(const float* query) override {

diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp
@@ -1199,6 +1199,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
               precompute_mode(precompute_mode),
               sel(sel) {
         this->store_pairs = store_pairs;
+        this->keep_max = is_similarity_metric(METRIC_TYPE);
     }
 
     void set_query(const float* query) override {

diff --git a/faiss/IndexIVFSpectralHash.cpp b/faiss/IndexIVFSpectralHash.cpp
@@ -224,6 +224,7 @@ struct IVFScanner : InvertedListScanner {
               hc(qcode.data(), index->code_size) {
         this->store_pairs = store_pairs;
         this->code_size = index->code_size;
+        this->keep_max = is_similarity_metric(index->metric_type);
     }
 
     void set_query(const float* query) override {

diff --git a/faiss/impl/ScalarQuantizer.cpp b/faiss/impl/ScalarQuantizer.cpp
@@ -1231,6 +1231,7 @@ struct IVFSQScannerIP : InvertedListScanner {
         this->store_pairs = store_pairs;
         this->sel = sel;
         this->code_size = code_size;
+        this->keep_max = true;
     }
 
     void set_query(const float* query) override {