diff --git a/faiss/IndexIVFAdditiveQuantizer.cpp b/faiss/IndexIVFAdditiveQuantizer.cpp index edaff51ec7..f0fde48b0e 100644 --- a/faiss/IndexIVFAdditiveQuantizer.cpp +++ b/faiss/IndexIVFAdditiveQuantizer.cpp @@ -116,6 +116,21 @@ void IndexIVFAdditiveQuantizer::sa_decode( } } +void IndexIVFAdditiveQuantizer::reconstruct_from_offset( + int64_t list_no, + int64_t offset, + float* recons) const { + const uint8_t* code = invlists->get_single_code(list_no, offset); + aq->decode(code, recons, 1); + if (by_residual) { + std::vector centroid(d); + quantizer->reconstruct(list_no, centroid.data()); + for (int i = 0; i < d; ++i) { + recons[i] += centroid[i]; + } + } +} + IndexIVFAdditiveQuantizer::~IndexIVFAdditiveQuantizer() = default; /********************************************* diff --git a/faiss/IndexIVFAdditiveQuantizer.h b/faiss/IndexIVFAdditiveQuantizer.h index d065947d09..ecd5a7bda6 100644 --- a/faiss/IndexIVFAdditiveQuantizer.h +++ b/faiss/IndexIVFAdditiveQuantizer.h @@ -56,6 +56,9 @@ struct IndexIVFAdditiveQuantizer : IndexIVF { void sa_decode(idx_t n, const uint8_t* codes, float* x) const override; + void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons) + const override; + ~IndexIVFAdditiveQuantizer() override; }; diff --git a/tests/test_index.py b/tests/test_index.py index 43db906e47..d266bdbcc0 100644 --- a/tests/test_index.py +++ b/tests/test_index.py @@ -493,6 +493,23 @@ def test_IndexIVFPQ(self): self.run_search_and_reconstruct(index, xb, xq, eps=1.0) + def test_IndexIVFRQ(self): + d = 32 + nb = 1000 + nt = 1500 + nq = 200 + + (xt, xb, xq) = get_dataset(d, nb, nt, nq) + + quantizer = faiss.IndexFlatL2(d) + index = faiss.IndexIVFResidualQuantizer(quantizer, d, 32, 8, 8) + index.cp.min_points_per_centroid = 5 # quiet warning + index.nprobe = 4 + index.train(xt) + index.add(xb) + + self.run_search_and_reconstruct(index, xb, xq, eps=1.0) + def test_MultiIndex(self): d = 32 nb = 1000