Support for Remove ids from IVFPQFastScan index

faiss/impl/FaissAssert.h

@@ -94,13 +94,15 @@
         }                                              \
     } while (false)
 
-#define FAISS_THROW_IF_NOT_MSG(X, MSG)                       \
+#define FAISS_THROW_IF_MSG(X, MSG)                           \
     do {                                                     \
-        if (!(X)) {                                          \
+        if (X) {                                             \
             FAISS_THROW_FMT("Error: '%s' failed: " MSG, #X); \
         }                                                    \
     } while (false)
 
+#define FAISS_THROW_IF_NOT_MSG(X, MSG) FAISS_THROW_IF_MSG(!(X), MSG)
+
 #define FAISS_THROW_IF_NOT_FMT(X, FMT, ...)                               \
     do {                                                                  \
         if (!(X)) {                                                       \

faiss/invlists/BlockInvertedLists.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/impl/CodePacker.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/IDSelector.h>
 
 #include <faiss/impl/io.h>
 #include <faiss/impl/io_macros.h>
@@ -54,7 +55,9 @@ size_t BlockInvertedLists::add_entries(
     codes[list_no].resize(n_block * block_size);
     if (o % block_size == 0) {
         // copy whole blocks
-        memcpy(&codes[list_no][o * code_size], code, n_block * block_size);
+        memcpy(&codes[list_no][o * packer->code_size],
+               code,
+               n_block * block_size);
     } else {
         FAISS_THROW_IF_NOT_MSG(packer, "missing code packer");
         std::vector<uint8_t> buffer(packer->code_size);
@@ -76,6 +79,29 @@ const uint8_t* BlockInvertedLists::get_codes(size_t list_no) const {
     return codes[list_no].get();
 }
 
+size_t BlockInvertedLists::remove_ids(const IDSelector& sel) {
+    idx_t nremove = 0;
+#pragma omp parallel for
+    for (idx_t i = 0; i < nlist; i++) {
+        std::vector<uint8_t> buffer(packer->code_size);
+        idx_t l = ids[i].size(), j = 0;
+        while (j < l) {
+            if (sel.is_member(ids[i][j])) {
+                l--;
+                ids[i][j] = ids[i][l];
+                packer->unpack_1(codes[i].data(), l, buffer.data());
+                packer->pack_1(buffer.data(), j, codes[i].data());
+            } else {
+                j++;
+            }
+        }
+        resize(i, l);
+        nremove += ids[i].size() - l;
+    }
+
+    return nremove;
+}
+
 const idx_t* BlockInvertedLists::get_ids(size_t list_no) const {
     assert(list_no < nlist);
     return ids[list_no].data();
@@ -102,12 +128,6 @@ void BlockInvertedLists::update_entries(
         const idx_t*,
         const uint8_t*) {
     FAISS_THROW_MSG("not impemented");
-    /*
-    assert (list_no < nlist);
-    assert (n_entry + offset <= ids[list_no].size());
-    memcpy (&ids[list_no][offset], ids_in, sizeof(ids_in[0]) * n_entry);
-    memcpy (&codes[list_no][offset * code_size], codes_in, code_size * n_entry);
-    */
 }
 
 BlockInvertedLists::~BlockInvertedLists() {

faiss/invlists/BlockInvertedLists.h

@@ -15,6 +15,7 @@
 namespace faiss {
 
 struct CodePacker;
+struct IDSelector;
 
 /** Inverted Lists that are organized by blocks.
  *
@@ -47,6 +48,8 @@ struct BlockInvertedLists : InvertedLists {
     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;
+    /// remove ids from the InvertedLists
+    size_t remove_ids(const IDSelector& sel);
 
     // works only on empty BlockInvertedLists
     // the codes should be of size ceil(n_entry / n_per_block) * block_size

faiss/invlists/DirectMap.cpp

@@ -15,6 +15,7 @@
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
 #include <faiss/impl/IDSelector.h>
+#include <faiss/invlists/BlockInvertedLists.h>
 
 namespace faiss {
 
@@ -148,8 +149,12 @@ size_t DirectMap::remove_ids(const IDSelector& sel, InvertedLists* invlists) {
     std::vector<idx_t> toremove(nlist);
 
     size_t nremove = 0;
-
+    BlockInvertedLists* block_invlists =
+            dynamic_cast<BlockInvertedLists*>(invlists);
     if (type == NoMap) {
+        if (block_invlists != nullptr) {
+            return block_invlists->remove_ids(sel);
+        }
         // exhaustive scan of IVF
 #pragma omp parallel for
         for (idx_t i = 0; i < nlist; i++) {
@@ -178,6 +183,9 @@ size_t DirectMap::remove_ids(const IDSelector& sel, InvertedLists* invlists) {
             }
         }
     } else if (type == Hashtable) {
+        FAISS_THROW_IF_MSG(
+                block_invlists,
+                "remove with hashtable is not supported with BlockInvertedLists");
         const IDSelectorArray* sela =
                 dynamic_cast<const IDSelectorArray*>(&sel);
         FAISS_THROW_IF_NOT_MSG(

tests/test_merge_index.py

@@ -246,19 +246,45 @@ def test_merge_IDMap2(self):
 
 class TestRemoveFastScan(unittest.TestCase):
 
-    def do_fast_scan_test(self, factory_key, size1):
+    def do_fast_scan_test(self,
+                          factory_key,
+                          with_ids=False,
+                          direct_map_type=faiss.DirectMap.NoMap):
         ds = SyntheticDataset(110, 1000, 1000, 100)
-        index1 = faiss.index_factory(ds.d, factory_key)
-        index1.train(ds.get_train())
-        index1.reset()
+        index = faiss.index_factory(ds.d, factory_key)
+        index.train(ds.get_train())
+
+        index.reset()
         tokeep = [i % 3 == 0 for i in range(ds.nb)]
-        index1.add(ds.get_database()[tokeep])
-        _, Iref = index1.search(ds.get_queries(), 5)
-        index1.reset()
-        index1.add(ds.get_database())
-        index1.remove_ids(np.where(np.logical_not(tokeep))[0])
-        _, Inew = index1.search(ds.get_queries(), 5)
+        if with_ids:
+            index.add_with_ids(ds.get_database()[tokeep], np.arange(ds.nb)[tokeep])
+            faiss.extract_index_ivf(index).nprobe = 5
+        else:
+            index.add(ds.get_database()[tokeep])
+        _, Iref = index.search(ds.get_queries(), 5)
+
+        index.reset()
+        if with_ids:
+            index.add_with_ids(ds.get_database(), np.arange(ds.nb))
+            index.set_direct_map_type(direct_map_type)
+            faiss.extract_index_ivf(index).nprobe = 5
+        else:
+            index.add(ds.get_database())
+        index.remove_ids(np.where(np.logical_not(tokeep))[0])
+        _, Inew = index.search(ds.get_queries(), 5)
         np.testing.assert_array_equal(Inew, Iref)
 
-    def test_remove(self):
-        self.do_fast_scan_test("PQ5x4fs", 320)
+    def test_remove_PQFastScan(self):
+        # with_ids is not support for this type of index
+        self.do_fast_scan_test("PQ5x4fs", False)
+
+    def test_remove_IVFPQFastScan(self):
+        self.do_fast_scan_test("IVF20,PQ5x4fs", True)
+
+    def test_remove_IVFPQFastScan_2(self):
+        self.assertRaisesRegex(Exception,
+                               ".*not supported.*",
+                               self.do_fast_scan_test,
+                               "IVF20,PQ5x4fs",
+                               True,
+                               faiss.DirectMap.Hashtable)