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<uint8_t>` 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/impl/zerocopy_io.h>

faiss::ZeroCopyIOReader reader(buffer.data(), buffer.size());
std::unique_ptr<faiss::Index> 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

@@ -2929,7 +2929,7 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
         auto idxq = std::make_unique<IndexRaBitQ>();
         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<Index> 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.

tests/test_zerocopy.cpp

@@ -7,13 +7,16 @@
 
 #include <gtest/gtest.h>
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <memory>
 #include <random>
 #include <vector>
 
 #include <faiss/IndexBinaryFlat.h>
 #include <faiss/IndexFlat.h>
+#include <faiss/IndexRaBitQ.h>
 #include <faiss/impl/io.h>
 #include <faiss/impl/zerocopy_io.h>
 #include <faiss/index_io.h>
@@ -45,6 +48,30 @@ std::vector<uint8_t> make_binary_data(
     return database;
 }
 
+std::unique_ptr<faiss::IndexRaBitQ> 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<faiss::IndexRaBitQ>(d, faiss::METRIC_L2, nb_bits);
+    index->qb = 4;
+
+    const std::vector<float> 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<uint8_t>& 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<float> xq = make_data(nq, d, 789);
+
+    // create reference results
+    std::vector<float> ref_dis_1(k * nq);
+    std::vector<faiss::idx_t> 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<uint8_t> 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(&copy_reader);
+    auto* index1c_rq = dynamic_cast<faiss::IndexRaBitQ*>(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<faiss::IndexRaBitQ*>(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<float> cand_dis_zc(k * nq);
+    std::vector<faiss::idx_t> cand_ids_zc(k * nq);
+    index1zc->search(nq, xq.data(), k, cand_dis_zc.data(), cand_ids_zc.data());
+
+    std::vector<float> cand_dis_c(k * nq);
+    std::vector<faiss::idx_t> 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);
+}