21 #include "FaissAssert.h"
22 #include "IndexFlat.h"
23 #include "AuxIndexStructures.h"
27 using ScopedIds = InvertedLists::ScopedIds;
28 using ScopedCodes = InvertedLists::ScopedCodes;
35 Level1Quantizer::Level1Quantizer (Index * quantizer,
size_t nlist):
36 quantizer (quantizer),
38 quantizer_trains_alone (0),
40 clustering_index (nullptr)
48 Level1Quantizer::Level1Quantizer ():
51 quantizer_trains_alone (0), own_fields (false),
52 clustering_index (nullptr)
55 Level1Quantizer::~Level1Quantizer ()
65 printf (
"IVF quantizer does not need training.\n");
68 printf (
"IVF quantizer trains alone...\n");
72 "nlist not consistent with quantizer size");
75 printf (
"Training level-1 quantizer on %ld vectors in %ldD\n",
90 "Training L2 quantizer on %ld vectors in %ldD%s\n",
93 FAISS_THROW_IF_NOT (metric_type == METRIC_L2);
97 clus.
train(n, x, assigner);
102 printf (
"Adding centroids to quantizer\n");
114 IndexIVF::IndexIVF (
Index * quantizer,
size_t d,
115 size_t nlist,
size_t code_size,
121 code_size (code_size),
125 maintain_direct_map (false)
127 FAISS_THROW_IF_NOT (d == quantizer->
d);
136 IndexIVF::IndexIVF ():
137 invlists (nullptr), own_invlists (false),
139 nprobe (1), max_codes (0), parallel_mode (0),
140 maintain_direct_map (false)
154 for (idx_t i0 = 0; i0 < n; i0 += bs) {
155 idx_t i1 = std::min (n, i0 + bs);
157 printf(
" IndexIVF::add_with_ids %ld:%ld\n", i0, i1);
160 xids ? xids + i0 :
nullptr);
166 std::unique_ptr<idx_t []> idx(
new idx_t[n]);
168 size_t nadd = 0, nminus1 = 0;
170 for (
size_t i = 0; i < n; i++) {
171 if (idx[i] < 0) nminus1++;
174 std::unique_ptr<uint8_t []> flat_codes(
new uint8_t [n *
code_size]);
177 #pragma omp parallel reduction(+: nadd)
179 int nt = omp_get_num_threads();
180 int rank = omp_get_thread_num();
183 for (
size_t i = 0; i < n; i++) {
184 long list_no = idx [i];
185 if (list_no >= 0 && list_no % nt == rank) {
186 long id = xids ? xids[i] :
ntotal + i;
195 printf(
" added %ld / %ld vectors (%ld -1s)\n", nadd, n, nminus1);
208 if (new_maintain_direct_map) {
209 direct_map.resize (
ntotal, -1);
210 for (
size_t key = 0; key <
nlist; key++) {
214 for (
long ofs = 0; ofs < list_size; ofs++) {
215 FAISS_THROW_IF_NOT_MSG (
216 0 <= idlist [ofs] && idlist[ofs] <
ntotal,
217 "direct map supported only for seuquential ids");
218 direct_map [idlist [ofs]] = key << 32 | ofs;
229 float *distances, idx_t *labels)
const
231 long * idx =
new long [n *
nprobe];
233 float * coarse_dis =
new float [n *
nprobe];
238 indexIVF_stats.quantization_time +=
getmillisecs() - t0;
244 distances, labels,
false);
252 const float *coarse_dis ,
253 float *distances, idx_t *labels,
260 size_t nlistv = 0, ndis = 0, nheap = 0;
268 for (idx_t i0 = 0; i0 < n; i0 += check_period) {
269 idx_t i1 = std::min(i0 + check_period, n);
271 #pragma omp parallel reduction(+: nlistv, ndis, nheap)
284 auto init_result = [&](
float *simi, idx_t *idxi) {
286 heap_heapify<HeapForIP> (k, simi, idxi);
288 heap_heapify<HeapForL2> (k, simi, idxi);
292 auto reorder_result = [&] (
float *simi, idx_t *idxi) {
294 heap_reorder<HeapForIP> (k, simi, idxi);
296 heap_reorder<HeapForL2> (k, simi, idxi);
302 auto scan_one_list = [&] (idx_t key,
float coarse_dis_i,
303 float *simi, idx_t *idxi) {
309 FAISS_THROW_IF_NOT_FMT (key < (idx_t)
nlist,
310 "Invalid key=%ld nlist=%ld\n",
316 if (list_size == 0) {
320 scanner->
set_list (key, coarse_dis_i);
326 std::unique_ptr<InvertedLists::ScopedIds> sids;
334 nheap += scanner->
scan_codes (list_size, scodes.get(),
347 for (
size_t i = i0; i < i1; i++) {
350 float * simi = distances + i * k;
351 idx_t * idxi = labels + i * k;
353 init_result (simi, idxi);
358 for (
size_t ik = 0; ik <
nprobe; ik++) {
360 nscan += scan_one_list (
361 keys [i * nprobe + ik],
362 coarse_dis[i * nprobe + ik],
366 if (max_codes && nscan >= max_codes) {
372 reorder_result (simi, idxi);
375 std::vector <idx_t> local_idx (k);
376 std::vector <float> local_dis (k);
378 for (
size_t i = i0; i < i1; i++) {
380 init_result (local_dis.data(), local_idx.data());
382 #pragma omp for schedule(dynamic)
383 for (
size_t ik = 0; ik <
nprobe; ik++) {
384 ndis += scan_one_list (
385 keys [i * nprobe + ik],
386 coarse_dis[i * nprobe + ik],
387 local_dis.data(), local_idx.data()
394 float * simi = distances + i * k;
395 idx_t * idxi = labels + i * k;
397 init_result (simi, idxi);
405 local_dis.data(), local_idx.data(), k);
409 local_dis.data(), local_idx.data(), k);
414 reorder_result (simi, idxi);
417 FAISS_THROW_FMT (
"parallel_mode %d not supported\n",
424 indexIVF_stats.nq += n;
425 indexIVF_stats.nlist += nlistv;
426 indexIVF_stats.ndis += ndis;
427 indexIVF_stats.nheap_updates += nheap;
437 std::unique_ptr<idx_t[]> keys (
new idx_t[nx *
nprobe]);
438 std::unique_ptr<float []> coarse_dis (
new float[nx * nprobe]);
442 indexIVF_stats.quantization_time +=
getmillisecs() - t0;
447 range_search_preassigned (nx, x, radius, keys.get (), coarse_dis.get (),
453 void IndexIVF::range_search_preassigned (
454 idx_t nx,
const float *x,
float radius,
455 const idx_t *keys,
const float *coarse_dis,
459 size_t nlistv = 0, ndis = 0;
460 bool store_pairs =
false;
462 std::vector<RangeSearchPartialResult *> all_pres (omp_get_max_threads());
464 #pragma omp parallel reduction(+: nlistv, ndis)
467 std::unique_ptr<InvertedListScanner> scanner
469 FAISS_THROW_IF_NOT (scanner.get ());
470 all_pres[omp_get_thread_num()] = &pres;
476 idx_t key = keys[i *
nprobe + ik];
478 FAISS_THROW_IF_NOT_FMT (
480 "Invalid key=%ld at ik=%ld nlist=%ld\n",
484 if (list_size == 0)
return;
489 scanner->set_list (key, coarse_dis[i *
nprobe + ik]);
492 scanner->scan_codes_range (list_size, scodes.get(),
493 ids.get(), radius, qres);
499 for (
size_t i = 0; i < nx; i++) {
500 scanner->set_query (x + i *
d);
502 RangeQueryResult & qres = pres.new_result (i);
504 for (
size_t ik = 0; ik <
nprobe; ik++) {
505 scan_list_func (i, ik, qres);
512 for (
size_t i = 0; i < nx; i++) {
513 scanner->set_query (x + i * d);
515 RangeQueryResult & qres = pres.new_result (i);
517 #pragma omp for schedule(dynamic)
518 for (
size_t ik = 0; ik <
nprobe; ik++) {
519 scan_list_func (i, ik, qres);
523 std::vector<RangeQueryResult *> all_qres (nx);
524 RangeQueryResult *qres =
nullptr;
526 #pragma omp for schedule(dynamic)
527 for (
size_t iik = 0; iik < nx *
nprobe; iik++) {
530 if (qres ==
nullptr || qres->qno != i) {
531 FAISS_ASSERT (!qres || i > qres->qno);
532 qres = &pres.new_result (i);
533 scanner->set_query (x + i * d);
535 scan_list_func (i, ik, *qres);
538 FAISS_THROW_FMT (
"parallel_mode %d not supported\n",
parallel_mode);
550 indexIVF_stats.nq += nx;
551 indexIVF_stats.nlist += nlistv;
552 indexIVF_stats.ndis += ndis;
564 FAISS_THROW_IF_NOT_MSG (direct_map.size() ==
ntotal,
565 "direct map is not initialized");
566 FAISS_THROW_IF_NOT_MSG (key >= 0 && key < direct_map.size(),
568 idx_t list_no = direct_map[key] >> 32;
569 idx_t offset = direct_map[key] & 0xffffffff;
576 FAISS_THROW_IF_NOT (ni == 0 || (i0 >= 0 && i0 + ni <=
ntotal));
578 for (idx_t list_no = 0; list_no <
nlist; list_no++) {
582 for (idx_t offset = 0; offset < list_size; offset++) {
583 idx_t
id = idlist[offset];
584 if (!(
id >= i0 &&
id < i0 + ni)) {
588 float* reconstructed = recons + (
id - i0) * d;
596 float *distances, idx_t *labels,
599 idx_t * idx =
new idx_t [n *
nprobe];
601 float * coarse_dis =
new float [n *
nprobe];
611 distances, labels,
true );
612 for (idx_t i = 0; i < n; ++i) {
613 for (idx_t j = 0; j < k; ++j) {
614 idx_t ij = i * k + j;
615 idx_t key = labels[ij];
616 float* reconstructed = recons + ij *
d;
619 memset(reconstructed, -1,
sizeof(*reconstructed) * d);
621 int list_no = key >> 32;
622 int offset = key & 0xffffffff;
637 FAISS_THROW_MSG (
"reconstruct_from_offset not implemented");
651 "direct map remove not implemented");
653 std::vector<idx_t> toremove(
nlist);
655 #pragma omp parallel for
656 for (idx_t i = 0; i <
nlist; i++) {
660 if (sel.is_member (idsi[j])) {
670 toremove[i] = l0 - l;
674 for (idx_t i = 0; i <
nlist; i++) {
675 if (toremove[i] > 0) {
676 nremove += toremove[i];
691 printf (
"Training level-1 quantizer\n");
696 printf (
"Training IVF residual\n");
705 printf(
"IndexIVF: no residual training\n");
713 FAISS_THROW_IF_NOT (other.
d == d);
716 FAISS_THROW_IF_NOT_MSG (
typeid (*
this) ==
typeid (other),
717 "can only merge indexes of the same type");
726 "direct map copy not implemented");
749 idx_t a1, idx_t a2)
const
755 FAISS_THROW_IF_NOT_FMT (
756 subset_type == 0 || subset_type == 1 || subset_type == 2,
757 "subset type %d not implemented", subset_type);
765 for (idx_t list_no = 0; list_no <
nlist; list_no++) {
769 if (subset_type == 0) {
770 for (idx_t i = 0; i < n; i++) {
771 idx_t
id = ids_in[i];
772 if (a1 <=
id &&
id < a2) {
779 }
else if (subset_type == 1) {
780 for (idx_t i = 0; i < n; i++) {
781 idx_t
id = ids_in[i];
789 }
else if (subset_type == 2) {
791 size_t next_accu_n = accu_n + n;
792 size_t next_accu_a1 = next_accu_n * a1 /
ntotal;
793 size_t i1 = next_accu_a1 - accu_a1;
794 size_t next_accu_a2 = next_accu_n * a2 /
ntotal;
795 size_t i2 = next_accu_a2 - accu_a2;
797 for (idx_t i = i1; i < i2; i++) {
804 accu_a1 = next_accu_a1;
805 accu_a2 = next_accu_a2;
809 FAISS_ASSERT(accu_n ==
ntotal);
814 IndexIVF::~IndexIVF()
822 void IndexIVFStats::reset()
824 memset ((
void*)
this, 0,
sizeof (*
this));
828 IndexIVFStats indexIVF_stats;
836 FAISS_THROW_MSG (
"scan_codes_range not implemented");
virtual void encode_vectors(idx_t n, const float *x, const idx_t *list_nos, uint8_t *codes) const =0
virtual void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const IVFSearchParameters *params=nullptr) const
result structure for a single query
simple (default) implementation as an array of inverted lists
void check_compatible_for_merge(const IndexIVF &other) const
void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons) const override
virtual void reset()=0
removes all elements from the database.
size_t nprobe
number of probes at query time
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
idx_t remove_ids(const IDSelector &sel) override
Dataset manipulation functions.
virtual size_t list_size(size_t list_no) const =0
get the size of a list
void reconstruct(idx_t key, float *recons) const override
virtual void train(idx_t n, const float *x)
virtual void train_residual(idx_t n, const float *x)
void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result) const override
void merge_from(InvertedLists *oivf, size_t add_id)
move all entries from oivf (empty on output)
virtual idx_t get_single_id(size_t list_no, size_t offset) const
long idx_t
all indices are this type
size_t code_size
code size per vector in bytes
virtual void copy_subset_to(IndexIVF &other, int subset_type, idx_t a1, idx_t a2) const
void train(idx_t n, const float *x) override
Trains the quantizer and calls train_residual to train sub-quantizers.
void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
char quantizer_trains_alone
virtual void add(idx_t n, const float *x)=0
virtual void set_list(idx_t list_no, float coarse_dis)=0
following codes come from this inverted list
void add_with_ids(idx_t n, const float *x, const idx_t *xids) override
default implementation that calls encode_vectors
virtual size_t add_entry(size_t list_no, idx_t theid, const uint8_t *code)
add one entry to an inverted list
virtual InvertedListScanner * get_InvertedListScanner(bool store_pairs=false) const
get a scanner for this index (store_pairs means ignore labels)
virtual size_t scan_codes(size_t n, const uint8_t *codes, const idx_t *ids, float *distances, idx_t *labels, size_t k) const =0
void replace_invlists(InvertedLists *il, bool own=false)
replace the inverted lists, old one is deallocated if own_invlists
ClusteringParameters cp
to override default clustering params
idx_t ntotal
total nb of indexed vectors
bool verbose
verbosity level
void reset() override
removes all elements from the database.
double getmillisecs()
ms elapsed since some arbitrary epoch
std::vector< float > centroids
centroids (k * d)
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const =0
Index * clustering_index
to override index used during clustering
void train_q1(size_t n, const float *x, bool verbose, MetricType metric_type)
Trains the quantizer and calls train_residual to train sub-quantizers.
size_t nprobe
number of probes at query time
the entries in the buffers are split per query
size_t nlist
number of possible key values
static void merge(std::vector< RangeSearchPartialResult * > &partial_results, bool do_delete=true)
void make_direct_map(bool new_maintain_direct_map=true)
MetricType metric_type
type of metric this index uses for search
InvertedLists * invlists
Acess to the actual data.
virtual void reconstruct_from_offset(idx_t list_no, idx_t offset, float *recons) const
static size_t get_period_hint(size_t flops)
void add(idx_t n, const float *x) override
Calls add_with_ids with NULL ids.
virtual void train(idx_t n, const float *x, faiss::Index &index)
Index is used during the assignment stage.
size_t max_codes
max nb of codes to visit to do a query
virtual void prefetch_lists(const idx_t *list_nos, int nlist) const
Index * quantizer
quantizer that maps vectors to inverted lists
bool is_trained
set if the Index does not require training, or if training is done already
size_t max_codes
max nb of codes to visit to do a query
bool maintain_direct_map
map for direct access to the elements. Enables reconstruct().
bool spherical
do we want normalized centroids?
bool own_fields
whether object owns the quantizer
virtual void set_query(const float *query_vector)=0
from now on we handle this query.
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
virtual void merge_from(IndexIVF &other, idx_t add_id)
size_t nlist
number of possible key values
size_t code_size
code size per vector in bytes
virtual void scan_codes_range(size_t n, const uint8_t *codes, const idx_t *ids, float radius, RangeQueryResult &result) const
MetricType
Some algorithms support both an inner product version and a L2 search version.