forked from facebookresearch/faiss
-
Notifications
You must be signed in to change notification settings - Fork 0
/
MetaIndexes.h
166 lines (113 loc) · 4.09 KB
/
MetaIndexes.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
// -*- c++ -*-
#ifndef META_INDEXES_H
#define META_INDEXES_H
#include <vector>
#include <unordered_map>
#include "Index.h"
namespace faiss {
/** Index that translates search results to ids */
struct IndexIDMap : Index {
Index * index; ///! the sub-index
bool own_fields; ///! whether pointers are deleted in destructo
std::vector<long> id_map;
explicit IndexIDMap (Index *index);
/// Same as add_core, but stores xids instead of sequential ids
/// @param xids if non-null, ids to store for the vectors (size n)
void add_with_ids(idx_t n, const float* x, const long* xids) override;
/// this will fail. Use add_with_ids
void add(idx_t n, const float* x) override;
void search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels) const override;
void train(idx_t n, const float* x) override;
void reset() override;
/// remove ids adapted to IndexFlat
long remove_ids(const IDSelector& sel) override;
void range_search (idx_t n, const float *x, float radius,
RangeSearchResult *result) const override;
~IndexIDMap() override;
IndexIDMap () {own_fields=false; index=nullptr; }
};
/** same as IndexIDMap but also provides an efficient reconstruction
implementation via a 2-way index */
struct IndexIDMap2 : IndexIDMap {
std::unordered_map<idx_t, idx_t> rev_map;
explicit IndexIDMap2 (Index *index);
/// make the rev_map from scratch
void construct_rev_map ();
void add_with_ids(idx_t n, const float* x, const long* xids) override;
long remove_ids(const IDSelector& sel) override;
void reconstruct (idx_t key, float * recons) const override;
~IndexIDMap2() override {}
IndexIDMap2 () {}
};
/** Index that concatenates the results from several sub-indexes
*
*/
struct IndexShards : Index {
std::vector<Index*> shard_indexes;
bool own_fields; /// should the sub-indexes be deleted along with this?
bool threaded;
bool successive_ids;
/**
* @param threaded do we use one thread per sub_index or do
* queries sequentially?
* @param successive_ids should we shift the returned ids by
* the size of each sub-index or return them
* as they are?
*/
explicit IndexShards (idx_t d, bool threaded = false,
bool successive_ids = true);
void add_shard (Index *);
// update metric_type and ntotal. Call if you changes something in
// the shard indexes.
void sync_with_shard_indexes ();
Index *at(int i) {return shard_indexes[i]; }
/// supported only for sub-indices that implement add_with_ids
void add(idx_t n, const float* x) override;
void add_with_ids(idx_t n, const float* x, const long* xids) override;
void search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels) const override;
void train(idx_t n, const float* x) override;
void reset() override;
~IndexShards() override;
};
/** splits input vectors in segments and assigns each segment to a sub-index
* used to distribute a MultiIndexQuantizer
*/
struct IndexSplitVectors: Index {
bool own_fields;
bool threaded;
std::vector<Index*> sub_indexes;
idx_t sum_d; /// sum of dimensions seen so far
explicit IndexSplitVectors (idx_t d, bool threaded = false);
void add_sub_index (Index *);
void sync_with_sub_indexes ();
void add(idx_t n, const float* x) override;
void search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels) const override;
void train(idx_t n, const float* x) override;
void reset() override;
~IndexSplitVectors() override;
};
}
#endif