19 #include "FaissAssert.h"
21 #include "IndexFlat.h"
22 #include "VectorTransform.h"
26 #include "IndexIVFPQ.h"
43 static uint32_t fourcc (
const char sx[4]) {
44 const unsigned char *x = (
unsigned char*)sx;
45 return x[0] | x[1] << 8 | x[2] << 16 | x[3] << 24;
56 #define WRITEANDCHECK(ptr, n) { \
57 size_t ret = fwrite (ptr, sizeof (* (ptr)), n, f); \
58 FAISS_ASSERT (ret == (n) || !"write error"); \
61 #define READANDCHECK(ptr, n) { \
62 size_t ret = fread (ptr, sizeof (* (ptr)), n, f); \
63 FAISS_ASSERT (ret == (n) || !"write error"); \
66 #define WRITE1(x) WRITEANDCHECK(&(x), 1)
67 #define READ1(x) READANDCHECK(&(x), 1)
69 #define WRITEVECTOR(vec) { \
70 size_t size = (vec).size (); \
71 WRITEANDCHECK (&size, 1); \
72 WRITEANDCHECK ((vec).data (), size); \
75 #define READVECTOR(vec) { \
77 READANDCHECK (&size, 1); \
78 FAISS_ASSERT (size >= 0 && size < (1L << 40)); \
79 (vec).resize (size); \
80 READANDCHECK ((vec).data (), size); \
87 #define WRITETABPAD16(tab, size_in) { \
88 size_t size = (size_in); \
89 WRITEANDCHECK (&size, 1); \
90 uint8_t padding[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; \
91 int idx = ftell(f) % 16; \
92 padding [idx] = 15 - idx; \
93 WRITEANDCHECK (padding + idx, 16 - idx); \
94 WRITEANDCHECK ((tab), size); \
97 #define READTABPAD16(tab, basetype, expected_size) { \
99 READANDCHECK (&size, 1); \
100 FAISS_ASSERT ((expected_size) == size); \
101 uint8_t padding[16], npad; \
103 FAISS_ASSERT (npad < 16); \
104 READANDCHECK (padding, npad); \
105 (tab) = new basetype [size]; \
106 READANDCHECK ((tab), size); \
110 #define TABOFFSETPAD16(taboffset, basetype, expected_size) { \
112 READANDCHECK (&size, 1); \
113 FAISS_ASSERT ((expected_size) == size); \
114 uint8_t padding[16], npad; \
116 FAISS_ASSERT (npad < 16); \
117 READANDCHECK (padding, npad); \
118 taboffset = ftell(f); \
119 fseek (f, sizeof(basetype) * size, SEEK_CUR); \
129 static void write_index_header (
const Index *idx, FILE *f) {
131 WRITE1 (idx->ntotal);
135 WRITE1 (idx->is_trained);
136 WRITE1 (idx->metric_type);
141 void write_VectorTransform (
const VectorTransform *vt, FILE *f) {
142 if (
const LinearTransform * lt =
143 dynamic_cast < const LinearTransform *> (vt)) {
144 if (dynamic_cast<const RandomRotationMatrix *>(lt)) {
145 uint32_t h = fourcc (
"rrot");
147 }
else if (
const PCAMatrix * pca =
148 dynamic_cast<const PCAMatrix *>(lt)) {
149 uint32_t h = fourcc (
"PcAm");
151 WRITE1 (pca->eigen_power);
152 WRITE1 (pca->random_rotation);
153 WRITE1 (pca->balanced_bins);
154 WRITEVECTOR (pca->mean);
155 WRITEVECTOR (pca->eigenvalues);
156 WRITEVECTOR (pca->PCAMat);
159 uint32_t h = fourcc (
"LTra");
162 WRITE1 (lt->have_bias);
165 }
else if (
const RemapDimensionsTransform *rdt =
166 dynamic_cast<const RemapDimensionsTransform *>(vt)) {
167 uint32_t h = fourcc (
"RmDT");
169 WRITEVECTOR (rdt->map);
170 }
else FAISS_ASSERT (!
"cannot serialize this");
174 WRITE1 (vt->is_trained);
177 static void write_ProductQuantizer (
const ProductQuantizer *pq, FILE *f) {
181 WRITEVECTOR (pq->centroids);
184 void write_ProductQuantizer (
const ProductQuantizer*pq,
const char *fname) {
185 FILE *f = fopen (fname,
"w");
187 fprintf (stderr,
"cannot open %s for writing:", fname);
191 write_ProductQuantizer (pq, f);
197 static void write_ivf_header (
const IndexIVF * ivf, FILE *f,
198 bool include_ids =
true) {
199 write_index_header (ivf, f);
201 WRITE1 (ivf->nprobe);
202 write_index (ivf->quantizer, f);
204 for (
size_t i = 0; i < ivf->nlist; i++)
205 WRITEVECTOR (ivf->ids[i]);
207 WRITE1 (ivf->maintain_direct_map);
208 WRITEVECTOR (ivf->direct_map);
211 void write_index (
const Index *idx, FILE *f) {
212 if (
const IndexFlat * idxf = dynamic_cast<const IndexFlat *> (idx)) {
213 uint32_t h = fourcc (
214 idxf->metric_type == METRIC_INNER_PRODUCT ?
"IxFI" :
215 idxf->metric_type == METRIC_L2 ?
"IxF2" :
nullptr);
217 write_index_header (idx, f);
218 WRITEVECTOR (idxf->xb);
219 }
else if(
const IndexLSH * idxl = dynamic_cast<const IndexLSH *> (idx)) {
220 uint32_t h = fourcc (
"IxHe");
222 write_index_header (idx, f);
223 WRITE1 (idxl->nbits);
224 WRITE1 (idxl->rotate_data);
225 WRITE1 (idxl->train_thresholds);
226 WRITEVECTOR (idxl->thresholds);
227 WRITE1 (idxl->bytes_per_vec);
228 write_VectorTransform (&idxl->rrot, f);
229 WRITEVECTOR (idxl->codes);
230 }
else if(
const IndexPQ * idxp = dynamic_cast<const IndexPQ *> (idx)) {
231 uint32_t h = fourcc (
"IxPq");
233 write_index_header (idx, f);
234 write_ProductQuantizer (&idxp->pq, f);
235 WRITEVECTOR (idxp->codes);
237 WRITE1 (idxp->search_type);
238 WRITE1 (idxp->encode_signs);
239 WRITE1 (idxp->polysemous_ht);
240 }
else if(
const IndexIVFFlat * ivfl =
241 dynamic_cast<const IndexIVFFlat *> (idx)) {
242 uint32_t h = fourcc (
"IvFl");
244 write_ivf_header (ivfl, f);
245 for(
int i = 0; i < ivfl->nlist; i++)
246 WRITEVECTOR (ivfl->vecs[i]);
247 }
else if(
const IndexIVFPQ * ivpq =
248 dynamic_cast<const IndexIVFPQ *> (idx)) {
249 const IndexIVFPQR * ivfpqr =
dynamic_cast<const IndexIVFPQR *
> (idx);
250 const IndexIVFPQCompact * ivfpqc =
251 dynamic_cast<const IndexIVFPQCompact *
> (idx);
252 uint32_t h = fourcc (ivfpqr ?
"IvQR" : ivfpqc ?
"IvPC" :
"IvPQ");
254 write_ivf_header (ivpq, f, !ivfpqc);
255 WRITE1 (ivpq->by_residual);
256 WRITE1 (ivpq->code_size);
257 write_ProductQuantizer (&ivpq->pq, f);
259 for(
int i = 0; i < ivpq->codes.size(); i++)
260 WRITEVECTOR (ivpq->codes[i]);
263 write_ProductQuantizer (&ivfpqr->refine_pq, f);
264 WRITEVECTOR (ivfpqr->refine_codes);
265 WRITE1 (ivfpqr->k_factor);
268 WRITETABPAD16 (ivfpqc->limits, ivfpqc->nlist + 1);
269 WRITETABPAD16 (ivfpqc->compact_ids, ivfpqc->ntotal);
270 WRITETABPAD16 (ivfpqc->compact_codes,
271 ivfpqc->ntotal * ivfpqc->code_size);
273 }
else if(
const IndexPreTransform * ixpt =
274 dynamic_cast<const IndexPreTransform *> (idx)) {
275 uint32_t h = fourcc (
"IxPT");
277 write_index_header (ixpt, f);
278 int nt = ixpt->chain.size();
280 for (
int i = 0; i < nt; i++)
281 write_VectorTransform (ixpt->chain[i], f);
282 write_index (ixpt->index, f);
283 }
else if(
const MultiIndexQuantizer * imiq =
284 dynamic_cast<const MultiIndexQuantizer *> (idx)) {
285 uint32_t h = fourcc (
"Imiq");
287 write_index_header (imiq, f);
288 write_ProductQuantizer (&imiq->pq, f);
289 }
else if(
const IndexRefineFlat * idxrf =
290 dynamic_cast<const IndexRefineFlat *> (idx)) {
291 uint32_t h = fourcc (
"IxRF");
293 write_index_header (idxrf, f);
294 write_index (idxrf->base_index, f);
295 write_index (&idxrf->refine_index, f);
296 WRITE1 (idxrf->k_factor);
298 FAISS_ASSERT (!
"don't know how to serialize this type of index");
302 void write_index (
const Index *idx,
const char *fname) {
303 FILE *f = fopen (fname,
"w");
305 fprintf (stderr,
"cannot open %s for writing:", fname);
309 write_index (idx, f);
313 void write_VectorTransform (
const VectorTransform *vt,
const char *fname) {
314 FILE *f = fopen (fname,
"w");
316 fprintf (stderr,
"cannot open %s for writing:", fname);
320 write_VectorTransform (vt, f);
328 static void read_index_header (Index *idx, FILE *f) {
334 READ1 (idx->is_trained);
335 READ1 (idx->metric_type);
336 idx->verbose =
false;
339 VectorTransform* read_VectorTransform (FILE *f) {
342 VectorTransform *vt =
nullptr;
344 if (h == fourcc (
"rrot") || h == fourcc (
"PCAm") ||
345 h == fourcc (
"LTra") || h == fourcc (
"PcAm")) {
346 LinearTransform *lt =
nullptr;
347 if (h == fourcc (
"rrot")) {
348 lt =
new RandomRotationMatrix ();
349 }
else if (h == fourcc (
"PCAm") ||
350 h == fourcc (
"PcAm")) {
351 PCAMatrix * pca =
new PCAMatrix ();
352 READ1 (pca->eigen_power);
353 READ1 (pca->random_rotation);
354 if (h == fourcc (
"PcAm"))
355 READ1 (pca->balanced_bins);
356 READVECTOR (pca->mean);
357 READVECTOR (pca->eigenvalues);
358 READVECTOR (pca->PCAMat);
360 }
else if (h == fourcc (
"LTra")) {
361 lt =
new LinearTransform ();
363 READ1 (lt->have_bias);
367 }
else if (h == fourcc (
"RmDT")) {
368 RemapDimensionsTransform *rdt =
new RemapDimensionsTransform ();
369 READVECTOR (rdt->map);
371 }
else FAISS_ASSERT(!
"fourcc not recognized");
374 READ1 (vt->is_trained);
378 static void read_ProductQuantizer (ProductQuantizer *pq, FILE *f) {
382 pq->set_derived_values ();
383 READVECTOR (pq->centroids);
386 ProductQuantizer * read_ProductQuantizer (
const char*fname) {
387 FILE *f = fopen (fname,
"r");
389 fprintf (stderr,
"cannot open %s for reading:", fname);
393 ProductQuantizer *pq =
new ProductQuantizer();
394 read_ProductQuantizer(pq, f);
399 static void read_ivf_header (IndexIVF * ivf, FILE *f,
400 bool include_ids =
true) {
401 read_index_header (ivf, f);
405 ivf->own_fields =
true;
407 ivf->ids.resize (ivf->nlist);
408 for (
size_t i = 0; i < ivf->nlist; i++)
409 READVECTOR (ivf->ids[i]);
411 READ1 (ivf->maintain_direct_map);
412 READVECTOR (ivf->direct_map);
415 static IndexIVFPQ *read_ivfpq (FILE *f, uint32_t h,
bool try_mmap)
418 IndexIVFPQR *ivfpqr =
419 h == fourcc (
"IvQR") ?
new IndexIVFPQR () : nullptr;
420 IndexIVFPQCompact *ivfpqc =
421 h == fourcc (
"IvPC") ?
new IndexIVFPQCompact () : nullptr;
422 IndexIVFPQ * ivpq = ivfpqr ? ivfpqr : ivfpqc ? ivfpqc :
new IndexIVFPQ ();
423 read_ivf_header (ivpq, f, !ivfpqc);
424 READ1 (ivpq->by_residual);
425 READ1 (ivpq->code_size);
426 read_ProductQuantizer (&ivpq->pq, f);
428 ivpq->codes.resize (ivpq->nlist);
429 for (
size_t i = 0; i < ivpq->nlist; i++)
430 READVECTOR (ivpq->codes[i]);
433 ivpq->use_precomputed_table = 0;
434 if (ivpq->by_residual)
435 ivpq->precompute_table ();
437 read_ProductQuantizer (&ivfpqr->refine_pq, f);
438 READVECTOR (ivfpqr->refine_codes);
439 READ1 (ivfpqr->k_factor);
443 READTABPAD16 (ivfpqc->limits, uint32_t, ivfpqc->nlist + 1);
444 READTABPAD16 (ivfpqc->compact_ids, uint32_t, ivfpqc->ntotal);
445 READTABPAD16 (ivfpqc->compact_codes, uint8_t,
446 ivfpqc->ntotal * ivfpqc->code_size);
448 long offset_limits, offset_compact_ids, offset_compact_codes;
449 TABOFFSETPAD16 (offset_limits, uint32_t, ivfpqc->nlist + 1);
450 TABOFFSETPAD16 (offset_compact_ids, uint32_t, ivfpqc->ntotal);
451 TABOFFSETPAD16 (offset_compact_codes, uint8_t,
452 ivfpqc->ntotal * ivfpqc->code_size);
453 ivfpqc->mmap_length = ftell (f);
455 ivfpqc->mmap_buffer = (
char*)mmap (
456 nullptr, ivfpqc->mmap_length,
457 PROT_READ, MAP_SHARED, fileno (f), 0);
458 if (!ivfpqc->mmap_buffer) {
459 perror (
"mmap failed");
464 ivfpqc->limits = (uint32_t*)(ivfpqc->mmap_buffer + offset_limits);
465 ivfpqc->compact_ids = (uint32_t*)(ivfpqc->mmap_buffer +
467 ivfpqc->compact_codes = (uint8_t*)(ivfpqc->mmap_buffer +
468 offset_compact_codes);
474 int read_old_fmt_hack = 0;
477 Index * idx =
nullptr;
480 if (h == fourcc (
"IxFI") || h == fourcc (
"IxF2")) {
482 if (h == fourcc (
"IxFI")) idxf =
new IndexFlatIP ();
484 read_index_header (idxf, f);
485 READVECTOR (idxf->
xb);
486 FAISS_ASSERT (idxf->
xb.size() == idxf->
ntotal * idxf->
d);
488 }
else if (h == fourcc(
"IxHE") || h == fourcc(
"IxHe")) {
490 read_index_header (idxl, f);
496 if (h == fourcc(
"IxHE")) {
497 FAISS_ASSERT (idxl->
nbits % 64 == 0 ||
498 !
"can only read old format IndexLSH with "
499 "nbits multiple of 64");
504 (read_VectorTransform (f));
505 FAISS_ASSERT(rrot || !
"expected a random rotation");
509 READVECTOR (idxl->
codes);
510 FAISS_ASSERT (idxl->
rrot.d_in == idxl->
d &&
514 }
else if (h == fourcc (
"IxPQ") || h == fourcc (
"IxPo") ||
515 h == fourcc (
"IxPq")) {
518 read_index_header (idxp, f);
519 read_ProductQuantizer (&idxp->
pq, f);
520 READVECTOR (idxp->
codes);
521 if (h == fourcc (
"IxPo") || h == fourcc (
"IxPq")) {
522 READ1 (idxp->search_type);
523 READ1 (idxp->encode_signs);
529 if (h == fourcc (
"IxPQ") || h == fourcc (
"IxPo")) {
533 }
else if(h == fourcc (
"IvFl")) {
535 read_ivf_header (ivfl, f);
537 for (
size_t i = 0; i < ivfl->
nlist; i++)
538 READVECTOR (ivfl->
vecs[i]);
541 }
else if(h == fourcc (
"IvPQ") || h == fourcc (
"IvQR") ||
542 h == fourcc (
"IvPC")) {
544 idx = read_ivfpq (f, h, try_mmap);
546 }
else if(h == fourcc (
"IxPT")) {
549 read_index_header (ixpt, f);
551 if (read_old_fmt_hack == 2) {
556 for (
int i = 0; i < nt; i++)
557 ixpt->chain.push_back (read_VectorTransform (f));
560 }
else if(h == fourcc (
"Imiq")) {
562 read_index_header (imiq, f);
563 read_ProductQuantizer (&imiq->pq, f);
565 }
else if(h == fourcc (
"IxRF")) {
567 read_index_header (idxrf, f);
576 fprintf (stderr,
"Index type 0x%08x not supported\n", h);
583 Index *
read_index (
const char *fname,
bool try_mmap) {
584 FILE *f = fopen (fname,
"r");
586 fprintf (stderr,
"cannot open %s for reading:", fname);
595 VectorTransform *read_VectorTransform (
const char *fname) {
596 FILE *f = fopen (fname,
"r");
598 fprintf (stderr,
"cannot open %s for reading:", fname);
602 VectorTransform *vt = read_VectorTransform (f);
613 Index * clone_index (
const Index *index)
616 return cl.clone_Index (index);
621 #define TRYCLONE(classname, obj) \
622 if (const classname *clo = dynamic_cast<const classname *>(obj)) { \
623 return new classname(*clo); \
626 VectorTransform *Cloner::clone_VectorTransform (
const VectorTransform *vt)
628 TRYCLONE (RemapDimensionsTransform, vt)
629 TRYCLONE (OPQMatrix, vt)
630 TRYCLONE (PCAMatrix, vt)
631 TRYCLONE (RandomRotationMatrix, vt)
632 TRYCLONE (LinearTransform, vt)
634 FAISS_ASSERT(!
"clone not supported for this type of VectorTransform");
639 IndexIVF * Cloner::clone_IndexIVF (
const IndexIVF *ivf)
641 TRYCLONE (IndexIVFPQR, ivf)
642 TRYCLONE (IndexIVFPQ, ivf)
643 TRYCLONE (IndexIVFFlat, ivf)
645 FAISS_ASSERT(!
"clone not supported for this type of IndexIVF");
650 Index *Cloner::clone_Index (
const Index *index)
652 TRYCLONE (IndexPQ, index)
653 TRYCLONE (IndexLSH, index)
654 TRYCLONE (IndexFlatL2, index)
655 TRYCLONE (IndexFlatIP, index)
656 TRYCLONE (IndexFlat, index)
657 TRYCLONE (MultiIndexQuantizer, index)
658 if (const IndexIVF * ivf = dynamic_cast<const IndexIVF*>(index)) {
659 IndexIVF *res = clone_IndexIVF (ivf);
660 res->own_fields =
true;
661 res->quantizer = clone_Index (ivf->quantizer);
663 }
else if (
const IndexPreTransform * ipt =
664 dynamic_cast<const IndexPreTransform*> (index)) {
665 IndexPreTransform *res =
new IndexPreTransform ();
667 res->index = clone_Index (ipt->index);
668 for (
int i = 0; i < ipt->chain.size(); i++)
669 res->chain.push_back (clone_VectorTransform (ipt->chain[i]));
670 res->own_fields =
true;
673 FAISS_ASSERT(!
"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
IndexFlat refine_index
storage for full vectors
bool own_fields
should the base index be deallocated?
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
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