facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Implement serialization for IndexIVFFlatPanorama (#4636)

Browse Source 
Summary:
## Changes
This PR implements `write_index` and `read_index` for `IndexIVFFlatPanorama`.

## Testing
Added a test to `test_panorama` to verify integrity of `search()` results after writing and reading back the index from disk.

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

Reviewed By: subhadeepkaran

Differential Revision: D85450597

Pulled By: mnorris11

fbshipit-source-id: 19e065f223b60f8abe14bb0bb8f67f2c8f339234
This commit is contained in:
Akash Nayar
2025-10-24 15:14:49 -07:00
committed by meta-codesync[bot]
parent 595c8aa8a2
commit 859127cd0b
6 changed files with 107 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/IndexIVFFlatPanorama.cpp
+2
View File
@@ -43,6 +43,8 @@ IndexIVFFlatPanorama::IndexIVFFlatPanorama(
this->own_invlists = own_invlists;
}
IndexIVFFlatPanorama::IndexIVFFlatPanorama() : n_levels(0) {}
namespace {
template <typename VectorDistance, bool use_sel>
faiss/IndexIVFFlatPanorama.h
+5 -3
View File
@@ -36,6 +36,10 @@ namespace faiss {
/// insertion logic. The code responsible for level-oriented storage is in
/// `ArrayInvertedListsPanorama`, which is a struct member of `IndexIVF`.
struct IndexIVFFlatPanorama : IndexIVFFlat {
size_t n_levels;
std::vector<MaybeOwnedVector<float>> cum_sums;
explicit IndexIVFFlatPanorama(
Index* quantizer,
size_t d,
@@ -52,9 +56,7 @@ struct IndexIVFFlatPanorama : IndexIVFFlat {
void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
const override;
private:
const size_t n_levels;
std::vector<MaybeOwnedVector<float>> cum_sums;
IndexIVFFlatPanorama();
};
} // namespace faiss
faiss/clone_index.cpp
+2
View File
@@ -23,6 +23,7 @@
#include <faiss/IndexIVF.h>
#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFFlatPanorama.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/IndexIVFPQFastScan.h>
#include <faiss/IndexIVFPQR.h>
@@ -97,6 +98,7 @@ IndexIVF* Cloner::clone_IndexIVF(const IndexIVF* ivf) {
TRYCLONE(IndexIVFFlatDedup, ivf)
TRYCLONE(IndexIVFFlat, ivf)
TRYCLONE(IndexIVFFlatPanorama, ivf)
TRYCLONE(IndexIVFSpectralHash, ivf)
faiss/impl/index_read.cpp
+36
View File
@@ -29,6 +29,7 @@
#include <faiss/IndexIVFAdditiveQuantizer.h>
#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFFlatPanorama.h>
#include <faiss/IndexIVFIndependentQuantizer.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/IndexIVFPQFastScan.h>
@@ -327,6 +328,34 @@ InvertedLists* read_InvertedLists(IOReader* f, int io_flags) {
"read_InvertedLists:"
" WARN! inverted lists not stored with IVF object\n");
return nullptr;
} else if (h == fourcc("ilpn") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
size_t nlist, code_size, n_levels;
READ1(nlist);
READ1(code_size);
READ1(n_levels);
auto ailp = new ArrayInvertedListsPanorama(nlist, code_size, n_levels);
std::vector<size_t> sizes(nlist);
read_ArrayInvertedLists_sizes(f, sizes);
for (size_t i = 0; i < nlist; i++) {
ailp->ids[i].resize(sizes[i]);
size_t num_elems =
((sizes[i] + ArrayInvertedListsPanorama::kBatchSize - 1) /
ArrayInvertedListsPanorama::kBatchSize) *
ArrayInvertedListsPanorama::kBatchSize;
ailp->codes[i].resize(num_elems * code_size);
ailp->cum_sums[i].resize(num_elems * (n_levels + 1));
}
for (size_t i = 0; i < nlist; i++) {
size_t n = sizes[i];
if (n > 0) {
read_vector_with_known_size(
ailp->codes[i], f, ailp->codes[i].size());
read_vector_with_known_size(ailp->ids[i], f, n);
read_vector_with_known_size(
ailp->cum_sums[i], f, ailp->cum_sums[i].size());
}
}
return ailp;
} else if (h == fourcc("ilar") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
auto ails = new ArrayInvertedLists(0, 0);
READ1(ails->nlist);
@@ -929,6 +958,13 @@ Index* read_index(IOReader* f, int io_flags) {
}
read_InvertedLists(ivfl, f, io_flags);
idx = ivfl;
} else if (h == fourcc("IwPn")) {
IndexIVFFlatPanorama* ivfp = new IndexIVFFlatPanorama();
read_ivf_header(ivfp, f);
ivfp->code_size = ivfp->d * sizeof(float);
READ1(ivfp->n_levels);
read_InvertedLists(ivfp, f, io_flags);
idx = ivfp;
} else if (h == fourcc("IwFl")) {
IndexIVFFlat* ivfl = new IndexIVFFlat();
read_ivf_header(ivfl, f);
faiss/impl/index_write.cpp
+36
View File
@@ -27,6 +27,7 @@
#include <faiss/IndexIVFAdditiveQuantizer.h>
#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFFlatPanorama.h>
#include <faiss/IndexIVFIndependentQuantizer.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/IndexIVFPQFastScan.h>
@@ -247,6 +248,33 @@ void write_InvertedLists(const InvertedLists* ils, IOWriter* f) {
if (ils == nullptr) {
uint32_t h = fourcc("il00");
WRITE1(h);
} else if (
const auto& ailp =
dynamic_cast<const ArrayInvertedListsPanorama*>(ils)) {
uint32_t h = fourcc("ilpn");
WRITE1(h);
WRITE1(ailp->nlist);
WRITE1(ailp->code_size);
WRITE1(ailp->n_levels);
uint32_t list_type = fourcc("full");
WRITE1(list_type);
std::vector<size_t> sizes;
sizes.reserve(ailp->nlist);
for (size_t i = 0; i < ailp->nlist; i++) {
sizes.push_back(ailp->ids[i].size());
}
WRITEVECTOR(sizes);
// Write codes, ids, and cum_sums
for (size_t i = 0; i < ailp->nlist; i++) {
size_t n = ailp->ids[i].size();
if (n > 0) {
WRITEANDCHECK(ailp->codes[i].data(), ailp->codes[i].size());
WRITEANDCHECK(ailp->ids[i].data(), n);
WRITEANDCHECK(
ailp->cum_sums[i].data(), ailp->cum_sums[i].size());
}
}
} else if (
const auto& ails = dynamic_cast<const ArrayInvertedLists*>(ils)) {
uint32_t h = fourcc("ilar");
@@ -640,6 +668,14 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
WRITEVECTOR(tab);
}
write_InvertedLists(ivfl->invlists, f);
} else if (
const IndexIVFFlatPanorama* ivfp =
dynamic_cast<const IndexIVFFlatPanorama*>(idx)) {
uint32_t h = fourcc("IwPn");
WRITE1(h);
write_ivf_header(ivfp, f);
WRITE1(ivfp->n_levels);
write_InvertedLists(ivfp->invlists, f);
} else if (
const IndexIVFFlat* ivfl_2 =
dynamic_cast<const IndexIVFFlat*>(idx)) {
tests/test_panorama.py
+26
View File
@@ -15,6 +15,8 @@ Paper: https://www.arxiv.org/pdf/2510.00566
"""
import unittest
import tempfile
import os
import faiss
import numpy as np
@@ -508,3 +510,27 @@ class TestIndexIVFFlatPanorama(unittest.TestCase):
D_reg_2, I_reg_2 = index_regular.search(xq, k)
D_pan_2, I_pan_2 = index_panorama.search(xq, k)
self.assert_search_results_equal(D_reg_2, I_reg_2, D_pan_2, I_pan_2)
def test_serialization(self):
"""Test that writing and reading Panorama indexes preserves search results"""
d, nb, nt, nq, nlist, nlevels, k = 128, 10000, 15000, 100, 128, 8, 20
xt, xb, xq = self.generate_data(d, nt, nb, nq, seed=2024)
indexes = [self.create_panorama(d, nlist, nlevels, xt, xb, nprobe=32)]
for index in indexes:
D_before, I_before = index.search(xq, k)
fd, fname = tempfile.mkstemp()
os.close(fd)
try:
faiss.write_index(index, fname)
index_after = faiss.read_index(fname)
finally:
if os.path.exists(fname):
os.unlink(fname)
D_after, I_after = index_after.search(xq, k)
np.testing.assert_array_equal(I_before, I_after)
np.testing.assert_array_equal(D_before, D_after)