From 88b6c47cf514d3bee7d1995c09a39234fa89d8f9 Mon Sep 17 00:00:00 2001 From: lyang24 Date: Wed, 29 Jul 2026 11:46:10 -0700 Subject: [PATCH] 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 | 97 ++++++++++++++++++++++++++++++-- tests/test_rabitq_fastscan.py | 46 +++++++++++++++ 2 files changed, 139 insertions(+), 4 deletions(-) diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index 9eebd5158..8b6ba6207 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/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(nlist); list_no++) { + const size_t nb = orig.invlists->list_size(list_no); + if (nb == 0) { + continue; + } + + AlignedTable 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(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 packer(get_CodePacker()); + const size_t nb2 = roundup(nb, bbs); + AlignedTable 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); diff --git a/tests/test_rabitq_fastscan.py b/tests/test_rabitq_fastscan.py index fd16e0bce..b4be13876 100644 --- a/tests/test_rabitq_fastscan.py +++ b/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):