forked from facebookresearch/faiss
-
Notifications
You must be signed in to change notification settings - Fork 0
/
IndexBinary.h
163 lines (132 loc) · 4.91 KB
/
IndexBinary.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
/**
* Copyright (c) Facebook, Inc. and its affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
// -*- c++ -*-
#ifndef FAISS_INDEX_BINARY_H
#define FAISS_INDEX_BINARY_H
#include <cstdio>
#include <typeinfo>
#include <string>
#include <sstream>
#include <faiss/impl/FaissAssert.h>
#include <faiss/Index.h>
namespace faiss {
/// Forward declarations see AuxIndexStructures.h
struct IDSelector;
struct RangeSearchResult;
/** Abstract structure for a binary index.
*
* Supports adding vertices and searching them.
*
* All queries are symmetric because there is no distinction between codes and
* vectors.
*/
struct IndexBinary {
using idx_t = Index::idx_t; ///< all indices are this type
using component_t = uint8_t;
using distance_t = int32_t;
int d; ///< vector dimension
int code_size; ///< number of bytes per vector ( = d / 8 )
idx_t ntotal; ///< total nb of indexed vectors
bool verbose; ///< verbosity level
/// set if the Index does not require training, or if training is done already
bool is_trained;
/// type of metric this index uses for search
MetricType metric_type;
explicit IndexBinary(idx_t d = 0, MetricType metric = METRIC_L2)
: d(d),
code_size(d / 8),
ntotal(0),
verbose(false),
is_trained(true),
metric_type(metric) {
FAISS_THROW_IF_NOT(d % 8 == 0);
}
virtual ~IndexBinary();
/** Perform training on a representative set of vectors.
*
* @param n nb of training vectors
* @param x training vecors, size n * d / 8
*/
virtual void train(idx_t n, const uint8_t *x);
/** Add n vectors of dimension d to the index.
*
* Vectors are implicitly assigned labels ntotal .. ntotal + n - 1
* @param x input matrix, size n * d / 8
*/
virtual void add(idx_t n, const uint8_t *x) = 0;
/** Same as add, but stores xids instead of sequential ids.
*
* The default implementation fails with an assertion, as it is
* not supported by all indexes.
*
* @param xids if non-null, ids to store for the vectors (size n)
*/
virtual void add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids);
/** Query n vectors of dimension d to the index.
*
* return at most k vectors. If there are not enough results for a
* query, the result array is padded with -1s.
*
* @param x input vectors to search, size n * d / 8
* @param labels output labels of the NNs, size n*k
* @param distances output pairwise distances, size n*k
*/
virtual void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels) const = 0;
/** Query n vectors of dimension d to the index.
*
* return all vectors with distance < radius. Note that many
* indexes do not implement the range_search (only the k-NN search
* is mandatory).
*
* @param x input vectors to search, size n * d / 8
* @param radius search radius
* @param result result table
*/
virtual void range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result) const;
/** Return the indexes of the k vectors closest to the query x.
*
* This function is identical to search but only returns labels of neighbors.
* @param x input vectors to search, size n * d / 8
* @param labels output labels of the NNs, size n*k
*/
void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1);
/// Removes all elements from the database.
virtual void reset() = 0;
/** Removes IDs from the index. Not supported by all indexes.
*/
virtual size_t remove_ids(const IDSelector& sel);
/** Reconstruct a stored vector.
*
* This function may not be defined for some indexes.
* @param key id of the vector to reconstruct
* @param recons reconstucted vector (size d / 8)
*/
virtual void reconstruct(idx_t key, uint8_t *recons) const;
/** Reconstruct vectors i0 to i0 + ni - 1.
*
* This function may not be defined for some indexes.
* @param recons reconstucted vectors (size ni * d / 8)
*/
virtual void reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const;
/** Similar to search, but also reconstructs the stored vectors (or an
* approximation in the case of lossy coding) for the search results.
*
* If there are not enough results for a query, the resulting array
* is padded with -1s.
*
* @param recons reconstructed vectors size (n, k, d)
**/
virtual void search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
uint8_t *recons) const;
/** Display the actual class name and some more info. */
void display() const;
};
} // namespace faiss
#endif // FAISS_INDEX_BINARY_H