Support for Remove ids from IVFPQFastScan index (#3354)

Summary:

**Change was previously reverted because of build failure as change D55577576 removed the definition of FAISS_THROW_IF_MSG**

**Context**
[Issue 3128](#3128) is an enhancement request to support remove_ids for IVFPQFastScan.

Existing mechanism use direct map and iterate over items in selector and use scopecodes and scopeIds to replace item to be removed. Given that codes are packed, it is hard to return single code how it is packed in CodePackerPQ4. Thus, we need a custom implementation to removed_ids.

**In this diff**,
1. We have added custom implementation of remove_ids from BlockInvertedLists which unpack code as it iterate and repack in new position. DirectMap use this remove_id function in BlockInvertedLists for type NoMap in DirectMap.

2. Also, we are throwing exception for other map type in DirectMap i.e. HashTable

Reviewed By: ramilbakhshyiev

Differential Revision: D55858959

Author: kuarora

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)