Implement IVFRaBitQ FastScan conversion constructor (#5422)

Summary:
Fix broken IndexIVFRaBitQFastScan(IndexIVFRaBitQ) conversion constructor by packing existing IVF RaBitQ codes into FastScan blocks while preserving aux data.

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

Reviewed By: alibeklfc

Differential Revision: D112829289

Pulled By: bshethmeta

fbshipit-source-id: 4d616a7f31b28678902f972b855599fae2decae3

faiss/IndexIVFRaBitQFastScan.cpp

@@ -65,7 +65,8 @@ IndexIVFRaBitQFastScan::IndexIVFRaBitQFastScan(
             metric == METRIC_L2 || metric == METRIC_INNER_PRODUCT,
             "RaBitQ only supports L2 and Inner Product metrics");
     FAISS_THROW_IF_NOT_MSG(
-            bbs_in % 32 == 0, "Batch size must be multiple of 32");
+            bbs_in > 0 && bbs_in % 32 == 0,
+            "Batch size must be positive and a multiple of 32");
     FAISS_THROW_IF_MSG(quantizer_in == nullptr, "Quantizer cannot be null");
 
     by_residual = true;
@@ -94,10 +95,12 @@ IndexIVFRaBitQFastScan::IndexIVFRaBitQFastScan(
     }
 }
 
-// Constructor that converts an existing IndexIVFRaBitQ to FastScan format
+// Constructor that converts an existing IndexIVFRaBitQ to FastScan format.
+// Like other IVF FastScan conversion constructors, this borrows orig's
+// quantizer and orig_invlists; orig must outlive the converted index.
 IndexIVFRaBitQFastScan::IndexIVFRaBitQFastScan(
         const IndexIVFRaBitQ& orig,
-        int /* bbs */)
+        int bbs_in)
         : IndexIVFFastScan(
                   orig.quantizer,
                   orig.d,
@@ -105,7 +108,93 @@ IndexIVFRaBitQFastScan::IndexIVFRaBitQFastScan(
                   0,
                   orig.metric_type,
                   false),
-          rabitq(orig.rabitq) {}
+          rabitq(orig.rabitq) {
+    FAISS_THROW_IF_NOT_MSG(orig.d > 0, "Dimension must be positive");
+    FAISS_THROW_IF_NOT_MSG(
+            orig.metric_type == METRIC_L2 ||
+                    orig.metric_type == METRIC_INNER_PRODUCT,
+            "RaBitQ only supports L2 and Inner Product metrics");
+    FAISS_THROW_IF_NOT_MSG(
+            bbs_in > 0 && bbs_in % 32 == 0,
+            "Batch size must be positive and a multiple of 32");
+    FAISS_THROW_IF_NOT_MSG(orig.invlists != nullptr, "Source invlists null");
+
+    by_residual = true;
+    qb = orig.qb;
+    centered = false;
+
+    const size_t M_fastscan = (orig.d + 3) / 4;
+    constexpr size_t nbits_fastscan = 4;
+
+    this->bbs = bbs_in;
+    this->fine_quantizer = &rabitq;
+    this->M = M_fastscan;
+    this->nbits = nbits_fastscan;
+    this->ksub = (1 << nbits_fastscan);
+    this->M2 = roundup(M_fastscan, 2);
+
+    const size_t bit_pattern_size = (d + 7) / 8;
+    const size_t storage_size = compute_per_vector_storage_size();
+    this->code_size = bit_pattern_size + storage_size;
+    FAISS_THROW_IF_NOT_MSG(
+            orig.code_size == code_size,
+            "Source IndexIVFRaBitQ code size is incompatible");
+
+    ntotal = orig.ntotal;
+    is_trained = orig.is_trained;
+    nprobe = orig.nprobe;
+
+    replace_invlists(new BlockInvertedLists(nlist, get_CodePacker()), true);
+
+#pragma omp parallel for if (nlist > 100)
+    for (idx_t list_no = 0; list_no < static_cast<idx_t>(nlist); list_no++) {
+        const size_t nb = orig.invlists->list_size(list_no);
+        if (nb == 0) {
+            continue;
+        }
+
+        AlignedTable<uint8_t> flat_codes(nb * code_size);
+        memset(flat_codes.get(), 0, nb * code_size);
+
+        InvertedLists::ScopedCodes orig_codes(orig.invlists, list_no);
+        for (size_t i = 0; i < nb; i++) {
+            const uint8_t* orig_code = orig_codes.get() + i * orig.code_size;
+            uint8_t* fs_code = flat_codes.get() + i * code_size;
+
+            for (size_t j = 0; j < static_cast<size_t>(d); j++) {
+                const size_t orig_byte_idx = j / 8;
+                const size_t orig_bit_offset = j % 8;
+                const bool bit_value =
+                        (orig_code[orig_byte_idx] >> orig_bit_offset) & 1;
+                if (bit_value) {
+                    rabitq_utils::set_bit_fastscan(fs_code, j);
+                }
+            }
+
+            memcpy(fs_code + bit_pattern_size,
+                   orig_code + bit_pattern_size,
+                   storage_size);
+        }
+
+        std::unique_ptr<CodePacker> packer(get_CodePacker());
+        const size_t nb2 = roundup(nb, bbs);
+        AlignedTable<uint8_t> block_codes(nb2 / bbs * packer->block_size);
+        memset(block_codes.get(), 0, block_codes.size());
+
+        for (size_t i = 0; i < nb; i++) {
+            packer->pack_1(
+                    flat_codes.get() + i * code_size, i, block_codes.get());
+        }
+
+        invlists->add_entries(
+                list_no,
+                nb,
+                InvertedLists::ScopedIds(orig.invlists, list_no).get(),
+                block_codes.get());
+    }
+
+    orig_invlists = orig.invlists;
+}
 
 size_t IndexIVFRaBitQFastScan::compute_per_vector_storage_size() const {
     return rabitq_utils::compute_per_vector_storage_size(rabitq.nb_bits, d);

tests/test_rabitq_fastscan.py

@@ -849,6 +849,52 @@ class TestMultiBitRaBitQFastScan(unittest.TestCase):
                     np.testing.assert_array_equal(I_rbq_simd, I_rbq_none)
                     np.testing.assert_array_equal(I_fs_simd, I_fs_none)
 
+    def test_ivf_conversion_constructor_matches_direct_fastscan(self):
+        """Converting IndexIVFRaBitQ must build a usable FastScan index."""
+        d, nlist, nprobe, k = 64, 16, 4, 10
+        ds = datasets.SyntheticDataset(d, 700, 300, 20)
+
+        for metric in [faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT]:
+            for nb_bits in [1, 4]:
+                with self.subTest(metric=metric, nb_bits=nb_bits):
+                    quantizer = faiss.IndexFlat(d, metric)
+                    index_rbq = faiss.IndexIVFRaBitQ(
+                        quantizer, d, nlist, metric, True, nb_bits
+                    )
+                    index_rbq.nprobe = nprobe
+                    index_rbq.train(ds.get_train())
+                    index_rbq.add(ds.get_database())
+
+                    index_converted = faiss.IndexIVFRaBitQFastScan(
+                        index_rbq, 32
+                    )
+
+                    direct_quantizer = faiss.clone_index(index_rbq.quantizer)
+                    index_direct = faiss.IndexIVFRaBitQFastScan(
+                        direct_quantizer, d, nlist, metric, 32, True, nb_bits
+                    )
+                    index_direct.nprobe = nprobe
+                    index_direct.qb = index_rbq.qb
+                    index_direct.centered = index_converted.centered
+                    index_direct.train(ds.get_train())
+                    index_direct.add(ds.get_database())
+
+                    D_converted, I_converted = index_converted.search(
+                        ds.get_queries(), k
+                    )
+                    D_direct, I_direct = index_direct.search(
+                        ds.get_queries(), k
+                    )
+
+                    self.assertEqual(index_converted.ntotal, index_rbq.ntotal)
+                    self.assertEqual(index_converted.bbs, 32)
+                    self.assertEqual(index_converted.nprobe, index_rbq.nprobe)
+                    self.assertTrue(index_converted.by_residual)
+                    np.testing.assert_array_equal(I_converted, I_direct)
+                    np.testing.assert_allclose(
+                        D_converted, D_direct, rtol=0, atol=1e-5
+                    )
+
     # ==================== Serialization Tests ====================
 
     def test_serialization(self):