Search in sharded invlists

Summary:
This diff adds a CombinedIndexSharded1T class to combined_index that uses the 30 shards from the Spark reducer.
The metadata is stored in pickle files on manifold.

Differential Revision: D24018824

fbshipit-source-id: be4ff8b38c3d6e1bb907e02b655d0e419b7a6fea

faiss/InvertedLists.cpp

@@ -616,6 +616,77 @@ void MaskedInvertedLists::prefetch_lists (
     il1->prefetch_lists (list1.data(), list1.size());
 }
 
+/*****************************************
+ * MaskedInvertedLists implementation
+ ******************************************/
+
+
+StopWordsInvertedLists::StopWordsInvertedLists (
+        const InvertedLists *il0, size_t maxsize):
+    ReadOnlyInvertedLists (il0->nlist, il0->code_size),
+    il0 (il0), maxsize (maxsize)
+{
+
+}
+
+size_t StopWordsInvertedLists::list_size(size_t list_no) const
+{
+    size_t sz = il0->list_size(list_no);
+    return sz < maxsize ? sz : 0;
+}
+
+const uint8_t * StopWordsInvertedLists::get_codes (size_t list_no) const
+{
+    return il0->list_size(list_no) < maxsize ? il0->get_codes(list_no) : nullptr;
+}
+
+const idx_t * StopWordsInvertedLists::get_ids (size_t list_no) const
+{
+    return il0->list_size(list_no) < maxsize ? il0->get_ids(list_no) : nullptr;
+}
+
+void StopWordsInvertedLists::release_codes (
+      size_t list_no, const uint8_t *codes) const
+{
+    if (il0->list_size (list_no) < maxsize) {
+        il0->release_codes (list_no, codes);
+    }
+}
+
+void StopWordsInvertedLists::release_ids (size_t list_no, const idx_t *ids) const
+{
+    if (il0->list_size (list_no) < maxsize) {
+        il0->release_ids (list_no, ids);
+    }
+}
+
+idx_t StopWordsInvertedLists::get_single_id (size_t list_no, size_t offset) const
+{
+    FAISS_THROW_IF_NOT(il0->list_size (list_no) < maxsize);
+    return il0->get_single_id (list_no, offset);
+}
+
+const uint8_t * StopWordsInvertedLists::get_single_code (
+           size_t list_no, size_t offset) const
+{
+    FAISS_THROW_IF_NOT(il0->list_size (list_no) < maxsize);
+    return il0->get_single_code (list_no, offset);
+}
+
+void StopWordsInvertedLists::prefetch_lists (
+       const idx_t *list_nos, int nlist) const
+{
+    std::vector<idx_t> list0;
+    for (int i = 0; i < nlist; i++) {
+        idx_t list_no = list_nos[i];
+        if (list_no < 0) continue;
+        if (il0->list_size(list_no) < maxsize) {
+            list0.push_back(list_no);
+        }
+    }
+    il0->prefetch_lists (list0.data(), list0.size());
+}
+
 
 
 } // namespace faiss

faiss/InvertedLists.h

@@ -329,6 +329,33 @@ struct MaskedInvertedLists: ReadOnlyInvertedLists {
 
 };
 
+
+/** if the inverted list in il is smaller than maxsize then return it,
+ *  otherwise return an empty invlist */
+struct StopWordsInvertedLists: ReadOnlyInvertedLists {
+
+    const InvertedLists *il0;
+    size_t maxsize;
+
+    StopWordsInvertedLists (const InvertedLists *il, size_t maxsize);
+
+    size_t list_size(size_t list_no) const override;
+    const uint8_t * get_codes (size_t list_no) const override;
+    const idx_t * get_ids (size_t list_no) const override;
+
+    void release_codes (size_t list_no, const uint8_t *codes) const override;
+    void release_ids (size_t list_no, const idx_t *ids) const override;
+
+    idx_t get_single_id (size_t list_no, size_t offset) const override;
+
+    const uint8_t * get_single_code (
+           size_t list_no, size_t offset) const override;
+
+    void prefetch_lists (const idx_t *list_nos, int nlist) const override;
+
+};
+
+
 } // namespace faiss
 
 

tests/test_index_composite.py

@@ -574,6 +574,48 @@ class TestInvlistMeta(unittest.TestCase):
         assert np.all(D == Dref)
         assert np.all(I == Iref)
 
+    def test_stop_words(self):
+        d = 10
+        nb = 1000
+        nq = 1
+        nt = 200
+
+        xt, xb, xq = get_dataset_2(d, nt, nb, nq)
+
+        index = faiss.index_factory(d, "IVF32,Flat")
+        index.nprobe = 4
+        index.train(xt)
+        index.add(xb)
+        Dref, Iref = index.search(xq, 10)
+
+        il = index.invlists
+        maxsz = max(il.list_size(i) for i in range(il.nlist))
+
+        il2 = faiss.StopWordsInvertedLists(il, maxsz + 1)
+        index.own_invlists
+        index.own_invlists = False
+
+        index.replace_invlists(il2, False)
+        D1, I1 = index.search(xq, 10)
+        np.testing.assert_array_equal(Dref, D1)
+        np.testing.assert_array_equal(Iref, I1)
+
+        # cleanup to avoid segfault on exit
+        index.replace_invlists(il, False)
+
+        # voluntarily unbalance one invlist
+        i = int(I1[0, 0])
+        index.add(np.vstack([xb[i]] * (maxsz + 10)))
+
+        # introduce stopwords again
+        index.replace_invlists(il2, False)
+
+        D2, I2 = index.search(xq, 10)
+        self.assertFalse(i in list(I2.ravel()))
+
+        # avoid mem leak
+        index.replace_invlists(il, True)
+