18 #include "FaissAssert.h"
20 #include "IndexFlat.h"
21 #include "VectorTransform.h"
25 #include "IndexIVFPQ.h"
26 #include "MetaIndexes.h"
27 #include "IndexIVFScalarQuantizer.h"
55 static uint32_t fourcc (
const char sx[4]) {
56 const unsigned char *x = (
unsigned char*)sx;
57 return x[0] | x[1] << 8 | x[2] << 16 | x[3] << 24;
68 #define WRITEANDCHECK(ptr, n) { \
69 size_t ret = fwrite (ptr, sizeof (* (ptr)), n, f); \
70 FAISS_THROW_IF_NOT_MSG (ret == (n), "write error"); \
73 #define READANDCHECK(ptr, n) { \
74 size_t ret = fread (ptr, sizeof (* (ptr)), n, f); \
75 FAISS_THROW_IF_NOT_MSG (ret == (n), "read error"); \
78 #define WRITE1(x) WRITEANDCHECK(&(x), 1)
79 #define READ1(x) READANDCHECK(&(x), 1)
81 #define WRITEVECTOR(vec) { \
82 size_t size = (vec).size (); \
83 WRITEANDCHECK (&size, 1); \
84 WRITEANDCHECK ((vec).data (), size); \
87 #define READVECTOR(vec) { \
89 READANDCHECK (&size, 1); \
90 FAISS_THROW_IF_NOT (size >= 0 && size < (1L << 40)); \
91 (vec).resize (size); \
92 READANDCHECK ((vec).data (), size); \
104 #define WRITETABPAD16(tab, size_in) { \
105 size_t size = (size_in); \
106 WRITEANDCHECK (&size, 1); \
107 uint8_t padding[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; \
108 int idx = ftell(f) % 16; \
109 padding [idx] = 15 - idx; \
110 WRITEANDCHECK (padding + idx, 16 - idx); \
111 WRITEANDCHECK ((tab), size); \
114 #define READTABPAD16(tab, basetype, expected_size) { \
116 READANDCHECK (&size, 1); \
117 FAISS_THROW_IF_NOT ((expected_size) == size); \
118 uint8_t padding[16], npad; \
120 FAISS_THROW_IF_NOT (npad < 16); \
121 READANDCHECK (padding, npad); \
122 (tab) = new basetype [size]; \
123 READANDCHECK ((tab), size); \
127 #define TABOFFSETPAD16(taboffset, basetype, expected_size) { \
129 READANDCHECK (&size, 1); \
130 FAISS_THROW_IF_NOT ((expected_size) == size); \
131 uint8_t padding[16], npad; \
133 FAISS_THROW_IF_NOT (npad < 16); \
134 READANDCHECK (padding, npad); \
135 taboffset = ftell(f); \
136 fseek (f, sizeof(basetype) * size, SEEK_CUR); \
146 static void write_index_header (
const Index *idx, FILE *f) {
158 void write_VectorTransform (
const VectorTransform *vt, FILE *f) {
159 if (
const LinearTransform * lt =
160 dynamic_cast < const LinearTransform *> (vt)) {
161 if (dynamic_cast<const RandomRotationMatrix *>(lt)) {
162 uint32_t h = fourcc (
"rrot");
164 }
else if (
const PCAMatrix * pca =
165 dynamic_cast<const PCAMatrix *>(lt)) {
166 uint32_t h = fourcc (
"PcAm");
168 WRITE1 (pca->eigen_power);
169 WRITE1 (pca->random_rotation);
170 WRITE1 (pca->balanced_bins);
171 WRITEVECTOR (pca->mean);
172 WRITEVECTOR (pca->eigenvalues);
173 WRITEVECTOR (pca->PCAMat);
176 uint32_t h = fourcc (
"LTra");
179 WRITE1 (lt->have_bias);
182 }
else if (
const RemapDimensionsTransform *rdt =
183 dynamic_cast<const RemapDimensionsTransform *>(vt)) {
184 uint32_t h = fourcc (
"RmDT");
186 WRITEVECTOR (rdt->map);
188 FAISS_THROW_MSG (
"cannot serialize this");
193 WRITE1 (vt->is_trained);
196 static void write_ProductQuantizer (
const ProductQuantizer *pq, FILE *f) {
200 WRITEVECTOR (pq->centroids);
203 static void write_ScalarQuantizer (
const ScalarQuantizer *ivsc, FILE *f) {
204 WRITE1 (ivsc->qtype);
205 WRITE1 (ivsc->rangestat);
206 WRITE1 (ivsc->rangestat_arg);
208 WRITE1 (ivsc->code_size);
209 WRITEVECTOR (ivsc->trained);
212 void write_ProductQuantizer (
const ProductQuantizer*pq,
const char *fname) {
213 FILE *f = fopen (fname,
"w");
214 FAISS_THROW_IF_NOT_FMT (f,
"cannot open %s for writing", fname);
215 ScopeFileCloser closer(f);
216 write_ProductQuantizer (pq, f);
221 static void write_ivf_header (
const IndexIVF * ivf, FILE *f,
222 bool include_ids =
true) {
223 write_index_header (ivf, f);
225 WRITE1 (ivf->nprobe);
226 write_index (ivf->quantizer, f);
228 for (
size_t i = 0; i < ivf->nlist; i++)
229 WRITEVECTOR (ivf->ids[i]);
231 WRITE1 (ivf->maintain_direct_map);
232 WRITEVECTOR (ivf->direct_map);
235 void write_index (
const Index *idx, FILE *f) {
236 if (
const IndexFlat * idxf = dynamic_cast<const IndexFlat *> (idx)) {
237 uint32_t h = fourcc (
238 idxf->metric_type == METRIC_INNER_PRODUCT ?
"IxFI" :
239 idxf->metric_type == METRIC_L2 ?
"IxF2" :
nullptr);
241 write_index_header (idx, f);
242 WRITEVECTOR (idxf->xb);
243 }
else if(
const IndexLSH * idxl = dynamic_cast<const IndexLSH *> (idx)) {
244 uint32_t h = fourcc (
"IxHe");
246 write_index_header (idx, f);
247 WRITE1 (idxl->nbits);
248 WRITE1 (idxl->rotate_data);
249 WRITE1 (idxl->train_thresholds);
250 WRITEVECTOR (idxl->thresholds);
251 WRITE1 (idxl->bytes_per_vec);
252 write_VectorTransform (&idxl->rrot, f);
253 WRITEVECTOR (idxl->codes);
254 }
else if(
const IndexPQ * idxp = dynamic_cast<const IndexPQ *> (idx)) {
255 uint32_t h = fourcc (
"IxPq");
257 write_index_header (idx, f);
258 write_ProductQuantizer (&idxp->pq, f);
259 WRITEVECTOR (idxp->codes);
261 WRITE1 (idxp->search_type);
262 WRITE1 (idxp->encode_signs);
263 WRITE1 (idxp->polysemous_ht);
264 }
else if(
const IndexIVFFlat * ivfl =
265 dynamic_cast<const IndexIVFFlat *> (idx)) {
266 uint32_t h = fourcc (
"IvFl");
268 write_ivf_header (ivfl, f);
269 for(
int i = 0; i < ivfl->nlist; i++)
270 WRITEVECTOR (ivfl->vecs[i]);
271 }
else if(
const IndexIVFScalarQuantizer * ivsc =
272 dynamic_cast<const IndexIVFScalarQuantizer *> (idx)) {
273 uint32_t h = fourcc (
"IvSQ");
275 write_ivf_header (ivsc, f);
276 write_ScalarQuantizer (&ivsc->sq, f);
277 WRITE1 (ivsc->code_size);
278 for(
int i = 0; i < ivsc->nlist; i++)
279 WRITEVECTOR (ivsc->codes[i]);
280 }
else if(
const IndexIVFPQ * ivpq =
281 dynamic_cast<const IndexIVFPQ *> (idx)) {
282 const IndexIVFPQR * ivfpqr =
dynamic_cast<const IndexIVFPQR *
> (idx);
283 const IndexIVFPQCompact * ivfpqc =
284 dynamic_cast<const IndexIVFPQCompact *
> (idx);
285 uint32_t h = fourcc (ivfpqr ?
"IvQR" : ivfpqc ?
"IvPC" :
"IvPQ");
287 write_ivf_header (ivpq, f, !ivfpqc);
288 WRITE1 (ivpq->by_residual);
289 WRITE1 (ivpq->code_size);
290 write_ProductQuantizer (&ivpq->pq, f);
292 for(
int i = 0; i < ivpq->codes.size(); i++)
293 WRITEVECTOR (ivpq->codes[i]);
296 write_ProductQuantizer (&ivfpqr->refine_pq, f);
297 WRITEVECTOR (ivfpqr->refine_codes);
298 WRITE1 (ivfpqr->k_factor);
301 WRITETABPAD16 (ivfpqc->limits, ivfpqc->nlist + 1);
302 WRITETABPAD16 (ivfpqc->compact_ids, ivfpqc->ntotal);
303 WRITETABPAD16 (ivfpqc->compact_codes,
304 ivfpqc->ntotal * ivfpqc->code_size);
306 }
else if(
const IndexPreTransform * ixpt =
307 dynamic_cast<const IndexPreTransform *> (idx)) {
308 uint32_t h = fourcc (
"IxPT");
310 write_index_header (ixpt, f);
311 int nt = ixpt->chain.size();
313 for (
int i = 0; i < nt; i++)
314 write_VectorTransform (ixpt->chain[i], f);
315 write_index (ixpt->index, f);
316 }
else if(
const MultiIndexQuantizer * imiq =
317 dynamic_cast<const MultiIndexQuantizer *> (idx)) {
318 uint32_t h = fourcc (
"Imiq");
320 write_index_header (imiq, f);
321 write_ProductQuantizer (&imiq->pq, f);
322 }
else if(
const IndexRefineFlat * idxrf =
323 dynamic_cast<const IndexRefineFlat *> (idx)) {
324 uint32_t h = fourcc (
"IxRF");
326 write_index_header (idxrf, f);
327 write_index (idxrf->base_index, f);
328 write_index (&idxrf->refine_index, f);
329 WRITE1 (idxrf->k_factor);
330 }
else if(
const IndexIDMap * idxmap =
331 dynamic_cast<const IndexIDMap *> (idx)) {
332 uint32_t h = fourcc (
"IxMp");
334 write_index_header (idxmap, f);
335 write_index (idxmap->index, f);
336 WRITEVECTOR (idxmap->id_map);
338 FAISS_THROW_MSG (
"don't know how to serialize this type of index");
342 void write_index (
const Index *idx,
const char *fname) {
343 FILE *f = fopen (fname,
"w");
344 FAISS_THROW_IF_NOT_FMT (f,
"cannot open %s for writing", fname);
345 ScopeFileCloser closer(f);
346 write_index (idx, f);
349 void write_VectorTransform (
const VectorTransform *vt,
const char *fname) {
350 FILE *f = fopen (fname,
"w");
351 FAISS_THROW_IF_NOT_FMT (f,
"cannot open %s for writing", fname);
352 ScopeFileCloser closer(f);
353 write_VectorTransform (vt, f);
360 static void read_index_header (Index *idx, FILE *f) {
366 READ1 (idx->is_trained);
367 READ1 (idx->metric_type);
368 idx->verbose =
false;
371 VectorTransform* read_VectorTransform (FILE *f) {
374 VectorTransform *vt =
nullptr;
376 if (h == fourcc (
"rrot") || h == fourcc (
"PCAm") ||
377 h == fourcc (
"LTra") || h == fourcc (
"PcAm")) {
378 LinearTransform *lt =
nullptr;
379 if (h == fourcc (
"rrot")) {
380 lt =
new RandomRotationMatrix ();
381 }
else if (h == fourcc (
"PCAm") ||
382 h == fourcc (
"PcAm")) {
383 PCAMatrix * pca =
new PCAMatrix ();
384 READ1 (pca->eigen_power);
385 READ1 (pca->random_rotation);
386 if (h == fourcc (
"PcAm"))
387 READ1 (pca->balanced_bins);
388 READVECTOR (pca->mean);
389 READVECTOR (pca->eigenvalues);
390 READVECTOR (pca->PCAMat);
392 }
else if (h == fourcc (
"LTra")) {
393 lt =
new LinearTransform ();
395 READ1 (lt->have_bias);
399 }
else if (h == fourcc (
"RmDT")) {
400 RemapDimensionsTransform *rdt =
new RemapDimensionsTransform ();
401 READVECTOR (rdt->map);
404 FAISS_THROW_MSG(
"fourcc not recognized");
408 READ1 (vt->is_trained);
412 static void read_ProductQuantizer (ProductQuantizer *pq, FILE *f) {
416 pq->set_derived_values ();
417 READVECTOR (pq->centroids);
420 static void read_ScalarQuantizer (ScalarQuantizer *ivsc, FILE *f) {
422 READ1 (ivsc->rangestat);
423 READ1 (ivsc->rangestat_arg);
425 READ1 (ivsc->code_size);
426 READVECTOR (ivsc->trained);
429 ProductQuantizer * read_ProductQuantizer (
const char*fname) {
430 FILE *f = fopen (fname,
"r");
431 FAISS_THROW_IF_NOT_FMT (f,
"cannot open %s for writing", fname);
432 ScopeFileCloser closer(f);
433 ProductQuantizer *pq =
new ProductQuantizer();
434 ScopeDeleter1<ProductQuantizer> del (pq);
435 read_ProductQuantizer(pq, f);
440 static void read_ivf_header (IndexIVF * ivf, FILE *f,
441 bool include_ids =
true) {
442 read_index_header (ivf, f);
446 ivf->own_fields =
true;
448 ivf->ids.resize (ivf->nlist);
449 for (
size_t i = 0; i < ivf->nlist; i++)
450 READVECTOR (ivf->ids[i]);
452 READ1 (ivf->maintain_direct_map);
453 READVECTOR (ivf->direct_map);
456 static IndexIVFPQ *read_ivfpq (FILE *f, uint32_t h,
bool try_mmap)
459 IndexIVFPQR *ivfpqr =
460 h == fourcc (
"IvQR") ?
new IndexIVFPQR () : nullptr;
461 IndexIVFPQCompact *ivfpqc =
462 h == fourcc (
"IvPC") ?
new IndexIVFPQCompact () : nullptr;
463 IndexIVFPQ * ivpq = ivfpqr ? ivfpqr : ivfpqc ? ivfpqc :
new IndexIVFPQ ();
464 read_ivf_header (ivpq, f, !ivfpqc);
465 READ1 (ivpq->by_residual);
466 READ1 (ivpq->code_size);
467 read_ProductQuantizer (&ivpq->pq, f);
469 ivpq->codes.resize (ivpq->nlist);
470 for (
size_t i = 0; i < ivpq->nlist; i++)
471 READVECTOR (ivpq->codes[i]);
474 ivpq->use_precomputed_table = 0;
475 if (ivpq->by_residual)
476 ivpq->precompute_table ();
478 read_ProductQuantizer (&ivfpqr->refine_pq, f);
479 READVECTOR (ivfpqr->refine_codes);
480 READ1 (ivfpqr->k_factor);
484 READTABPAD16 (ivfpqc->limits, uint32_t, ivfpqc->nlist + 1);
485 READTABPAD16 (ivfpqc->compact_ids, uint32_t, ivfpqc->ntotal);
486 READTABPAD16 (ivfpqc->compact_codes, uint8_t,
487 ivfpqc->ntotal * ivfpqc->code_size);
489 long offset_limits, offset_compact_ids, offset_compact_codes;
490 TABOFFSETPAD16 (offset_limits, uint32_t, ivfpqc->nlist + 1);
491 TABOFFSETPAD16 (offset_compact_ids, uint32_t, ivfpqc->ntotal);
492 TABOFFSETPAD16 (offset_compact_codes, uint8_t,
493 ivfpqc->ntotal * ivfpqc->code_size);
494 ivfpqc->mmap_length = ftell (f);
496 ivfpqc->mmap_buffer = (
char*)mmap (
497 nullptr, ivfpqc->mmap_length,
498 PROT_READ, MAP_SHARED, fileno (f), 0);
499 if (!ivfpqc->mmap_buffer) {
500 perror (
"mmap failed");
505 ivfpqc->limits = (uint32_t*)(ivfpqc->mmap_buffer + offset_limits);
506 ivfpqc->compact_ids = (uint32_t*)(ivfpqc->mmap_buffer +
508 ivfpqc->compact_codes = (uint8_t*)(ivfpqc->mmap_buffer +
509 offset_compact_codes);
515 int read_old_fmt_hack = 0;
518 Index * idx =
nullptr;
521 if (h == fourcc (
"IxFI") || h == fourcc (
"IxF2")) {
523 if (h == fourcc (
"IxFI")) idxf =
new IndexFlatIP ();
525 read_index_header (idxf, f);
526 READVECTOR (idxf->
xb);
527 FAISS_THROW_IF_NOT (idxf->
xb.size() == idxf->
ntotal * idxf->
d);
530 }
else if (h == fourcc(
"IxHE") || h == fourcc(
"IxHe")) {
532 read_index_header (idxl, f);
538 if (h == fourcc(
"IxHE")) {
539 FAISS_THROW_IF_NOT_FMT (idxl->
nbits % 64 == 0,
540 "can only read old format IndexLSH with "
541 "nbits multiple of 64 (got %d)",
548 (read_VectorTransform (f));
549 FAISS_THROW_IF_NOT_MSG(rrot,
"expected a random rotation");
553 READVECTOR (idxl->
codes);
554 FAISS_THROW_IF_NOT (idxl->
rrot.d_in == idxl->
d &&
559 }
else if (h == fourcc (
"IxPQ") || h == fourcc (
"IxPo") ||
560 h == fourcc (
"IxPq")) {
563 read_index_header (idxp, f);
564 read_ProductQuantizer (&idxp->
pq, f);
565 READVECTOR (idxp->
codes);
566 if (h == fourcc (
"IxPo") || h == fourcc (
"IxPq")) {
567 READ1 (idxp->search_type);
568 READ1 (idxp->encode_signs);
574 if (h == fourcc (
"IxPQ") || h == fourcc (
"IxPo")) {
578 }
else if(h == fourcc (
"IvFl")) {
580 read_ivf_header (ivfl, f);
582 for (
size_t i = 0; i < ivfl->
nlist; i++)
583 READVECTOR (ivfl->
vecs[i]);
585 }
else if(h == fourcc (
"IvSQ")) {
587 read_ivf_header (ivsc, f);
589 read_ScalarQuantizer (&ivsc->sq, f);
590 READ1 (ivsc->code_size);
591 for(
int i = 0; i < ivsc->
nlist; i++)
592 READVECTOR (ivsc->
codes[i]);
594 }
else if(h == fourcc (
"IvPQ") || h == fourcc (
"IvQR") ||
595 h == fourcc (
"IvPC")) {
597 idx = read_ivfpq (f, h, try_mmap);
599 }
else if(h == fourcc (
"IxPT")) {
602 read_index_header (ixpt, f);
604 if (read_old_fmt_hack == 2) {
609 for (
int i = 0; i < nt; i++)
610 ixpt->chain.push_back (read_VectorTransform (f));
613 }
else if(h == fourcc (
"Imiq")) {
615 read_index_header (imiq, f);
616 read_ProductQuantizer (&imiq->pq, f);
618 }
else if(h == fourcc (
"IxRF")) {
620 read_index_header (idxrf, f);
628 }
else if(h == fourcc (
"IxMp")) {
630 read_index_header (idxmap, f);
633 READVECTOR (idxmap->
id_map);
636 fprintf (stderr,
"Index type 0x%08x not supported\n", h);
644 Index *
read_index (
const char *fname,
bool try_mmap) {
645 FILE *f = fopen (fname,
"r");
646 FAISS_THROW_IF_NOT_FMT (f,
"cannot open %s for reading:", fname);
652 VectorTransform *read_VectorTransform (
const char *fname) {
653 FILE *f = fopen (fname,
"r");
655 fprintf (stderr,
"cannot open %s for reading:", fname);
659 VectorTransform *vt = read_VectorTransform (f);
670 Index * clone_index (
const Index *index)
673 return cl.clone_Index (index);
678 #define TRYCLONE(classname, obj) \
679 if (const classname *clo = dynamic_cast<const classname *>(obj)) { \
680 return new classname(*clo); \
683 VectorTransform *Cloner::clone_VectorTransform (
const VectorTransform *vt)
685 TRYCLONE (RemapDimensionsTransform, vt)
686 TRYCLONE (OPQMatrix, vt)
687 TRYCLONE (PCAMatrix, vt)
688 TRYCLONE (RandomRotationMatrix, vt)
689 TRYCLONE (LinearTransform, vt)
691 FAISS_THROW_MSG(
"clone not supported for this type of VectorTransform");
696 IndexIVF * Cloner::clone_IndexIVF (
const IndexIVF *ivf)
698 TRYCLONE (IndexIVFPQR, ivf)
699 TRYCLONE (IndexIVFPQ, ivf)
700 TRYCLONE (IndexIVFFlat, ivf)
702 FAISS_THROW_MSG(
"clone not supported for this type of IndexIVF");
707 Index *Cloner::clone_Index (
const Index *index)
709 TRYCLONE (IndexPQ, index)
710 TRYCLONE (IndexLSH, index)
711 TRYCLONE (IndexFlatL2, index)
712 TRYCLONE (IndexFlatIP, index)
713 TRYCLONE (IndexFlat, index)
714 TRYCLONE (MultiIndexQuantizer, index)
715 if (const IndexIVF * ivf = dynamic_cast<const IndexIVF*>(index)) {
716 IndexIVF *res = clone_IndexIVF (ivf);
717 res->own_fields =
true;
718 res->quantizer = clone_Index (ivf->quantizer);
720 }
else if (
const IndexPreTransform * ipt =
721 dynamic_cast<const IndexPreTransform*> (index)) {
722 IndexPreTransform *res =
new IndexPreTransform ();
724 res->index = clone_Index (ipt->index);
725 for (
int i = 0; i < ipt->chain.size(); i++)
726 res->chain.push_back (clone_VectorTransform (ipt->chain[i]));
727 res->own_fields =
true;
730 FAISS_THROW_MSG(
"clone not supported for this type of Index");
std::vector< uint8_t > codes
Codes. Size ntotal * pq.code_size.
Randomly rotate a set of vectors.
Index * read_index(FILE *f, bool try_mmap)
int bytes_per_vec
nb of 8-bits per encoded vector
std::vector< float > thresholds
thresholds to compare with
bool train_thresholds
whether we train thresholds or use 0
Index * base_index
faster index to pre-select the vectors that should be filtered
std::vector< std::vector< uint8_t > > codes
inverted list codes.
IndexFlat refine_index
storage for full vectors
bool own_fields
should the base index be deallocated?
std::vector< long > id_map
! whether pointers are deleted in destructo
RandomRotationMatrix rrot
optional random rotation
long idx_t
all indices are this type
ProductQuantizer pq
The product quantizer used to encode the vectors.
idx_t ntotal
total nb of indexed vectors
MetricType metric_type
type of metric this index uses for search
size_t nlist
number of possible key values
int nbits
nb of bits per vector
bool is_trained
set if the Index does not require training, or if training is done already
std::vector< float > xb
database vectors, size ntotal * d
int polysemous_ht
Hamming threshold used for polysemy.
bool rotate_data
whether to apply a random rotation to input
std::vector< uint8_t > codes
encoded dataset
std::vector< std::vector< float > > vecs
bool own_fields
! the sub-index