Fix IVF quantizer centroid sharding so IDs are generated (facebookresearch#4197)

Summary:

Ivan and I discussed 2 problems:
1. We may want to try to offload/shard PQ or SQ table data if there is a big enough win (pending)
2. IDs seem to be random after sharding.

This diff solves 2.

Root cause is that we add to quantizer without IDs.

Instead, we wrap in IndexIDMap2 (which provides reconstruction, whereas IndexIDMap does not).

Laser's quantizers are Flat and HNSW, so we can wrap like this.

Differential Revision: D69832788

Author: Michael Norris

faiss/IVFlib.cpp

@@ -9,6 +9,7 @@
 #include <omp.h>
 
 #include <memory>
+#include <numeric>
 
 #include <faiss/IndexAdditiveQuantizer.h>
 #include <faiss/IndexIVFAdditiveQuantizer.h>
@@ -529,20 +530,30 @@ void handle_ivf(
         faiss::IndexIVF* index,
         int64_t shard_count,
         const std::string& filename_template,
-        ShardingFunction* sharding_function) {
+        ShardingFunction* sharding_function,
+        bool generate_ids) {
     std::vector<faiss::IndexIVF*> sharded_indexes(shard_count);
     auto clone = static_cast<faiss::IndexIVF*>(faiss::clone_index(index));
     clone->quantizer->reset();
     for (int64_t i = 0; i < shard_count; i++) {
         sharded_indexes[i] =
                 static_cast<faiss::IndexIVF*>(faiss::clone_index(clone));
+        if (generate_ids) {
+            // Assume the quantizer does not natively support add_with_ids.
+            sharded_indexes[i]->quantizer =
+                    new IndexIDMap2(sharded_indexes[i]->quantizer);
+        }
     }
 
     // assign centroids to each sharded Index based on sharding_function, and
     // add them to the quantizer of each sharded index
     std::vector<std::vector<float>> sharded_centroids(shard_count);
+    std::vector<std::vector<idx_t>> xids(shard_count);
     for (int64_t i = 0; i < index->quantizer->ntotal; i++) {
         int64_t shard_id = (*sharding_function)(i, shard_count);
+        // Since the quantizer does not natively support add_with_ids, we simply
+        // generate them.
+        xids[shard_id].push_back(i);
         float* reconstructed = new float[index->quantizer->d];
         index->quantizer->reconstruct(i, reconstructed);
         sharded_centroids[shard_id].insert(
@@ -552,9 +563,16 @@ void handle_ivf(
         delete[] reconstructed;
     }
     for (int64_t i = 0; i < shard_count; i++) {
-        sharded_indexes[i]->quantizer->add(
-                sharded_centroids[i].size() / index->quantizer->d,
-                sharded_centroids[i].data());
+        if (generate_ids) {
+            sharded_indexes[i]->quantizer->add_with_ids(
+                    sharded_centroids[i].size() / index->quantizer->d,
+                    sharded_centroids[i].data(),
+                    xids[i].data());
+        } else {
+            sharded_indexes[i]->quantizer->add(
+                    sharded_centroids[i].size() / index->quantizer->d,
+                    sharded_centroids[i].data());
+        }
     }
 
     for (int64_t i = 0; i < shard_count; i++) {
@@ -572,7 +590,8 @@ void handle_binary_ivf(
         faiss::IndexBinaryIVF* index,
         int64_t shard_count,
         const std::string& filename_template,
-        ShardingFunction* sharding_function) {
+        ShardingFunction* sharding_function,
+        bool generate_ids) {
     std::vector<faiss::IndexBinaryIVF*> sharded_indexes(shard_count);
 
     auto clone = static_cast<faiss::IndexBinaryIVF*>(
@@ -582,14 +601,23 @@ void handle_binary_ivf(
     for (int64_t i = 0; i < shard_count; i++) {
         sharded_indexes[i] = static_cast<faiss::IndexBinaryIVF*>(
                 faiss::clone_binary_index(clone));
+        if (generate_ids) {
+            // Assume the quantizer does not natively support add_with_ids.
+            sharded_indexes[i]->quantizer =
+                    new IndexBinaryIDMap2(sharded_indexes[i]->quantizer);
+        }
     }
 
     // assign centroids to each sharded Index based on sharding_function, and
     // add them to the quantizer of each sharded index
     int64_t reconstruction_size = index->quantizer->d / 8;
     std::vector<std::vector<uint8_t>> sharded_centroids(shard_count);
+    std::vector<std::vector<idx_t>> xids(shard_count);
     for (int64_t i = 0; i < index->quantizer->ntotal; i++) {
         int64_t shard_id = (*sharding_function)(i, shard_count);
+        // Since the quantizer does not natively support add_with_ids, we simply
+        // generate them.
+        xids[shard_id].push_back(i);
         uint8_t* reconstructed = new uint8_t[reconstruction_size];
         index->quantizer->reconstruct(i, reconstructed);
         sharded_centroids[shard_id].insert(
@@ -599,9 +627,16 @@ void handle_binary_ivf(
         delete[] reconstructed;
     }
     for (int64_t i = 0; i < shard_count; i++) {
-        sharded_indexes[i]->quantizer->add(
-                sharded_centroids[i].size() / reconstruction_size,
-                sharded_centroids[i].data());
+        if (generate_ids) {
+            sharded_indexes[i]->quantizer->add_with_ids(
+                    sharded_centroids[i].size() / reconstruction_size,
+                    sharded_centroids[i].data(),
+                    xids[i].data());
+        } else {
+            sharded_indexes[i]->quantizer->add(
+                    sharded_centroids[i].size() / reconstruction_size,
+                    sharded_centroids[i].data());
+        }
     }
 
     for (int64_t i = 0; i < shard_count; i++) {
@@ -620,7 +655,8 @@ void sharding_helper(
         IndexType* index,
         int64_t shard_count,
         const std::string& filename_template,
-        ShardingFunction* sharding_function) {
+        ShardingFunction* sharding_function,
+        bool generate_ids) {
     FAISS_THROW_IF_MSG(index->quantizer->ntotal == 0, "No centroids to shard.");
     FAISS_THROW_IF_MSG(
             filename_template.find("%d") == std::string::npos,
@@ -636,30 +672,44 @@ void sharding_helper(
                 dynamic_cast<faiss::IndexIVF*>(index),
                 shard_count,
                 filename_template,
-                sharding_function);
+                sharding_function,
+                generate_ids);
     } else if (typeid(IndexType) == typeid(faiss::IndexBinaryIVF)) {
         handle_binary_ivf(
                 dynamic_cast<faiss::IndexBinaryIVF*>(index),
                 shard_count,
                 filename_template,
-                sharding_function);
+                sharding_function,
+                generate_ids);
     }
 }
 
 void shard_ivf_index_centroids(
         faiss::IndexIVF* index,
         int64_t shard_count,
         const std::string& filename_template,
-        ShardingFunction* sharding_function) {
-    sharding_helper(index, shard_count, filename_template, sharding_function);
+        ShardingFunction* sharding_function,
+        bool generate_ids) {
+    sharding_helper(
+            index,
+            shard_count,
+            filename_template,
+            sharding_function,
+            generate_ids);
 }
 
 void shard_binary_ivf_index_centroids(
         faiss::IndexBinaryIVF* index,
         int64_t shard_count,
         const std::string& filename_template,
-        ShardingFunction* sharding_function) {
-    sharding_helper(index, shard_count, filename_template, sharding_function);
+        ShardingFunction* sharding_function,
+        bool generate_ids) {
+    sharding_helper(
+            index,
+            shard_count,
+            filename_template,
+            sharding_function,
+            generate_ids);
 }
 
 } // namespace ivflib

faiss/IVFlib.h

@@ -191,19 +191,23 @@ struct DefaultShardingFunction : ShardingFunction {
  * @param filename_template Template for shard filenames.
  * @param sharding_function The function to shard by. The default is ith vector
  *                          mod shard_count.
+ * @param generate_ids      Generates ids using IndexIDMap2. If false, ids
+ *                          will not match to original quantizer.
  * @return                  The number of shards written.
  */
 void shard_ivf_index_centroids(
         IndexIVF* index,
         int64_t shard_count = 20,
         const std::string& filename_template = "shard.%d.index",
-        ShardingFunction* sharding_function = nullptr);
+        ShardingFunction* sharding_function = nullptr,
+        bool generate_ids = true);
 
 void shard_binary_ivf_index_centroids(
         faiss::IndexBinaryIVF* index,
         int64_t shard_count = 20,
         const std::string& filename_template = "shard.%d.index",
-        ShardingFunction* sharding_function = nullptr);
+        ShardingFunction* sharding_function = nullptr,
+        bool generate_ids = true);
 
 } // namespace ivflib
 } // namespace faiss

tests/test_ivflib.py

@@ -199,7 +199,7 @@ def default_sharding_function(self, i, shard_count):
         return i % shard_count
 
     def verify_sharded_ivf_indexes(
-            self, template, xb, shard_count, sharding_function):
+            self, template, xb, shard_count, sharding_function, generate_ids=True):
         sharded_indexes_counters = [0] * shard_count
         sharded_indexes = []
         for i in range(shard_count):
@@ -208,15 +208,21 @@ def verify_sharded_ivf_indexes(
             else:
                 index = faiss.read_index(template % i)
             sharded_indexes.append(index)
+
         # Reconstruct and verify each centroid
-        nb = len(xb)
-        for i in range(nb):
-            shard_id = sharding_function(i, shard_count)
-            reconstructed = sharded_indexes[shard_id].quantizer.reconstruct(
-                sharded_indexes_counters[shard_id])
-            sharded_indexes_counters[shard_id] += 1
-            print(f"reconstructed: {reconstructed}   xb[i]: {xb[i]}")
-            np.testing.assert_array_equal(reconstructed, xb[i])
+        if generate_ids:
+            for i in range(len(xb)):
+                shard_id = sharding_function(i, shard_count)
+                reconstructed = sharded_indexes[shard_id].quantizer.reconstruct(i)
+                np.testing.assert_array_equal(reconstructed, xb[i])
+        else:
+            for i in range(len(xb)):
+                shard_id = sharding_function(i, shard_count)
+                reconstructed = sharded_indexes[shard_id].quantizer.reconstruct(
+                    sharded_indexes_counters[shard_id])
+                sharded_indexes_counters[shard_id] += 1
+                np.testing.assert_array_equal(reconstructed, xb[i])
+
         # Clean up
         for i in range(shard_count):
             os.remove(template % i)
@@ -264,7 +270,8 @@ def test_save_index_shards_by_centroids_flat_quantizer_custom_sharding(
             index,
             shard_count,
             template,
-            self.custom_sharding_function
+            self.custom_sharding_function,
+            True
         )
         self.verify_sharded_ivf_indexes(
             template, xb, shard_count, self.custom_sharding_function)
@@ -282,7 +289,8 @@ def test_save_index_shards_by_centroids_hnsw_quantizer(self):
             index,
             shard_count,
             template,
-            None
+            None,
+            True
         )
         self.verify_sharded_ivf_indexes(
             template, xb, shard_count, self.default_sharding_function)
@@ -320,3 +328,40 @@ def test_save_index_shards_by_centroids_binary_hnsw_quantizer(self):
         )
         self.verify_sharded_ivf_indexes(
             template, xb, shard_count, self.default_sharding_function)
+
+    def test_save_index_shards_without_id_generation(self):
+        xb = np.random.randint(256, size=(self.nb, int(self.d / 8))).astype('uint8')
+        quantizer = faiss.IndexBinaryHNSW(self.d, 32)
+        index = faiss.IndexBinaryIVF(quantizer, self.d, self.nlist)
+        shard_count = 5
+
+        index.quantizer.add(xb)
+
+        template = str(random.randint(0, 100000)) + "shard.%d.index"
+        faiss.shard_binary_ivf_index_centroids(
+            index,
+            shard_count,
+            template,
+            None,
+            False
+        )
+        self.verify_sharded_ivf_indexes(
+            template, xb, shard_count, self.default_sharding_function, False)
+
+        xb = np.random.rand(self.nb, self.d).astype('float32')
+        quantizer = faiss.IndexHNSWFlat(self.d, 32)
+        index = faiss.IndexIVFFlat(quantizer, self.d, self.nlist)
+        shard_count = 23
+
+        index.quantizer.add(xb)
+
+        template = str(random.randint(0, 100000)) + "shard.%d.index"
+        faiss.shard_ivf_index_centroids(
+            index,
+            shard_count,
+            template,
+            None,
+            False
+        )
+        self.verify_sharded_ivf_indexes(
+            template, xb, shard_count, self.default_sharding_function, False)