19 #include "FaissAssert.h"
21 #include "IndexFlat.h"
22 #include "VectorTransform.h"
26 #include "IndexIVFPQ.h"
27 #include "MetaIndexes.h"
52 static uint32_t fourcc (
const char sx[4]) {
53 const unsigned char *x = (
unsigned char*)sx;
54 return x[0] | x[1] << 8 | x[2] << 16 | x[3] << 24;
65 #define WRITEANDCHECK(ptr, n) { \
66 size_t ret = fwrite (ptr, sizeof (* (ptr)), n, f); \
67 FAISS_ASSERT (ret == (n) || !"write error"); \
70 #define READANDCHECK(ptr, n) { \
71 size_t ret = fread (ptr, sizeof (* (ptr)), n, f); \
72 FAISS_ASSERT (ret == (n) || !"write error"); \
75 #define WRITE1(x) WRITEANDCHECK(&(x), 1)
76 #define READ1(x) READANDCHECK(&(x), 1)
78 #define WRITEVECTOR(vec) { \
79 size_t size = (vec).size (); \
80 WRITEANDCHECK (&size, 1); \
81 WRITEANDCHECK ((vec).data (), size); \
84 #define READVECTOR(vec) { \
86 READANDCHECK (&size, 1); \
87 FAISS_ASSERT (size >= 0 && size < (1L << 40)); \
88 (vec).resize (size); \
89 READANDCHECK ((vec).data (), size); \
96 #define WRITETABPAD16(tab, size_in) { \
97 size_t size = (size_in); \
98 WRITEANDCHECK (&size, 1); \
99 uint8_t padding[16] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}; \
100 int idx = ftell(f) % 16; \
101 padding [idx] = 15 - idx; \
102 WRITEANDCHECK (padding + idx, 16 - idx); \
103 WRITEANDCHECK ((tab), size); \
106 #define READTABPAD16(tab, basetype, expected_size) { \
108 READANDCHECK (&size, 1); \
109 FAISS_ASSERT ((expected_size) == size); \
110 uint8_t padding[16], npad; \
112 FAISS_ASSERT (npad < 16); \
113 READANDCHECK (padding, npad); \
114 (tab) = new basetype [size]; \
115 READANDCHECK ((tab), size); \
119 #define TABOFFSETPAD16(taboffset, basetype, expected_size) { \
121 READANDCHECK (&size, 1); \
122 FAISS_ASSERT ((expected_size) == size); \
123 uint8_t padding[16], npad; \
125 FAISS_ASSERT (npad < 16); \
126 READANDCHECK (padding, npad); \
127 taboffset = ftell(f); \
128 fseek (f, sizeof(basetype) * size, SEEK_CUR); \
138 static void write_index_header (
const Index *idx, FILE *f) {
140 WRITE1 (idx->ntotal);
144 WRITE1 (idx->is_trained);
145 WRITE1 (idx->metric_type);
150 void write_VectorTransform (
const VectorTransform *vt, FILE *f) {
151 if (
const LinearTransform * lt =
152 dynamic_cast < const LinearTransform *> (vt)) {
153 if (dynamic_cast<const RandomRotationMatrix *>(lt)) {
154 uint32_t h = fourcc (
"rrot");
156 }
else if (
const PCAMatrix * pca =
157 dynamic_cast<const PCAMatrix *>(lt)) {
158 uint32_t h = fourcc (
"PcAm");
160 WRITE1 (pca->eigen_power);
161 WRITE1 (pca->random_rotation);
162 WRITE1 (pca->balanced_bins);
163 WRITEVECTOR (pca->mean);
164 WRITEVECTOR (pca->eigenvalues);
165 WRITEVECTOR (pca->PCAMat);
168 uint32_t h = fourcc (
"LTra");
171 WRITE1 (lt->have_bias);
174 }
else if (
const RemapDimensionsTransform *rdt =
175 dynamic_cast<const RemapDimensionsTransform *>(vt)) {
176 uint32_t h = fourcc (
"RmDT");
178 WRITEVECTOR (rdt->map);
179 }
else FAISS_ASSERT (!
"cannot serialize this");
183 WRITE1 (vt->is_trained);
186 static void write_ProductQuantizer (
const ProductQuantizer *pq, FILE *f) {
190 WRITEVECTOR (pq->centroids);
193 void write_ProductQuantizer (
const ProductQuantizer*pq,
const char *fname) {
194 FILE *f = fopen (fname,
"w");
196 fprintf (stderr,
"cannot open %s for writing:", fname);
200 write_ProductQuantizer (pq, f);
206 static void write_ivf_header (
const IndexIVF * ivf, FILE *f,
207 bool include_ids =
true) {
208 write_index_header (ivf, f);
210 WRITE1 (ivf->nprobe);
211 write_index (ivf->quantizer, f);
213 for (
size_t i = 0; i < ivf->nlist; i++)
214 WRITEVECTOR (ivf->ids[i]);
216 WRITE1 (ivf->maintain_direct_map);
217 WRITEVECTOR (ivf->direct_map);
220 void write_index (
const Index *idx, FILE *f) {
221 if (
const IndexFlat * idxf = dynamic_cast<const IndexFlat *> (idx)) {
222 uint32_t h = fourcc (
223 idxf->metric_type == METRIC_INNER_PRODUCT ?
"IxFI" :
224 idxf->metric_type == METRIC_L2 ?
"IxF2" :
nullptr);
226 write_index_header (idx, f);
227 WRITEVECTOR (idxf->xb);
228 }
else if(
const IndexLSH * idxl = dynamic_cast<const IndexLSH *> (idx)) {
229 uint32_t h = fourcc (
"IxHe");
231 write_index_header (idx, f);
232 WRITE1 (idxl->nbits);
233 WRITE1 (idxl->rotate_data);
234 WRITE1 (idxl->train_thresholds);
235 WRITEVECTOR (idxl->thresholds);
236 WRITE1 (idxl->bytes_per_vec);
237 write_VectorTransform (&idxl->rrot, f);
238 WRITEVECTOR (idxl->codes);
239 }
else if(
const IndexPQ * idxp = dynamic_cast<const IndexPQ *> (idx)) {
240 uint32_t h = fourcc (
"IxPq");
242 write_index_header (idx, f);
243 write_ProductQuantizer (&idxp->pq, f);
244 WRITEVECTOR (idxp->codes);
246 WRITE1 (idxp->search_type);
247 WRITE1 (idxp->encode_signs);
248 WRITE1 (idxp->polysemous_ht);
249 }
else if(
const IndexIVFFlat * ivfl =
250 dynamic_cast<const IndexIVFFlat *> (idx)) {
251 uint32_t h = fourcc (
"IvFl");
253 write_ivf_header (ivfl, f);
254 for(
int i = 0; i < ivfl->nlist; i++)
255 WRITEVECTOR (ivfl->vecs[i]);
256 }
else if(
const IndexIVFPQ * ivpq =
257 dynamic_cast<const IndexIVFPQ *> (idx)) {
258 const IndexIVFPQR * ivfpqr =
dynamic_cast<const IndexIVFPQR *
> (idx);
259 const IndexIVFPQCompact * ivfpqc =
260 dynamic_cast<const IndexIVFPQCompact *
> (idx);
261 uint32_t h = fourcc (ivfpqr ?
"IvQR" : ivfpqc ?
"IvPC" :
"IvPQ");
263 write_ivf_header (ivpq, f, !ivfpqc);
264 WRITE1 (ivpq->by_residual);
265 WRITE1 (ivpq->code_size);
266 write_ProductQuantizer (&ivpq->pq, f);
268 for(
int i = 0; i < ivpq->codes.size(); i++)
269 WRITEVECTOR (ivpq->codes[i]);
272 write_ProductQuantizer (&ivfpqr->refine_pq, f);
273 WRITEVECTOR (ivfpqr->refine_codes);
274 WRITE1 (ivfpqr->k_factor);
277 WRITETABPAD16 (ivfpqc->limits, ivfpqc->nlist + 1);
278 WRITETABPAD16 (ivfpqc->compact_ids, ivfpqc->ntotal);
279 WRITETABPAD16 (ivfpqc->compact_codes,
280 ivfpqc->ntotal * ivfpqc->code_size);
282 }
else if(
const IndexPreTransform * ixpt =
283 dynamic_cast<const IndexPreTransform *> (idx)) {
284 uint32_t h = fourcc (
"IxPT");
286 write_index_header (ixpt, f);
287 int nt = ixpt->chain.size();
289 for (
int i = 0; i < nt; i++)
290 write_VectorTransform (ixpt->chain[i], f);
291 write_index (ixpt->index, f);
292 }
else if(
const MultiIndexQuantizer * imiq =
293 dynamic_cast<const MultiIndexQuantizer *> (idx)) {
294 uint32_t h = fourcc (
"Imiq");
296 write_index_header (imiq, f);
297 write_ProductQuantizer (&imiq->pq, f);
298 }
else if(
const IndexRefineFlat * idxrf =
299 dynamic_cast<const IndexRefineFlat *> (idx)) {
300 uint32_t h = fourcc (
"IxRF");
302 write_index_header (idxrf, f);
303 write_index (idxrf->base_index, f);
304 write_index (&idxrf->refine_index, f);
305 WRITE1 (idxrf->k_factor);
306 }
else if(
const IndexIDMap * idxmap =
307 dynamic_cast<const IndexIDMap *> (idx)) {
308 uint32_t h = fourcc (
"IxMp");
310 write_index_header (idxmap, f);
311 write_index (idxmap->index, f);
312 WRITEVECTOR (idxmap->id_map);
314 FAISS_ASSERT (!
"don't know how to serialize this type of index");
318 void write_index (
const Index *idx,
const char *fname) {
319 FILE *f = fopen (fname,
"w");
321 fprintf (stderr,
"cannot open %s for writing:", fname);
325 write_index (idx, f);
329 void write_VectorTransform (
const VectorTransform *vt,
const char *fname) {
330 FILE *f = fopen (fname,
"w");
332 fprintf (stderr,
"cannot open %s for writing:", fname);
336 write_VectorTransform (vt, f);
344 static void read_index_header (Index *idx, FILE *f) {
350 READ1 (idx->is_trained);
351 READ1 (idx->metric_type);
352 idx->verbose =
false;
355 VectorTransform* read_VectorTransform (FILE *f) {
358 VectorTransform *vt =
nullptr;
360 if (h == fourcc (
"rrot") || h == fourcc (
"PCAm") ||
361 h == fourcc (
"LTra") || h == fourcc (
"PcAm")) {
362 LinearTransform *lt =
nullptr;
363 if (h == fourcc (
"rrot")) {
364 lt =
new RandomRotationMatrix ();
365 }
else if (h == fourcc (
"PCAm") ||
366 h == fourcc (
"PcAm")) {
367 PCAMatrix * pca =
new PCAMatrix ();
368 READ1 (pca->eigen_power);
369 READ1 (pca->random_rotation);
370 if (h == fourcc (
"PcAm"))
371 READ1 (pca->balanced_bins);
372 READVECTOR (pca->mean);
373 READVECTOR (pca->eigenvalues);
374 READVECTOR (pca->PCAMat);
376 }
else if (h == fourcc (
"LTra")) {
377 lt =
new LinearTransform ();
379 READ1 (lt->have_bias);
383 }
else if (h == fourcc (
"RmDT")) {
384 RemapDimensionsTransform *rdt =
new RemapDimensionsTransform ();
385 READVECTOR (rdt->map);
387 }
else FAISS_ASSERT(!
"fourcc not recognized");
390 READ1 (vt->is_trained);
394 static void read_ProductQuantizer (ProductQuantizer *pq, FILE *f) {
398 pq->set_derived_values ();
399 READVECTOR (pq->centroids);
402 ProductQuantizer * read_ProductQuantizer (
const char*fname) {
403 FILE *f = fopen (fname,
"r");
405 fprintf (stderr,
"cannot open %s for reading:", fname);
409 ProductQuantizer *pq =
new ProductQuantizer();
410 read_ProductQuantizer(pq, f);
415 static void read_ivf_header (IndexIVF * ivf, FILE *f,
416 bool include_ids =
true) {
417 read_index_header (ivf, f);
421 ivf->own_fields =
true;
423 ivf->ids.resize (ivf->nlist);
424 for (
size_t i = 0; i < ivf->nlist; i++)
425 READVECTOR (ivf->ids[i]);
427 READ1 (ivf->maintain_direct_map);
428 READVECTOR (ivf->direct_map);
431 static IndexIVFPQ *read_ivfpq (FILE *f, uint32_t h,
bool try_mmap)
434 IndexIVFPQR *ivfpqr =
435 h == fourcc (
"IvQR") ?
new IndexIVFPQR () : nullptr;
436 IndexIVFPQCompact *ivfpqc =
437 h == fourcc (
"IvPC") ?
new IndexIVFPQCompact () : nullptr;
438 IndexIVFPQ * ivpq = ivfpqr ? ivfpqr : ivfpqc ? ivfpqc :
new IndexIVFPQ ();
439 read_ivf_header (ivpq, f, !ivfpqc);
440 READ1 (ivpq->by_residual);
441 READ1 (ivpq->code_size);
442 read_ProductQuantizer (&ivpq->pq, f);
444 ivpq->codes.resize (ivpq->nlist);
445 for (
size_t i = 0; i < ivpq->nlist; i++)
446 READVECTOR (ivpq->codes[i]);
449 ivpq->use_precomputed_table = 0;
450 if (ivpq->by_residual)
451 ivpq->precompute_table ();
453 read_ProductQuantizer (&ivfpqr->refine_pq, f);
454 READVECTOR (ivfpqr->refine_codes);
455 READ1 (ivfpqr->k_factor);
459 READTABPAD16 (ivfpqc->limits, uint32_t, ivfpqc->nlist + 1);
460 READTABPAD16 (ivfpqc->compact_ids, uint32_t, ivfpqc->ntotal);
461 READTABPAD16 (ivfpqc->compact_codes, uint8_t,
462 ivfpqc->ntotal * ivfpqc->code_size);
464 long offset_limits, offset_compact_ids, offset_compact_codes;
465 TABOFFSETPAD16 (offset_limits, uint32_t, ivfpqc->nlist + 1);
466 TABOFFSETPAD16 (offset_compact_ids, uint32_t, ivfpqc->ntotal);
467 TABOFFSETPAD16 (offset_compact_codes, uint8_t,
468 ivfpqc->ntotal * ivfpqc->code_size);
469 ivfpqc->mmap_length = ftell (f);
471 ivfpqc->mmap_buffer = (
char*)mmap (
472 nullptr, ivfpqc->mmap_length,
473 PROT_READ, MAP_SHARED, fileno (f), 0);
474 if (!ivfpqc->mmap_buffer) {
475 perror (
"mmap failed");
480 ivfpqc->limits = (uint32_t*)(ivfpqc->mmap_buffer + offset_limits);
481 ivfpqc->compact_ids = (uint32_t*)(ivfpqc->mmap_buffer +
483 ivfpqc->compact_codes = (uint8_t*)(ivfpqc->mmap_buffer +
484 offset_compact_codes);
490 int read_old_fmt_hack = 0;
493 Index * idx =
nullptr;
496 if (h == fourcc (
"IxFI") || h == fourcc (
"IxF2")) {
498 if (h == fourcc (
"IxFI")) idxf =
new IndexFlatIP ();
500 read_index_header (idxf, f);
501 READVECTOR (idxf->
xb);
502 FAISS_ASSERT (idxf->
xb.size() == idxf->
ntotal * idxf->
d);
504 }
else if (h == fourcc(
"IxHE") || h == fourcc(
"IxHe")) {
506 read_index_header (idxl, f);
512 if (h == fourcc(
"IxHE")) {
513 FAISS_ASSERT (idxl->
nbits % 64 == 0 ||
514 !
"can only read old format IndexLSH with "
515 "nbits multiple of 64");
520 (read_VectorTransform (f));
521 FAISS_ASSERT(rrot || !
"expected a random rotation");
525 READVECTOR (idxl->
codes);
526 FAISS_ASSERT (idxl->
rrot.d_in == idxl->
d &&
530 }
else if (h == fourcc (
"IxPQ") || h == fourcc (
"IxPo") ||
531 h == fourcc (
"IxPq")) {
534 read_index_header (idxp, f);
535 read_ProductQuantizer (&idxp->
pq, f);
536 READVECTOR (idxp->
codes);
537 if (h == fourcc (
"IxPo") || h == fourcc (
"IxPq")) {
538 READ1 (idxp->search_type);
539 READ1 (idxp->encode_signs);
545 if (h == fourcc (
"IxPQ") || h == fourcc (
"IxPo")) {
549 }
else if(h == fourcc (
"IvFl")) {
551 read_ivf_header (ivfl, f);
553 for (
size_t i = 0; i < ivfl->
nlist; i++)
554 READVECTOR (ivfl->
vecs[i]);
557 }
else if(h == fourcc (
"IvPQ") || h == fourcc (
"IvQR") ||
558 h == fourcc (
"IvPC")) {
560 idx = read_ivfpq (f, h, try_mmap);
562 }
else if(h == fourcc (
"IxPT")) {
565 read_index_header (ixpt, f);
567 if (read_old_fmt_hack == 2) {
572 for (
int i = 0; i < nt; i++)
573 ixpt->chain.push_back (read_VectorTransform (f));
576 }
else if(h == fourcc (
"Imiq")) {
578 read_index_header (imiq, f);
579 read_ProductQuantizer (&imiq->pq, f);
581 }
else if(h == fourcc (
"IxRF")) {
583 read_index_header (idxrf, f);
591 }
else if(h == fourcc (
"IxMp")) {
593 read_index_header (idxmap, f);
596 READVECTOR (idxmap->
id_map);
599 fprintf (stderr,
"Index type 0x%08x not supported\n", h);
606 Index *
read_index (
const char *fname,
bool try_mmap) {
607 FILE *f = fopen (fname,
"r");
609 fprintf (stderr,
"cannot open %s for reading:", fname);
618 VectorTransform *read_VectorTransform (
const char *fname) {
619 FILE *f = fopen (fname,
"r");
621 fprintf (stderr,
"cannot open %s for reading:", fname);
625 VectorTransform *vt = read_VectorTransform (f);
636 Index * clone_index (
const Index *index)
639 return cl.clone_Index (index);
644 #define TRYCLONE(classname, obj) \
645 if (const classname *clo = dynamic_cast<const classname *>(obj)) { \
646 return new classname(*clo); \
649 VectorTransform *Cloner::clone_VectorTransform (
const VectorTransform *vt)
651 TRYCLONE (RemapDimensionsTransform, vt)
652 TRYCLONE (OPQMatrix, vt)
653 TRYCLONE (PCAMatrix, vt)
654 TRYCLONE (RandomRotationMatrix, vt)
655 TRYCLONE (LinearTransform, vt)
657 FAISS_ASSERT(!
"clone not supported for this type of VectorTransform");
662 IndexIVF * Cloner::clone_IndexIVF (
const IndexIVF *ivf)
664 TRYCLONE (IndexIVFPQR, ivf)
665 TRYCLONE (IndexIVFPQ, ivf)
666 TRYCLONE (IndexIVFFlat, ivf)
668 FAISS_ASSERT(!
"clone not supported for this type of IndexIVF");
673 Index *Cloner::clone_Index (
const Index *index)
675 TRYCLONE (IndexPQ, index)
676 TRYCLONE (IndexLSH, index)
677 TRYCLONE (IndexFlatL2, index)
678 TRYCLONE (IndexFlatIP, index)
679 TRYCLONE (IndexFlat, index)
680 TRYCLONE (MultiIndexQuantizer, index)
681 if (const IndexIVF * ivf = dynamic_cast<const IndexIVF*>(index)) {
682 IndexIVF *res = clone_IndexIVF (ivf);
683 res->own_fields =
true;
684 res->quantizer = clone_Index (ivf->quantizer);
686 }
else if (
const IndexPreTransform * ipt =
687 dynamic_cast<const IndexPreTransform*> (index)) {
688 IndexPreTransform *res =
new IndexPreTransform ();
690 res->index = clone_Index (ipt->index);
691 for (
int i = 0; i < ipt->chain.size(); i++)
692 res->chain.push_back (clone_VectorTransform (ipt->chain[i]));
693 res->own_fields =
true;
696 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?
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
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