faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/index_read.cpp

@@ -25,9 +25,13 @@
 #include <faiss/invlists/InvertedListsIOHook.h>
 
 #include <faiss/Index2Layer.h>
+#include <faiss/IndexAdditiveQuantizer.h>
+#include <faiss/IndexAdditiveQuantizerFastScan.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVF.h>
+#include <faiss/IndexIVFAdditiveQuantizer.h>
+#include <faiss/IndexIVFAdditiveQuantizerFastScan.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexIVFPQFastScan.h>
@@ -35,12 +39,13 @@
 #include <faiss/IndexIVFSpectralHash.h>
 #include <faiss/IndexLSH.h>
 #include <faiss/IndexLattice.h>
+#include <faiss/IndexNNDescent.h>
 #include <faiss/IndexNSG.h>
 #include <faiss/IndexPQ.h>
 #include <faiss/IndexPQFastScan.h>
 #include <faiss/IndexPreTransform.h>
 #include <faiss/IndexRefine.h>
-#include <faiss/IndexResidual.h>
+#include <faiss/IndexRowwiseMinMax.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/MetaIndexes.h>
 #include <faiss/VectorTransform.h>
@@ -77,16 +82,22 @@ VectorTransform* read_VectorTransform(IOReader* f) {
     VectorTransform* vt = nullptr;
 
     if (h == fourcc("rrot") || h == fourcc("PCAm") || h == fourcc("LTra") ||
-        h == fourcc("PcAm") || h == fourcc("Viqm")) {
+        h == fourcc("PcAm") || h == fourcc("Viqm") || h == fourcc("Pcam")) {
         LinearTransform* lt = nullptr;
         if (h == fourcc("rrot")) {
             lt = new RandomRotationMatrix();
-        } else if (h == fourcc("PCAm") || h == fourcc("PcAm")) {
+        } else if (
+                h == fourcc("PCAm") || h == fourcc("PcAm") ||
+                h == fourcc("Pcam")) {
             PCAMatrix* pca = new PCAMatrix();
             READ1(pca->eigen_power);
+            if (h == fourcc("Pcam")) {
+                READ1(pca->epsilon);
+            }
             READ1(pca->random_rotation);
-            if (h == fourcc("PcAm"))
+            if (h != fourcc("PCAm")) {
                 READ1(pca->balanced_bins);
+            }
             READVECTOR(pca->mean);
             READVECTOR(pca->eigenvalues);
             READVECTOR(pca->PCAMat);
@@ -139,9 +150,10 @@ VectorTransform* read_VectorTransform(IOReader* f) {
         vt = itqt;
     } else {
         FAISS_THROW_FMT(
-                "fourcc %ud (\"%s\") not recognized",
+                "fourcc %ud (\"%s\") not recognized in %s",
                 h,
-                fourcc_inv_printable(h).c_str());
+                fourcc_inv_printable(h).c_str(),
+                f->name.c_str());
     }
     READ1(vt->d_in);
     READ1(vt->d_out);
@@ -239,15 +251,97 @@ static void read_ProductQuantizer(ProductQuantizer* pq, IOReader* f) {
     READVECTOR(pq->centroids);
 }
 
-static void read_ResidualQuantizer(ResidualQuantizer* rq, IOReader* f) {
+static void read_ResidualQuantizer_old(ResidualQuantizer* rq, IOReader* f) {
     READ1(rq->d);
     READ1(rq->M);
     READVECTOR(rq->nbits);
-    rq->set_derived_values();
     READ1(rq->is_trained);
     READ1(rq->train_type);
     READ1(rq->max_beam_size);
     READVECTOR(rq->codebooks);
+    READ1(rq->search_type);
+    READ1(rq->norm_min);
+    READ1(rq->norm_max);
+    rq->set_derived_values();
+}
+
+static void read_AdditiveQuantizer(AdditiveQuantizer* aq, IOReader* f) {
+    READ1(aq->d);
+    READ1(aq->M);
+    READVECTOR(aq->nbits);
+    READ1(aq->is_trained);
+    READVECTOR(aq->codebooks);
+    READ1(aq->search_type);
+    READ1(aq->norm_min);
+    READ1(aq->norm_max);
+    if (aq->search_type == AdditiveQuantizer::ST_norm_cqint8 ||
+        aq->search_type == AdditiveQuantizer::ST_norm_cqint4 ||
+        aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
+        aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) {
+        READXBVECTOR(aq->qnorm.codes);
+    }
+
+    if (aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
+        aq->search_type == AdditiveQuantizer::ST_norm_rq2x4) {
+        READVECTOR(aq->norm_tabs);
+    }
+
+    aq->set_derived_values();
+}
+
+static void read_ResidualQuantizer(ResidualQuantizer* rq, IOReader* f) {
+    read_AdditiveQuantizer(rq, f);
+    READ1(rq->train_type);
+    READ1(rq->max_beam_size);
+    if (!(rq->train_type & ResidualQuantizer::Skip_codebook_tables)) {
+        rq->compute_codebook_tables();
+    }
+}
+
+static void read_LocalSearchQuantizer(LocalSearchQuantizer* lsq, IOReader* f) {
+    read_AdditiveQuantizer(lsq, f);
+    READ1(lsq->K);
+    READ1(lsq->train_iters);
+    READ1(lsq->encode_ils_iters);
+    READ1(lsq->train_ils_iters);
+    READ1(lsq->icm_iters);
+    READ1(lsq->p);
+    READ1(lsq->lambd);
+    READ1(lsq->chunk_size);
+    READ1(lsq->random_seed);
+    READ1(lsq->nperts);
+    READ1(lsq->update_codebooks_with_double);
+}
+
+static void read_ProductAdditiveQuantizer(
+        ProductAdditiveQuantizer* paq,
+        IOReader* f) {
+    read_AdditiveQuantizer(paq, f);
+    READ1(paq->nsplits);
+}
+
+static void read_ProductResidualQuantizer(
+        ProductResidualQuantizer* prq,
+        IOReader* f) {
+    read_ProductAdditiveQuantizer(prq, f);
+
+    for (size_t i = 0; i < prq->nsplits; i++) {
+        auto rq = new ResidualQuantizer();
+        read_ResidualQuantizer(rq, f);
+        prq->quantizers.push_back(rq);
+    }
+}
+
+static void read_ProductLocalSearchQuantizer(
+        ProductLocalSearchQuantizer* plsq,
+        IOReader* f) {
+    read_ProductAdditiveQuantizer(plsq, f);
+
+    for (size_t i = 0; i < plsq->nsplits; i++) {
+        auto lsq = new LocalSearchQuantizer();
+        read_LocalSearchQuantizer(lsq, f);
+        plsq->quantizers.push_back(lsq);
+    }
 }
 
 static void read_ScalarQuantizer(ScalarQuantizer* ivsc, IOReader* f) {
@@ -310,6 +404,21 @@ static void read_NSG(NSG* nsg, IOReader* f) {
     }
 }
 
+static void read_NNDescent(NNDescent* nnd, IOReader* f) {
+    READ1(nnd->ntotal);
+    READ1(nnd->d);
+    READ1(nnd->K);
+    READ1(nnd->S);
+    READ1(nnd->R);
+    READ1(nnd->L);
+    READ1(nnd->iter);
+    READ1(nnd->search_L);
+    READ1(nnd->random_seed);
+    READ1(nnd->has_built);
+
+    READVECTOR(nnd->final_graph);
+}
+
 ProductQuantizer* read_ProductQuantizer(const char* fname) {
     FileIOReader reader(fname);
     return read_ProductQuantizer(&reader);
@@ -393,10 +502,14 @@ static IndexIVFPQ* read_ivfpq(IOReader* f, uint32_t h, int io_flags) {
     }
 
     if (ivpq->is_trained) {
-        // precomputed table not stored. It is cheaper to recompute it
+        // precomputed table not stored. It is cheaper to recompute it.
+        // precompute_table() may be disabled with a flag.
         ivpq->use_precomputed_table = 0;
-        if (ivpq->by_residual)
-            ivpq->precompute_table();
+        if (ivpq->by_residual) {
+            if ((io_flags & IO_FLAG_SKIP_PRECOMPUTE_TABLE) == 0) {
+                ivpq->precompute_table();
+            }
+        }
         if (ivfpqr) {
             read_ProductQuantizer(&ivfpqr->refine_pq, f);
             READVECTOR(ivfpqr->refine_codes);
@@ -422,8 +535,10 @@ Index* read_index(IOReader* f, int io_flags) {
             idxf = new IndexFlat();
         }
         read_index_header(idxf, f);
-        READVECTOR(idxf->xb);
-        FAISS_THROW_IF_NOT(idxf->xb.size() == idxf->ntotal * idxf->d);
+        idxf->code_size = idxf->d * sizeof(float);
+        READXBVECTOR(idxf->codes);
+        FAISS_THROW_IF_NOT(
+                idxf->codes.size() == idxf->ntotal * idxf->code_size);
         // leak!
         idx = idxf;
     } else if (h == fourcc("IxHE") || h == fourcc("IxHe")) {
@@ -433,7 +548,9 @@ Index* read_index(IOReader* f, int io_flags) {
         READ1(idxl->rotate_data);
         READ1(idxl->train_thresholds);
         READVECTOR(idxl->thresholds);
-        READ1(idxl->bytes_per_vec);
+        int code_size_i;
+        READ1(code_size_i);
+        idxl->code_size = code_size_i;
         if (h == fourcc("IxHE")) {
             FAISS_THROW_IF_NOT_FMT(
                     idxl->nbits % 64 == 0,
@@ -441,7 +558,7 @@ Index* read_index(IOReader* f, int io_flags) {
                     "nbits multiple of 64 (got %d)",
                     (int)idxl->nbits);
             // leak
-            idxl->bytes_per_vec *= 8;
+            idxl->code_size *= 8;
         }
         {
             RandomRotationMatrix* rrot = dynamic_cast<RandomRotationMatrix*>(
@@ -454,7 +571,7 @@ Index* read_index(IOReader* f, int io_flags) {
         FAISS_THROW_IF_NOT(
                 idxl->rrot.d_in == idxl->d && idxl->rrot.d_out == idxl->nbits);
         FAISS_THROW_IF_NOT(
-                idxl->codes.size() == idxl->ntotal * idxl->bytes_per_vec);
+                idxl->codes.size() == idxl->ntotal * idxl->code_size);
         idx = idxl;
     } else if (
             h == fourcc("IxPQ") || h == fourcc("IxPo") || h == fourcc("IxPq")) {
@@ -462,6 +579,7 @@ Index* read_index(IOReader* f, int io_flags) {
         IndexPQ* idxp = new IndexPQ();
         read_index_header(idxp, f);
         read_ProductQuantizer(&idxp->pq, f);
+        idxp->code_size = idxp->pq.code_size;
         READVECTOR(idxp->codes);
         if (h == fourcc("IxPo") || h == fourcc("IxPq")) {
             READ1(idxp->search_type);
@@ -475,16 +593,38 @@ Index* read_index(IOReader* f, int io_flags) {
             idxp->metric_type = METRIC_L2;
         }
         idx = idxp;
-    } else if (h == fourcc("IxRQ")) {
-        IndexResidual* idxr = new IndexResidual();
+    } else if (h == fourcc("IxRQ") || h == fourcc("IxRq")) {
+        IndexResidualQuantizer* idxr = new IndexResidualQuantizer();
         read_index_header(idxr, f);
-        read_ResidualQuantizer(&idxr->rq, f);
-        READ1(idxr->search_type);
-        READ1(idxr->norm_min);
-        READ1(idxr->norm_max);
+        if (h == fourcc("IxRQ")) {
+            read_ResidualQuantizer_old(&idxr->rq, f);
+        } else {
+            read_ResidualQuantizer(&idxr->rq, f);
+        }
+        READ1(idxr->code_size);
+        READVECTOR(idxr->codes);
+        idx = idxr;
+    } else if (h == fourcc("IxLS")) {
+        auto idxr = new IndexLocalSearchQuantizer();
+        read_index_header(idxr, f);
+        read_LocalSearchQuantizer(&idxr->lsq, f);
         READ1(idxr->code_size);
         READVECTOR(idxr->codes);
         idx = idxr;
+    } else if (h == fourcc("IxPR")) {
+        auto idxpr = new IndexProductResidualQuantizer();
+        read_index_header(idxpr, f);
+        read_ProductResidualQuantizer(&idxpr->prq, f);
+        READ1(idxpr->code_size);
+        READVECTOR(idxpr->codes);
+        idx = idxpr;
+    } else if (h == fourcc("IxPL")) {
+        auto idxpl = new IndexProductLocalSearchQuantizer();
+        read_index_header(idxpl, f);
+        read_ProductLocalSearchQuantizer(&idxpl->plsq, f);
+        READ1(idxpl->code_size);
+        READVECTOR(idxpl->codes);
+        idx = idxpl;
     } else if (h == fourcc("ImRQ")) {
         ResidualCoarseQuantizer* idxr = new ResidualCoarseQuantizer();
         read_index_header(idxr, f);
@@ -492,6 +632,103 @@ Index* read_index(IOReader* f, int io_flags) {
         READ1(idxr->beam_factor);
         idxr->set_beam_factor(idxr->beam_factor);
         idx = idxr;
+    } else if (
+            h == fourcc("ILfs") || h == fourcc("IRfs") || h == fourcc("IPRf") ||
+            h == fourcc("IPLf")) {
+        bool is_LSQ = h == fourcc("ILfs");
+        bool is_RQ = h == fourcc("IRfs");
+        bool is_PLSQ = h == fourcc("IPLf");
+
+        IndexAdditiveQuantizerFastScan* idxaqfs;
+        if (is_LSQ) {
+            idxaqfs = new IndexLocalSearchQuantizerFastScan();
+        } else if (is_RQ) {
+            idxaqfs = new IndexResidualQuantizerFastScan();
+        } else if (is_PLSQ) {
+            idxaqfs = new IndexProductLocalSearchQuantizerFastScan();
+        } else {
+            idxaqfs = new IndexProductResidualQuantizerFastScan();
+        }
+        read_index_header(idxaqfs, f);
+
+        if (is_LSQ) {
+            read_LocalSearchQuantizer((LocalSearchQuantizer*)idxaqfs->aq, f);
+        } else if (is_RQ) {
+            read_ResidualQuantizer((ResidualQuantizer*)idxaqfs->aq, f);
+        } else if (is_PLSQ) {
+            read_ProductLocalSearchQuantizer(
+                    (ProductLocalSearchQuantizer*)idxaqfs->aq, f);
+        } else {
+            read_ProductResidualQuantizer(
+                    (ProductResidualQuantizer*)idxaqfs->aq, f);
+        }
+
+        READ1(idxaqfs->implem);
+        READ1(idxaqfs->bbs);
+        READ1(idxaqfs->qbs);
+
+        READ1(idxaqfs->M);
+        READ1(idxaqfs->nbits);
+        READ1(idxaqfs->ksub);
+        READ1(idxaqfs->code_size);
+        READ1(idxaqfs->ntotal2);
+        READ1(idxaqfs->M2);
+
+        READ1(idxaqfs->rescale_norm);
+        READ1(idxaqfs->norm_scale);
+        READ1(idxaqfs->max_train_points);
+
+        READVECTOR(idxaqfs->codes);
+        idx = idxaqfs;
+    } else if (
+            h == fourcc("IVLf") || h == fourcc("IVRf") || h == fourcc("NPLf") ||
+            h == fourcc("NPRf")) {
+        bool is_LSQ = h == fourcc("IVLf");
+        bool is_RQ = h == fourcc("IVRf");
+        bool is_PLSQ = h == fourcc("NPLf");
+
+        IndexIVFAdditiveQuantizerFastScan* ivaqfs;
+        if (is_LSQ) {
+            ivaqfs = new IndexIVFLocalSearchQuantizerFastScan();
+        } else if (is_RQ) {
+            ivaqfs = new IndexIVFResidualQuantizerFastScan();
+        } else if (is_PLSQ) {
+            ivaqfs = new IndexIVFProductLocalSearchQuantizerFastScan();
+        } else {
+            ivaqfs = new IndexIVFProductResidualQuantizerFastScan();
+        }
+        read_ivf_header(ivaqfs, f);
+
+        if (is_LSQ) {
+            read_LocalSearchQuantizer((LocalSearchQuantizer*)ivaqfs->aq, f);
+        } else if (is_RQ) {
+            read_ResidualQuantizer((ResidualQuantizer*)ivaqfs->aq, f);
+        } else if (is_PLSQ) {
+            read_ProductLocalSearchQuantizer(
+                    (ProductLocalSearchQuantizer*)ivaqfs->aq, f);
+        } else {
+            read_ProductResidualQuantizer(
+                    (ProductResidualQuantizer*)ivaqfs->aq, f);
+        }
+
+        READ1(ivaqfs->by_residual);
+        READ1(ivaqfs->implem);
+        READ1(ivaqfs->bbs);
+        READ1(ivaqfs->qbs);
+
+        READ1(ivaqfs->M);
+        READ1(ivaqfs->nbits);
+        READ1(ivaqfs->ksub);
+        READ1(ivaqfs->code_size);
+        READ1(ivaqfs->qbs2);
+        READ1(ivaqfs->M2);
+
+        READ1(ivaqfs->rescale_norm);
+        READ1(ivaqfs->norm_scale);
+        READ1(ivaqfs->max_train_points);
+
+        read_InvertedLists(ivaqfs, f, io_flags);
+        idx = ivaqfs;
     } else if (h == fourcc("IvFl") || h == fourcc("IvFL")) { // legacy
         IndexIVFFlat* ivfl = new IndexIVFFlat();
         std::vector<std::vector<Index::idx_t>> ids;
@@ -571,6 +808,39 @@ Index* read_index(IOReader* f, int io_flags) {
         }
         read_InvertedLists(ivsc, f, io_flags);
         idx = ivsc;
+    } else if (
+            h == fourcc("IwLS") || h == fourcc("IwRQ") || h == fourcc("IwPL") ||
+            h == fourcc("IwPR")) {
+        bool is_LSQ = h == fourcc("IwLS");
+        bool is_RQ = h == fourcc("IwRQ");
+        bool is_PLSQ = h == fourcc("IwPL");
+        IndexIVFAdditiveQuantizer* iva;
+        if (is_LSQ) {
+            iva = new IndexIVFLocalSearchQuantizer();
+        } else if (is_RQ) {
+            iva = new IndexIVFResidualQuantizer();
+        } else if (is_PLSQ) {
+            iva = new IndexIVFProductLocalSearchQuantizer();
+        } else {
+            iva = new IndexIVFProductResidualQuantizer();
+        }
+        read_ivf_header(iva, f);
+        READ1(iva->code_size);
+        if (is_LSQ) {
+            read_LocalSearchQuantizer((LocalSearchQuantizer*)iva->aq, f);
+        } else if (is_RQ) {
+            read_ResidualQuantizer((ResidualQuantizer*)iva->aq, f);
+        } else if (is_PLSQ) {
+            read_ProductLocalSearchQuantizer(
+                    (ProductLocalSearchQuantizer*)iva->aq, f);
+        } else {
+            read_ProductResidualQuantizer(
+                    (ProductResidualQuantizer*)iva->aq, f);
+        }
+        READ1(iva->by_residual);
+        READ1(iva->use_precomputed_table);
+        read_InvertedLists(iva, f, io_flags);
+        idx = iva;
     } else if (h == fourcc("IwSh")) {
         IndexIVFSpectralHash* ivsp = new IndexIVFSpectralHash();
         read_ivf_header(ivsp, f);
@@ -668,8 +938,15 @@ Index* read_index(IOReader* f, int io_flags) {
             dynamic_cast<IndexPQ*>(idxhnsw->storage)->pq.compute_sdc_table();
         }
         idx = idxhnsw;
-    } else if (h == fourcc("INSf")) {
-        IndexNSG* idxnsg = new IndexNSGFlat();
+    } else if (
+            h == fourcc("INSf") || h == fourcc("INSp") || h == fourcc("INSs")) {
+        IndexNSG* idxnsg;
+        if (h == fourcc("INSf"))
+            idxnsg = new IndexNSGFlat();
+        if (h == fourcc("INSp"))
+            idxnsg = new IndexNSGPQ();
+        if (h == fourcc("INSs"))
+            idxnsg = new IndexNSGSQ();
         read_index_header(idxnsg, f);
         READ1(idxnsg->GK);
         READ1(idxnsg->build_type);
@@ -681,6 +958,13 @@ Index* read_index(IOReader* f, int io_flags) {
         idxnsg->storage = read_index(f, io_flags);
         idxnsg->own_fields = true;
         idx = idxnsg;
+    } else if (h == fourcc("INNf")) {
+        IndexNNDescent* idxnnd = new IndexNNDescentFlat();
+        read_index_header(idxnnd, f);
+        read_NNDescent(&idxnnd->nndescent, f);
+        idxnnd->storage = read_index(f, io_flags);
+        idxnnd->own_fields = true;
+        idx = idxnnd;
     } else if (h == fourcc("IPfs")) {
         IndexPQFastScan* idxpqfs = new IndexPQFastScan();
         read_index_header(idxpqfs, f);
@@ -691,6 +975,13 @@ Index* read_index(IOReader* f, int io_flags) {
         READ1(idxpqfs->ntotal2);
         READ1(idxpqfs->M2);
         READVECTOR(idxpqfs->codes);
+
+        const auto& pq = idxpqfs->pq;
+        idxpqfs->M = pq.M;
+        idxpqfs->nbits = pq.nbits;
+        idxpqfs->ksub = (1 << pq.nbits);
+        idxpqfs->code_size = pq.code_size;
+
         idx = idxpqfs;
 
     } else if (h == fourcc("IwPf")) {
@@ -705,7 +996,30 @@ Index* read_index(IOReader* f, int io_flags) {
         read_ProductQuantizer(&ivpq->pq, f);
         read_InvertedLists(ivpq, f, io_flags);
         ivpq->precompute_table();
+
+        const auto& pq = ivpq->pq;
+        ivpq->M = pq.M;
+        ivpq->nbits = pq.nbits;
+        ivpq->ksub = (1 << pq.nbits);
+        ivpq->code_size = pq.code_size;
+
         idx = ivpq;
+    } else if (h == fourcc("IRMf")) {
+        IndexRowwiseMinMax* imm = new IndexRowwiseMinMax();
+        read_index_header(imm, f);
+
+        imm->index = read_index(f, io_flags);
+        imm->own_fields = true;
+
+        idx = imm;
+    } else if (h == fourcc("IRMh")) {
+        IndexRowwiseMinMaxFP16* imm = new IndexRowwiseMinMaxFP16();
+        read_index_header(imm, f);
+
+        imm->index = read_index(f, io_flags);
+        imm->own_fields = true;
+
+        idx = imm;
     } else {
         FAISS_THROW_FMT(
                 "Index type 0x%08x (\"%s\") not recognized",