in faiss/impl/index_write.cpp [340:592]
void write_index(const Index* idx, IOWriter* f) {
if (const IndexFlat* idxf = dynamic_cast<const IndexFlat*>(idx)) {
uint32_t h =
fourcc(idxf->metric_type == METRIC_INNER_PRODUCT ? "IxFI"
: idxf->metric_type == METRIC_L2 ? "IxF2"
: "IxFl");
WRITE1(h);
write_index_header(idx, f);
WRITEXBVECTOR(idxf->codes);
} else if (const IndexLSH* idxl = dynamic_cast<const IndexLSH*>(idx)) {
uint32_t h = fourcc("IxHe");
WRITE1(h);
write_index_header(idx, f);
WRITE1(idxl->nbits);
WRITE1(idxl->rotate_data);
WRITE1(idxl->train_thresholds);
WRITEVECTOR(idxl->thresholds);
int code_size_i = idxl->code_size;
WRITE1(code_size_i);
write_VectorTransform(&idxl->rrot, f);
WRITEVECTOR(idxl->codes);
} else if (const IndexPQ* idxp = dynamic_cast<const IndexPQ*>(idx)) {
uint32_t h = fourcc("IxPq");
WRITE1(h);
write_index_header(idx, f);
write_ProductQuantizer(&idxp->pq, f);
WRITEVECTOR(idxp->codes);
// search params -- maybe not useful to store?
WRITE1(idxp->search_type);
WRITE1(idxp->encode_signs);
WRITE1(idxp->polysemous_ht);
} else if (
const IndexResidualQuantizer* idxr =
dynamic_cast<const IndexResidualQuantizer*>(idx)) {
uint32_t h = fourcc("IxRq");
WRITE1(h);
write_index_header(idx, f);
write_ResidualQuantizer(&idxr->rq, f);
WRITE1(idxr->code_size);
WRITEVECTOR(idxr->codes);
} else if (
auto* idxr = dynamic_cast<const IndexLocalSearchQuantizer*>(idx)) {
uint32_t h = fourcc("IxLS");
WRITE1(h);
write_index_header(idx, f);
write_LocalSearchQuantizer(&idxr->lsq, f);
WRITE1(idxr->code_size);
WRITEVECTOR(idxr->codes);
} else if (
const ResidualCoarseQuantizer* idxr =
dynamic_cast<const ResidualCoarseQuantizer*>(idx)) {
uint32_t h = fourcc("ImRQ");
WRITE1(h);
write_index_header(idx, f);
write_ResidualQuantizer(&idxr->rq, f);
WRITE1(idxr->beam_factor);
} else if (
const Index2Layer* idxp = dynamic_cast<const Index2Layer*>(idx)) {
uint32_t h = fourcc("Ix2L");
WRITE1(h);
write_index_header(idx, f);
write_index(idxp->q1.quantizer, f);
WRITE1(idxp->q1.nlist);
WRITE1(idxp->q1.quantizer_trains_alone);
write_ProductQuantizer(&idxp->pq, f);
WRITE1(idxp->code_size_1);
WRITE1(idxp->code_size_2);
WRITE1(idxp->code_size);
WRITEVECTOR(idxp->codes);
} else if (
const IndexScalarQuantizer* idxs =
dynamic_cast<const IndexScalarQuantizer*>(idx)) {
uint32_t h = fourcc("IxSQ");
WRITE1(h);
write_index_header(idx, f);
write_ScalarQuantizer(&idxs->sq, f);
WRITEVECTOR(idxs->codes);
} else if (
const IndexLattice* idxl = dynamic_cast<const IndexLattice*>(idx)) {
uint32_t h = fourcc("IxLa");
WRITE1(h);
WRITE1(idxl->d);
WRITE1(idxl->nsq);
WRITE1(idxl->scale_nbit);
WRITE1(idxl->zn_sphere_codec.r2);
write_index_header(idx, f);
WRITEVECTOR(idxl->trained);
} else if (
const IndexIVFFlatDedup* ivfl =
dynamic_cast<const IndexIVFFlatDedup*>(idx)) {
uint32_t h = fourcc("IwFd");
WRITE1(h);
write_ivf_header(ivfl, f);
{
std::vector<Index::idx_t> tab(2 * ivfl->instances.size());
long i = 0;
for (auto it = ivfl->instances.begin(); it != ivfl->instances.end();
++it) {
tab[i++] = it->first;
tab[i++] = it->second;
}
WRITEVECTOR(tab);
}
write_InvertedLists(ivfl->invlists, f);
} else if (
const IndexIVFFlat* ivfl = dynamic_cast<const IndexIVFFlat*>(idx)) {
uint32_t h = fourcc("IwFl");
WRITE1(h);
write_ivf_header(ivfl, f);
write_InvertedLists(ivfl->invlists, f);
} else if (
const IndexIVFScalarQuantizer* ivsc =
dynamic_cast<const IndexIVFScalarQuantizer*>(idx)) {
uint32_t h = fourcc("IwSq");
WRITE1(h);
write_ivf_header(ivsc, f);
write_ScalarQuantizer(&ivsc->sq, f);
WRITE1(ivsc->code_size);
WRITE1(ivsc->by_residual);
write_InvertedLists(ivsc->invlists, f);
} else if (auto iva = dynamic_cast<const IndexIVFAdditiveQuantizer*>(idx)) {
bool is_LSQ = dynamic_cast<const IndexIVFLocalSearchQuantizer*>(iva);
uint32_t h = fourcc(is_LSQ ? "IwLS" : "IwRQ");
WRITE1(h);
write_ivf_header(iva, f);
WRITE1(iva->code_size);
if (is_LSQ) {
write_LocalSearchQuantizer((LocalSearchQuantizer*)iva->aq, f);
} else {
write_ResidualQuantizer((ResidualQuantizer*)iva->aq, f);
}
WRITE1(iva->by_residual);
WRITE1(iva->use_precomputed_table);
write_InvertedLists(iva->invlists, f);
} else if (
const IndexIVFSpectralHash* ivsp =
dynamic_cast<const IndexIVFSpectralHash*>(idx)) {
uint32_t h = fourcc("IwSh");
WRITE1(h);
write_ivf_header(ivsp, f);
write_VectorTransform(ivsp->vt, f);
WRITE1(ivsp->nbit);
WRITE1(ivsp->period);
WRITE1(ivsp->threshold_type);
WRITEVECTOR(ivsp->trained);
write_InvertedLists(ivsp->invlists, f);
} else if (const IndexIVFPQ* ivpq = dynamic_cast<const IndexIVFPQ*>(idx)) {
const IndexIVFPQR* ivfpqr = dynamic_cast<const IndexIVFPQR*>(idx);
uint32_t h = fourcc(ivfpqr ? "IwQR" : "IwPQ");
WRITE1(h);
write_ivf_header(ivpq, f);
WRITE1(ivpq->by_residual);
WRITE1(ivpq->code_size);
write_ProductQuantizer(&ivpq->pq, f);
write_InvertedLists(ivpq->invlists, f);
if (ivfpqr) {
write_ProductQuantizer(&ivfpqr->refine_pq, f);
WRITEVECTOR(ivfpqr->refine_codes);
WRITE1(ivfpqr->k_factor);
}
} else if (
const IndexPreTransform* ixpt =
dynamic_cast<const IndexPreTransform*>(idx)) {
uint32_t h = fourcc("IxPT");
WRITE1(h);
write_index_header(ixpt, f);
int nt = ixpt->chain.size();
WRITE1(nt);
for (int i = 0; i < nt; i++)
write_VectorTransform(ixpt->chain[i], f);
write_index(ixpt->index, f);
} else if (
const MultiIndexQuantizer* imiq =
dynamic_cast<const MultiIndexQuantizer*>(idx)) {
uint32_t h = fourcc("Imiq");
WRITE1(h);
write_index_header(imiq, f);
write_ProductQuantizer(&imiq->pq, f);
} else if (
const IndexRefine* idxrf = dynamic_cast<const IndexRefine*>(idx)) {
uint32_t h = fourcc("IxRF");
WRITE1(h);
write_index_header(idxrf, f);
write_index(idxrf->base_index, f);
write_index(idxrf->refine_index, f);
WRITE1(idxrf->k_factor);
} else if (
const IndexIDMap* idxmap = dynamic_cast<const IndexIDMap*>(idx)) {
uint32_t h = dynamic_cast<const IndexIDMap2*>(idx) ? fourcc("IxM2")
: fourcc("IxMp");
// no need to store additional info for IndexIDMap2
WRITE1(h);
write_index_header(idxmap, f);
write_index(idxmap->index, f);
WRITEVECTOR(idxmap->id_map);
} else if (const IndexHNSW* idxhnsw = dynamic_cast<const IndexHNSW*>(idx)) {
uint32_t h = dynamic_cast<const IndexHNSWFlat*>(idx) ? fourcc("IHNf")
: dynamic_cast<const IndexHNSWPQ*>(idx) ? fourcc("IHNp")
: dynamic_cast<const IndexHNSWSQ*>(idx) ? fourcc("IHNs")
: dynamic_cast<const IndexHNSW2Level*>(idx) ? fourcc("IHN2")
: 0;
FAISS_THROW_IF_NOT(h != 0);
WRITE1(h);
write_index_header(idxhnsw, f);
write_HNSW(&idxhnsw->hnsw, f);
write_index(idxhnsw->storage, f);
} else if (const IndexNSG* idxnsg = dynamic_cast<const IndexNSG*>(idx)) {
uint32_t h =
dynamic_cast<const IndexNSGFlat*>(idx) ? fourcc("INSf") : 0;
FAISS_THROW_IF_NOT(h != 0);
WRITE1(h);
write_index_header(idxnsg, f);
WRITE1(idxnsg->GK);
WRITE1(idxnsg->build_type);
WRITE1(idxnsg->nndescent_S);
WRITE1(idxnsg->nndescent_R);
WRITE1(idxnsg->nndescent_L);
WRITE1(idxnsg->nndescent_iter);
write_NSG(&idxnsg->nsg, f);
write_index(idxnsg->storage, f);
} else if (
const IndexPQFastScan* idxpqfs =
dynamic_cast<const IndexPQFastScan*>(idx)) {
uint32_t h = fourcc("IPfs");
WRITE1(h);
write_index_header(idxpqfs, f);
write_ProductQuantizer(&idxpqfs->pq, f);
WRITE1(idxpqfs->implem);
WRITE1(idxpqfs->bbs);
WRITE1(idxpqfs->qbs);
WRITE1(idxpqfs->ntotal2);
WRITE1(idxpqfs->M2);
WRITEVECTOR(idxpqfs->codes);
} else if (
const IndexIVFPQFastScan* ivpq =
dynamic_cast<const IndexIVFPQFastScan*>(idx)) {
uint32_t h = fourcc("IwPf");
WRITE1(h);
write_ivf_header(ivpq, f);
WRITE1(ivpq->by_residual);
WRITE1(ivpq->code_size);
WRITE1(ivpq->bbs);
WRITE1(ivpq->M2);
WRITE1(ivpq->implem);
WRITE1(ivpq->qbs2);
write_ProductQuantizer(&ivpq->pq, f);
write_InvertedLists(ivpq->invlists, f);
} else {
FAISS_THROW_MSG("don't know how to serialize this type of index");
}
}