From 4fdbc9dc18adbfef3cd410a2446aa4c757320a8a Mon Sep 17 00:00:00 2001 From: Ge Gao Date: Mon, 14 Sep 2026 13:03:13 -0700 Subject: [PATCH] Support ZeroCopyIOReader for IndexRaBitQ deserialization (#5629) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5629 `IndexRaBitQ` keeps its database in `codes`, a `MaybeOwnedVector` inherited from `IndexFlatCodes`, but both read branches (`Ixrq` for single-bit and `Ixrr` for multi-bit) used the raw `READVECTOR` macro. So `codes` was always an owning heap vector even when the caller passed a view-capable reader, while every other flat-codes index already goes through `read_vector`. This switches both branches to `read_vector`, which dispatches on the reader type in `read_vector_base` and makes `codes` a view of the caller's buffer. Any other reader falls back to `READVECTOR`, so existing callers are unaffected. A zero-copy RaBitQ index now works like this: ``` #include faiss::ZeroCopyIOReader reader(buffer.data(), buffer.size()); std::unique_ptr index_zc(faiss::read_index(&reader)); ``` The buffer must outlive the index, as with any `faiss::ZeroCopyIOReader`. Reviewed By: mnorris11 Differential Revision: D119297402 fbshipit-source-id: 95c5ba050cb3a0636a0899efbe886616a7a8a2fd --- faiss/impl/index_read.cpp | 4 +- tests/test_zerocopy.cpp | 126 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 2 deletions(-) diff --git a/faiss/impl/index_read.cpp b/faiss/impl/index_read.cpp index 2899f76fe..2898529f2 100644 --- a/faiss/impl/index_read.cpp +++ b/faiss/impl/index_read.cpp @@ -2929,7 +2929,7 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { auto idxq = std::make_unique(); read_index_header(*idxq, f); read_RaBitQuantizer(idxq->rabitq, f, idxq->d, false); - READVECTOR(idxq->codes); + read_vector(idxq->codes, f); READVECTOR(idxq->center); READ1(idxq->qb); // qb=0: Not quantized - direct distance computation on given float32s. @@ -2948,7 +2948,7 @@ std::unique_ptr read_index_up(IOReader* f, int io_flags) { read_index_header(*idxq, f); read_RaBitQuantizer( idxq->rabitq, f, idxq->d, true); // Reads nb_bits from file - READVECTOR(idxq->codes); + read_vector(idxq->codes, f); READVECTOR(idxq->center); READ1(idxq->qb); // qb=0: Not quantized - direct distance computation on given float32s. diff --git a/tests/test_zerocopy.cpp b/tests/test_zerocopy.cpp index 2ba0dc781..0df5757bd 100644 --- a/tests/test_zerocopy.cpp +++ b/tests/test_zerocopy.cpp @@ -7,13 +7,16 @@ #include +#include #include #include +#include #include #include #include #include +#include #include #include #include @@ -45,6 +48,30 @@ std::vector make_binary_data( return database; } +std::unique_ptr make_rabitq_index( + const size_t nt, + const size_t d, + const uint8_t nb_bits, + const size_t seed) { + auto index = + std::make_unique(d, faiss::METRIC_L2, nb_bits); + index->qb = 4; + + const std::vector xt = make_data(nt, d, seed); + index->train(nt, xt.data()); + index->add(nt, xt.data()); + return index; +} + +// true if [begin, begin + size) lies inside the buffer +bool points_into( + const uint8_t* begin, + const size_t size, + const std::vector& buffer) { + return begin >= buffer.data() && + begin + size <= buffer.data() + buffer.size(); +} + } // namespace // the logic is the following: @@ -240,3 +267,102 @@ TEST(TestZeroCopy, zerocopy_binary_flatcodes) { ASSERT_EQ(ref_ids_1, cand_ids_3); ASSERT_EQ(ref_dis_1, cand_dis_3); } + +namespace { + +// the logic is the following: +// 1. generate a rabitq index, Index1, and create reference results +// 2. serialize it into a std::vector<> buffer, Buf1 +// 3. deserialize Buf1 twice: with a copying reader into Index1C, and with +// the zero-copy feature into Index1ZC +// 4. ensure that Index1C owns its codes while Index1ZC views Buf1 +// 5. ensure that both act as Index1 +// 6. ensure that Index1ZC follows Buf1 if we overwrite a code byte in +// place, while Index1C does not +void check_rabitq_zerocopy(const uint8_t nb_bits) { + // generate data + const size_t nt = 500; + const size_t nq = 10; + const size_t d = 32; + const size_t k = 10; + + const auto index1 = make_rabitq_index(nt, d, nb_bits, 123); + const std::vector xq = make_data(nq, d, 789); + + // create reference results + std::vector ref_dis_1(k * nq); + std::vector ref_ids_1(k * nq); + index1->search(nq, xq.data(), k, ref_dis_1.data(), ref_ids_1.data()); + + // serialize in a form of a vector + faiss::VectorIOWriter wr1; + faiss::write_index(index1.get(), &wr1); + + // clone a buffer + std::vector buffer = wr1.data; + + // create a copying index, it owns its codes + faiss::VectorIOReader copy_reader; + copy_reader.data = buffer; + auto index1c = faiss::read_index_up(©_reader); + auto* index1c_rq = dynamic_cast(index1c.get()); + ASSERT_NE(index1c_rq, nullptr); + ASSERT_TRUE(index1c_rq->codes.is_owned); + + // create a zero-copy index, its codes view the buffer + faiss::ZeroCopyIOReader reader(buffer.data(), buffer.size()); + auto index1zc = faiss::read_index_up(&reader); + auto* index1zc_rq = dynamic_cast(index1zc.get()); + ASSERT_NE(index1zc_rq, nullptr); + ASSERT_FALSE(index1zc_rq->codes.is_owned); + ASSERT_TRUE(points_into( + index1zc_rq->codes.data(), index1zc_rq->codes.size(), buffer)); + + // match the remaining fields vs index1 + ASSERT_EQ(index1zc_rq->d, index1->d); + ASSERT_EQ(index1zc_rq->ntotal, index1->ntotal); + ASSERT_EQ(index1zc_rq->metric_type, index1->metric_type); + ASSERT_EQ(index1zc_rq->code_size, index1->code_size); + ASSERT_EQ(index1zc_rq->qb, index1->qb); + ASSERT_EQ(index1zc_rq->rabitq.nb_bits, index1->rabitq.nb_bits); + ASSERT_EQ(index1zc_rq->center, index1->center); + ASSERT_EQ(index1zc_rq->codes.size(), index1->codes.size()); + ASSERT_TRUE( + std::equal( + index1zc_rq->codes.data(), + index1zc_rq->codes.data() + index1zc_rq->codes.size(), + index1->codes.data())); + + // perform a search on both + std::vector cand_dis_zc(k * nq); + std::vector cand_ids_zc(k * nq); + index1zc->search(nq, xq.data(), k, cand_dis_zc.data(), cand_ids_zc.data()); + + std::vector cand_dis_c(k * nq); + std::vector cand_ids_c(k * nq); + index1c->search(nq, xq.data(), k, cand_dis_c.data(), cand_ids_c.data()); + + // match vs ref1 + ASSERT_EQ(ref_ids_1, cand_ids_zc); + ASSERT_EQ(ref_dis_1, cand_dis_zc); + ASSERT_EQ(ref_ids_1, cand_ids_c); + ASSERT_EQ(ref_dis_1, cand_dis_c); + + // overwrite a code byte without moving the buffer + const size_t codes_offset = index1zc_rq->codes.data() - buffer.data(); + buffer[codes_offset] ^= 0xFF; + + // the zero-copy index follows the buffer, the copying one does not + ASSERT_EQ(index1zc_rq->codes[0], buffer[codes_offset]); + ASSERT_NE(index1c_rq->codes[0], index1zc_rq->codes[0]); +} + +} // namespace + +TEST(TestZeroCopy, zerocopy_rabitq_single_bit) { + check_rabitq_zerocopy(1); +} + +TEST(TestZeroCopy, zerocopy_rabitq_multi_bit) { + check_rabitq_zerocopy(4); +}