Add CappedInvertedLists for per-list size limiting (#4880)

Summary:
Following the suggestion in https://github.com/facebookresearch/faiss/issues/4848, this adds CappedInvertedLists which:

- Caps list_size() to maxsize (partial scan instead of skip)
- Allows writes (passes through to underlying invlists)
- Provides real_list_size() to query actual size

Unlike StopWordsInvertedLists which skips large lists entirely, this caps each list while maintaining partial recall from all clusters.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4880

Reviewed By: mnorris11

Differential Revision: D96039869

Pulled By: limqiying

fbshipit-source-id: 7964c5d14c49148c9a1c7c26668b48a6ebb784a8

faiss/invlists/InvertedLists.cpp

@@ -911,4 +911,77 @@ void StopWordsInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
     il0->prefetch_lists(list0.data(), list0.size());
 }
 
+/*****************************************
+ * CappedInvertedLists implementation
+ ******************************************/
+
+CappedInvertedLists::CappedInvertedLists(InvertedLists* il0, size_t maxsize)
+        : InvertedLists(il0->nlist, il0->code_size),
+          il0(il0),
+          maxsize(maxsize) {}
+
+size_t CappedInvertedLists::list_size(size_t list_no) const {
+    size_t sz = il0->list_size(list_no);
+    return sz < maxsize ? sz : maxsize;
+}
+
+size_t CappedInvertedLists::real_list_size(size_t list_no) const {
+    return il0->list_size(list_no);
+}
+
+const uint8_t* CappedInvertedLists::get_codes(size_t list_no) const {
+    return il0->get_codes(list_no);
+}
+
+const idx_t* CappedInvertedLists::get_ids(size_t list_no) const {
+    return il0->get_ids(list_no);
+}
+
+void CappedInvertedLists::release_codes(size_t list_no, const uint8_t* codes)
+        const {
+    il0->release_codes(list_no, codes);
+}
+
+void CappedInvertedLists::release_ids(size_t list_no, const idx_t* ids) const {
+    il0->release_ids(list_no, ids);
+}
+
+idx_t CappedInvertedLists::get_single_id(size_t list_no, size_t offset) const {
+    FAISS_THROW_IF_NOT(offset < list_size(list_no));
+    return il0->get_single_id(list_no, offset);
+}
+
+const uint8_t* CappedInvertedLists::get_single_code(
+        size_t list_no,
+        size_t offset) const {
+    FAISS_THROW_IF_NOT(offset < list_size(list_no));
+    return il0->get_single_code(list_no, offset);
+}
+
+void CappedInvertedLists::prefetch_lists(const idx_t* list_nos, int nlist)
+        const {
+    il0->prefetch_lists(list_nos, nlist);
+}
+
+size_t CappedInvertedLists::add_entries(
+        size_t list_no,
+        size_t n_entry,
+        const idx_t* ids,
+        const uint8_t* code) {
+    return il0->add_entries(list_no, n_entry, ids, code);
+}
+
+void CappedInvertedLists::update_entries(
+        size_t list_no,
+        size_t offset,
+        size_t n_entry,
+        const idx_t* ids,
+        const uint8_t* code) {
+    il0->update_entries(list_no, offset, n_entry, ids, code);
+}
+
+void CappedInvertedLists::resize(size_t list_no, size_t new_size) {
+    il0->resize(list_no, new_size);
+}
+
 } // namespace faiss

faiss/invlists/InvertedLists.h

@@ -467,6 +467,46 @@ struct StopWordsInvertedLists : ReadOnlyInvertedLists {
     void prefetch_lists(const idx_t* list_nos, int nlist) const override;
 };
 
+/** Cap list sizes to maxsize for searching, while allowing writes.
+ *  Unlike StopWordsInvertedLists which skips large lists entirely,
+ *  this caps each list to maxsize entries (partial scan). */
+struct CappedInvertedLists : InvertedLists {
+    InvertedLists* il0;
+    size_t maxsize;
+
+    CappedInvertedLists(InvertedLists* il, size_t maxsize);
+
+    size_t list_size(size_t list_no) const override;
+    size_t real_list_size(size_t list_no) const;
+
+    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;
+
+    size_t add_entries(
+            size_t list_no,
+            size_t n_entry,
+            const idx_t* ids,
+            const uint8_t* code) override;
+
+    void update_entries(
+            size_t list_no,
+            size_t offset,
+            size_t n_entry,
+            const idx_t* ids,
+            const uint8_t* code) override;
+
+    void resize(size_t list_no, size_t new_size) override;
+};
+
 } // namespace faiss
 
 #endif

faiss/python/swigfaiss.swig

@@ -916,6 +916,8 @@ void gpu_sync_all_devices()
     DOWNCAST (VStackInvertedLists)
     DOWNCAST (HStackInvertedLists)
     DOWNCAST (MaskedInvertedLists)
+    DOWNCAST (StopWordsInvertedLists)
+    DOWNCAST (CappedInvertedLists)
     DOWNCAST (InvertedLists)
     {
         assert(false);

tests/test_index_composite.py

@@ -662,6 +662,52 @@ class TestInvlistMeta(unittest.TestCase):
         # avoid mem leak
         index.replace_invlists(il, True)
 
+    def test_capped_invlists(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))
+
+        # cap at half the max size
+        cap = max(maxsz // 2, 1)
+        il2 = faiss.CappedInvertedLists(il, cap)
+
+        # verify capping works
+        for i in range(il.nlist):
+            orig_sz = il.list_size(i)
+            capped_sz = il2.list_size(i)
+            self.assertEqual(capped_sz, min(orig_sz, cap))
+            self.assertEqual(il2.real_list_size(i), orig_sz)
+
+        index.own_invlists = False
+        index.replace_invlists(il2, False)
+
+        # search with capped lists - should still return valid results
+        D1, I1 = index.search(xq, 10)
+        # results may differ due to capping, but should be valid
+        self.assertTrue(np.all(I1 >= -1))
+
+        # test that writes pass through
+        index.replace_invlists(il, False)
+        orig_ntotal = index.ntotal
+        index.replace_invlists(il2, False)
+        index.add(xb[:10])  # add through capped wrapper
+        self.assertEqual(index.ntotal, orig_ntotal + 10)
+
+        # cleanup
+        index.replace_invlists(il, True)
+
 
 class TestSplitMerge(unittest.TestCase):