From 67fc664fe0ac1774bd4dcf2ab1aaa488cc7d0545 Mon Sep 17 00:00:00 2001 From: DJ <44884520+Tickdack@users.noreply.github.com> Date: Thu, 19 Mar 2026 11:55:22 -0700 Subject: [PATCH] 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 | 73 ++++++++++++++++++++++++++++++++ faiss/invlists/InvertedLists.h | 40 +++++++++++++++++ faiss/python/swigfaiss.swig | 2 + tests/test_index_composite.py | 46 ++++++++++++++++++++ 4 files changed, 161 insertions(+) diff --git a/faiss/invlists/InvertedLists.cpp b/faiss/invlists/InvertedLists.cpp index 448b96973..07a22fe00 100644 --- a/faiss/invlists/InvertedLists.cpp +++ b/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 diff --git a/faiss/invlists/InvertedLists.h b/faiss/invlists/InvertedLists.h index 43c1ecc0c..d7137408b 100644 --- a/faiss/invlists/InvertedLists.h +++ b/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 diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 033dc8d07..af9b87e98 100644 --- a/faiss/python/swigfaiss.swig +++ b/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); diff --git a/tests/test_index_composite.py b/tests/test_index_composite.py index b6832ceb7..666e7c74f 100644 --- a/tests/test_index_composite.py +++ b/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):