usearch 0.1.9 → 0.2.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/.vscode/settings.json +6 -0
- package/CMakeLists.txt +26 -8
- package/CMakeListsInstructions.md +49 -33
- package/Cargo.toml +2 -2
- package/LICENSE +201 -0
- package/README.md +30 -19
- package/VERSION +1 -1
- package/assets/usearch-neighbor-types.png +0 -0
- package/build.gradle +1 -1
- package/include/usearch/usearch.hpp +26 -1
- package/package.json +2 -2
- package/python/bench.py +55 -73
- package/python/lib.cpp +365 -49
- package/python/test.py +35 -2
- package/python/usearch/client.py +34 -0
- package/python/usearch/io.py +67 -0
- package/python/usearch/server.py +51 -0
- package/setup.py +1 -1
- package/simsimd/.vscode/settings.json +3 -1
- package/simsimd/include/simsimd/simsimd.h +14 -7
- package/src/advanced.hpp +86 -17
- package/src/bench.cpp +45 -2
- package/wolfram/CMakeLists.txt +17 -0
- package/wolfram/lib.cpp +139 -4
- package/wolfram/test.wls +27 -0
- package/docs/Makefile +0 -20
package/python/bench.py
CHANGED
|
@@ -1,59 +1,12 @@
|
|
|
1
1
|
import datetime
|
|
2
2
|
import numpy as np
|
|
3
|
+
import fire
|
|
3
4
|
|
|
4
5
|
import faiss
|
|
5
6
|
import usearch
|
|
6
7
|
|
|
7
|
-
|
|
8
|
-
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
def read_matrix(filename: str, start_row: int = 0, count_rows: int = None):
|
|
12
|
-
"""
|
|
13
|
-
Read *.ibin, *.hbin, *.fbin, *.dbin files with matrixes.
|
|
14
|
-
Args:
|
|
15
|
-
:param filename (str): path to the matrix file
|
|
16
|
-
:param start_row (int): start reading vectors from this index
|
|
17
|
-
:param count_rows (int): number of vectors to read. If None, read all vectors
|
|
18
|
-
Returns:
|
|
19
|
-
Parsed matrix (numpy.ndarray)
|
|
20
|
-
"""
|
|
21
|
-
dtype = np.float32
|
|
22
|
-
scalar_size = 4
|
|
23
|
-
if filename.endswith('.fbin'):
|
|
24
|
-
dtype = np.float32
|
|
25
|
-
scalar_size = 4
|
|
26
|
-
elif filename.endswith('.dbin'):
|
|
27
|
-
dtype = np.float64
|
|
28
|
-
scalar_size = 8
|
|
29
|
-
elif filename.endswith('.hbin'):
|
|
30
|
-
dtype = np.float16
|
|
31
|
-
scalar_size = 2
|
|
32
|
-
elif filename.endswith('.ibin'):
|
|
33
|
-
dtype = np.int32
|
|
34
|
-
scalar_size = 4
|
|
35
|
-
else:
|
|
36
|
-
raise Exception('Unknown file type')
|
|
37
|
-
with open(filename, 'rb') as f:
|
|
38
|
-
rows, cols = np.fromfile(f, count=2, dtype=np.int32)
|
|
39
|
-
rows = (rows - start_row) if count_rows is None else count_rows
|
|
40
|
-
arr = np.fromfile(
|
|
41
|
-
f, count=rows * cols, dtype=dtype,
|
|
42
|
-
offset=start_row * scalar_size * cols)
|
|
43
|
-
return arr.reshape(rows, cols)
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
print('Will read datasets!')
|
|
47
|
-
xb = read_matrix('datasets/wiki_1M/base.1M.fbin')
|
|
48
|
-
xq = read_matrix('datasets/wiki_1M/query.public.100K.fbin')
|
|
49
|
-
|
|
50
|
-
d = 256 # vector size
|
|
51
|
-
M = 16
|
|
52
|
-
efSearch = 32 # number of entry points (neighbors) we use on each layer
|
|
53
|
-
efConstruction = 128 # number of entry points used on each layer
|
|
54
|
-
labels = np.arange(xb.shape[0], dtype=np.longlong)
|
|
55
|
-
print('Will benchmark algorithms!')
|
|
56
|
-
|
|
8
|
+
import struct
|
|
9
|
+
import numpy as np
|
|
57
10
|
|
|
58
11
|
def measure(f) -> float:
|
|
59
12
|
a = datetime.datetime.now()
|
|
@@ -64,29 +17,58 @@ def measure(f) -> float:
|
|
|
64
17
|
return c.seconds
|
|
65
18
|
|
|
66
19
|
|
|
67
|
-
def
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
20
|
+
def bench_faiss(index, vectors: np.array, queries: np.array, neighbors: np.array):
|
|
21
|
+
dt = measure(lambda: index.add(vectors))
|
|
22
|
+
print(f'- FAISS: {vectors.shape[0]/dt:.2f} vectors/s')
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def bench_usearch(index, vectors: np.array, queries: np.array, neighbors: np.array):
|
|
26
|
+
labels = np.arange(vectors.shape[0], dtype=np.longlong)
|
|
27
|
+
assert len(index) == 0
|
|
28
|
+
dt = measure(lambda: index.add(labels, vectors))
|
|
29
|
+
assert len(index) == vectors.shape[0]
|
|
30
|
+
print(f'- USearch: {vectors.shape[0]/dt:.2f} vectors/s')
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def main(
|
|
34
|
+
vectors: str,
|
|
35
|
+
queries: str,
|
|
36
|
+
neighbors: str,
|
|
37
|
+
connectivity: int = 16,
|
|
38
|
+
expansion_add: int = 128,
|
|
39
|
+
expansion_search: int = 64,
|
|
40
|
+
):
|
|
41
|
+
|
|
42
|
+
vectors_mat = load_matrix(vectors)
|
|
43
|
+
queries_mat = load_matrix(queries)
|
|
44
|
+
neighbors_mat = load_matrix(neighbors)
|
|
45
|
+
dim = vectors_mat.shape[1]
|
|
46
|
+
|
|
47
|
+
# Test FAISS
|
|
48
|
+
index = faiss.IndexHNSWFlat(dim, connectivity)
|
|
49
|
+
index.hnsw.efSearch = expansion_search
|
|
50
|
+
index.hnsw.efConstruction = expansion_add
|
|
51
|
+
bench_faiss(index, vectors_mat, queries_mat, neighbors_mat)
|
|
52
|
+
|
|
53
|
+
# Test FAISS with default Product Quantization
|
|
54
|
+
nlist = 10000 # Number of inverted lists (number of partitions or cells).
|
|
55
|
+
nsegment = 16 # Number of segments for PQ (number of subquantizers).
|
|
56
|
+
nbit = 8 # Number of bits to encode each segment.
|
|
57
|
+
coarse_quantizer = faiss.IndexHNSWFlat(dim, connectivity)
|
|
58
|
+
index = faiss.IndexIVFPQ(coarse_quantizer, dim, nlist, nsegment, nbit)
|
|
59
|
+
index.train(vectors_mat)
|
|
60
|
+
index.nprobe = 10
|
|
61
|
+
bench_faiss(index, vectors_mat, queries_mat, neighbors_mat)
|
|
62
|
+
|
|
63
|
+
# Test USearch `Index`
|
|
64
|
+
index = usearch.Index(
|
|
65
|
+
ndim=dim,
|
|
66
|
+
expansion_add=expansion_add,
|
|
67
|
+
expansion_search=expansion_search,
|
|
68
|
+
connectivity=connectivity,
|
|
79
69
|
)
|
|
80
|
-
|
|
81
|
-
dt_unum = measure(lambda: index_unum.add(labels, xb))
|
|
82
|
-
assert len(index_unum) == xb.shape[0]
|
|
83
|
-
|
|
84
|
-
return dt_meta, dt_unum
|
|
70
|
+
bench_usearch(index, vectors_mat, queries_mat, neighbors_mat)
|
|
85
71
|
|
|
86
72
|
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
dt_meta = sum(x[0] for x in durations)
|
|
90
|
-
dt_unum = sum(x[1] for x in durations)
|
|
91
|
-
print(f'- FAISS vectors per second: {xb.shape[0]*experiments/dt_meta:.2f}')
|
|
92
|
-
print(f'- USearch vectors per second: {xb.shape[0]*experiments/dt_unum:.2f}')
|
|
73
|
+
if __name__ == '__main__':
|
|
74
|
+
fire.Fire(main)
|
package/python/lib.cpp
CHANGED
|
@@ -24,11 +24,41 @@ using namespace unum::usearch;
|
|
|
24
24
|
using namespace unum;
|
|
25
25
|
|
|
26
26
|
namespace py = pybind11;
|
|
27
|
+
using py_shape_t = py::array::ShapeContainer;
|
|
27
28
|
|
|
28
29
|
using label_t = Py_ssize_t;
|
|
29
30
|
using distance_t = punned_distance_t;
|
|
31
|
+
using id_t = std::uint32_t;
|
|
30
32
|
using native_index_t = auto_index_gt<label_t>;
|
|
31
33
|
|
|
34
|
+
using set_member_t = std::uint32_t;
|
|
35
|
+
using set_view_t = span_gt<set_member_t const>;
|
|
36
|
+
using sets_index_t = index_gt<jaccard_gt<set_member_t>, label_t, id_t, set_member_t>;
|
|
37
|
+
|
|
38
|
+
using hash_word_t = std::uint64_t;
|
|
39
|
+
using hash_index_t = index_gt<bit_hamming_gt<hash_word_t>, label_t, id_t, hash_word_t>;
|
|
40
|
+
static constexpr std::size_t bits_per_hash_word_k = sizeof(hash_word_t) * CHAR_BIT;
|
|
41
|
+
|
|
42
|
+
struct hash_index_w_meta_t : public hash_index_t {
|
|
43
|
+
using hash_index_t::add;
|
|
44
|
+
using hash_index_t::capacity;
|
|
45
|
+
using hash_index_t::reserve;
|
|
46
|
+
using hash_index_t::search;
|
|
47
|
+
using hash_index_t::size;
|
|
48
|
+
|
|
49
|
+
std::vector<hash_word_t> buffer_;
|
|
50
|
+
std::size_t words_;
|
|
51
|
+
std::size_t bits_;
|
|
52
|
+
|
|
53
|
+
hash_index_w_meta_t(config_t config, std::size_t bits) : hash_index_t(config) {
|
|
54
|
+
words_ = divide_round_up<bits_per_hash_word_k>(bits);
|
|
55
|
+
bits_ = words_ * bits_per_hash_word_k;
|
|
56
|
+
buffer_.resize(words_);
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
span_gt<hash_word_t const> buffer() const noexcept { return {buffer_.data(), buffer_.size()}; }
|
|
60
|
+
};
|
|
61
|
+
|
|
32
62
|
static native_index_t make_index( //
|
|
33
63
|
std::size_t dimensions, //
|
|
34
64
|
std::size_t capacity, //
|
|
@@ -36,7 +66,8 @@ static native_index_t make_index( //
|
|
|
36
66
|
std::string const& metric, //
|
|
37
67
|
std::size_t connectivity, //
|
|
38
68
|
std::size_t expansion_add, //
|
|
39
|
-
std::size_t expansion_search
|
|
69
|
+
std::size_t expansion_search, //
|
|
70
|
+
std::size_t metric_uintptr //
|
|
40
71
|
) {
|
|
41
72
|
|
|
42
73
|
config_t config;
|
|
@@ -48,36 +79,98 @@ static native_index_t make_index( //
|
|
|
48
79
|
config.max_threads_search = std::thread::hardware_concurrency();
|
|
49
80
|
|
|
50
81
|
accuracy_t accuracy = accuracy_from_name(scalar_type.c_str(), scalar_type.size());
|
|
51
|
-
|
|
82
|
+
punned_metric_t metric_ptr = reinterpret_cast<punned_metric_t>(metric_uintptr);
|
|
83
|
+
if (metric_ptr)
|
|
84
|
+
return native_index_t::udf(dimensions, metric_ptr, accuracy, config);
|
|
85
|
+
else
|
|
86
|
+
return index_from_name<native_index_t>(metric.c_str(), metric.size(), dimensions, accuracy, config);
|
|
87
|
+
}
|
|
88
|
+
|
|
89
|
+
static std::unique_ptr<sets_index_t> make_sets_index( //
|
|
90
|
+
std::size_t capacity, //
|
|
91
|
+
std::size_t connectivity, //
|
|
92
|
+
std::size_t expansion_add, //
|
|
93
|
+
std::size_t expansion_search //
|
|
94
|
+
) {
|
|
95
|
+
config_t config;
|
|
96
|
+
config.expansion_add = expansion_add;
|
|
97
|
+
config.expansion_search = expansion_search;
|
|
98
|
+
config.connectivity = connectivity;
|
|
99
|
+
config.max_elements = capacity;
|
|
100
|
+
config.max_threads_add = 1;
|
|
101
|
+
config.max_threads_search = 1;
|
|
102
|
+
|
|
103
|
+
return std::unique_ptr<sets_index_t>(new sets_index_t(config));
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
static std::unique_ptr<hash_index_w_meta_t> make_hash_index( //
|
|
107
|
+
std::size_t bits, //
|
|
108
|
+
std::size_t capacity, //
|
|
109
|
+
std::size_t connectivity, //
|
|
110
|
+
std::size_t expansion_add, //
|
|
111
|
+
std::size_t expansion_search //
|
|
112
|
+
) {
|
|
113
|
+
config_t config;
|
|
114
|
+
config.expansion_add = expansion_add;
|
|
115
|
+
config.expansion_search = expansion_search;
|
|
116
|
+
config.connectivity = connectivity;
|
|
117
|
+
config.max_elements = capacity;
|
|
118
|
+
config.max_threads_add = 1;
|
|
119
|
+
config.max_threads_search = 1;
|
|
120
|
+
|
|
121
|
+
return std::unique_ptr<hash_index_w_meta_t>(new hash_index_w_meta_t(config, bits));
|
|
52
122
|
}
|
|
53
123
|
|
|
54
|
-
static void
|
|
124
|
+
static void add_one_to_index(native_index_t& index, label_t label, py::buffer vector, bool copy) {
|
|
125
|
+
|
|
126
|
+
py::buffer_info vector_info = vector.request();
|
|
127
|
+
if (vector_info.ndim != 1)
|
|
128
|
+
throw std::invalid_argument("Expects a vector, not a higher-rank tensor!");
|
|
129
|
+
|
|
130
|
+
ssize_t vector_dimensions = vector_info.shape[0];
|
|
131
|
+
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
132
|
+
if (vector_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
133
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
134
|
+
|
|
135
|
+
if (index.size() + 1 >= index.capacity())
|
|
136
|
+
index.reserve(ceil2(index.size() + 1));
|
|
137
|
+
|
|
138
|
+
// https://docs.python.org/3/library/struct.html#format-characters
|
|
139
|
+
if (vector_info.format == "e")
|
|
140
|
+
index.add(label, reinterpret_cast<f16_converted_t const*>(vector_data), 0, copy);
|
|
141
|
+
else if (vector_info.format == "f")
|
|
142
|
+
index.add(label, reinterpret_cast<float const*>(vector_data), 0, copy);
|
|
143
|
+
else if (vector_info.format == "d")
|
|
144
|
+
index.add(label, reinterpret_cast<double const*>(vector_data), 0, copy);
|
|
145
|
+
else
|
|
146
|
+
throw std::invalid_argument("Incompatible scalars in the vector!");
|
|
147
|
+
}
|
|
148
|
+
|
|
149
|
+
static void add_many_to_index(native_index_t& index, py::buffer labels, py::buffer vectors, bool copy) {
|
|
55
150
|
|
|
56
151
|
py::buffer_info labels_info = labels.request();
|
|
57
152
|
py::buffer_info vectors_info = vectors.request();
|
|
58
153
|
|
|
59
154
|
if (labels_info.format != py::format_descriptor<label_t>::format())
|
|
60
|
-
throw std::
|
|
155
|
+
throw std::invalid_argument("Incompatible label type!");
|
|
61
156
|
|
|
62
157
|
if (labels_info.ndim != 1)
|
|
63
|
-
throw std::
|
|
158
|
+
throw std::invalid_argument("Labels must be placed in a single-dimensional array!");
|
|
64
159
|
|
|
65
160
|
if (vectors_info.ndim != 2)
|
|
66
|
-
throw std::
|
|
161
|
+
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
67
162
|
|
|
68
163
|
ssize_t labels_count = labels_info.shape[0];
|
|
69
164
|
ssize_t vectors_count = vectors_info.shape[0];
|
|
70
165
|
ssize_t vectors_dimensions = vectors_info.shape[1];
|
|
71
|
-
if (vectors_dimensions != index.dimensions())
|
|
72
|
-
throw std::
|
|
166
|
+
if (vectors_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
167
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
73
168
|
|
|
74
169
|
if (labels_count != vectors_count)
|
|
75
|
-
throw std::
|
|
170
|
+
throw std::invalid_argument("Number of labels and vectors must match!");
|
|
76
171
|
|
|
77
|
-
if (index.size() + vectors_count >= index.capacity())
|
|
78
|
-
|
|
79
|
-
index.reserve(next_capacity);
|
|
80
|
-
}
|
|
172
|
+
if (index.size() + vectors_count >= index.capacity())
|
|
173
|
+
index.reserve(ceil2(index.size() + vectors_count));
|
|
81
174
|
|
|
82
175
|
char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
|
|
83
176
|
char const* labels_data = reinterpret_cast<char const*>(labels_info.ptr);
|
|
@@ -103,11 +196,48 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
|
|
|
103
196
|
index.add(label, vector, thread_idx, copy);
|
|
104
197
|
});
|
|
105
198
|
else
|
|
106
|
-
throw std::
|
|
199
|
+
throw std::invalid_argument("Incompatible scalars in the vectors matrix!");
|
|
200
|
+
}
|
|
201
|
+
|
|
202
|
+
static py::tuple search_one_in_index(native_index_t& index, py::buffer vector, std::size_t wanted) {
|
|
203
|
+
|
|
204
|
+
py::buffer_info vector_info = vector.request();
|
|
205
|
+
ssize_t vector_dimensions = vector_info.shape[0];
|
|
206
|
+
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
207
|
+
if (vector_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
208
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
209
|
+
|
|
210
|
+
py::array_t<label_t> labels_py(static_cast<ssize_t>(wanted));
|
|
211
|
+
py::array_t<distance_t> distances_py(static_cast<ssize_t>(wanted));
|
|
212
|
+
std::size_t count{};
|
|
213
|
+
auto labels_py1d = labels_py.mutable_unchecked<1>();
|
|
214
|
+
auto distances_py1d = distances_py.mutable_unchecked<1>();
|
|
215
|
+
|
|
216
|
+
// https://docs.python.org/3/library/struct.html#format-characters
|
|
217
|
+
if (vector_info.format == "e")
|
|
218
|
+
count = index.search( //
|
|
219
|
+
reinterpret_cast<f16_converted_t const*>(vector_data), wanted, &labels_py1d(0), &distances_py1d(0), 0);
|
|
220
|
+
else if (vector_info.format == "f")
|
|
221
|
+
count = index.search( //
|
|
222
|
+
reinterpret_cast<float const*>(vector_data), wanted, &labels_py1d(0), &distances_py1d(0), 0);
|
|
223
|
+
else if (vector_info.format == "d")
|
|
224
|
+
count = index.search( //
|
|
225
|
+
reinterpret_cast<double const*>(vector_data), wanted, &labels_py1d(0), &distances_py1d(0), 0);
|
|
226
|
+
else
|
|
227
|
+
throw std::invalid_argument("Incompatible scalars in the query vector!");
|
|
228
|
+
|
|
229
|
+
labels_py.resize(py_shape_t{static_cast<ssize_t>(count)});
|
|
230
|
+
distances_py.resize(py_shape_t{static_cast<ssize_t>(count)});
|
|
231
|
+
|
|
232
|
+
py::tuple results(3);
|
|
233
|
+
results[0] = labels_py;
|
|
234
|
+
results[1] = distances_py;
|
|
235
|
+
results[2] = static_cast<Py_ssize_t>(count);
|
|
236
|
+
return results;
|
|
107
237
|
}
|
|
108
238
|
|
|
109
239
|
/**
|
|
110
|
-
* @param vectors Matrix of vectors to search for
|
|
240
|
+
* @param vectors Matrix of vectors to search for.
|
|
111
241
|
* @param wanted Number of matches per request.
|
|
112
242
|
*
|
|
113
243
|
* @return Tuple with:
|
|
@@ -115,26 +245,30 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
|
|
|
115
245
|
* 2. matrix of distances,
|
|
116
246
|
* 3. array with match counts.
|
|
117
247
|
*/
|
|
118
|
-
static py::tuple
|
|
248
|
+
static py::tuple search_many_in_index(native_index_t& index, py::buffer vectors, std::size_t wanted) {
|
|
249
|
+
|
|
250
|
+
if (wanted == 0)
|
|
251
|
+
return py::tuple(3);
|
|
119
252
|
|
|
120
253
|
py::buffer_info vectors_info = vectors.request();
|
|
254
|
+
if (vectors_info.ndim == 1)
|
|
255
|
+
return search_one_in_index(index, vectors, wanted);
|
|
121
256
|
if (vectors_info.ndim != 2)
|
|
122
|
-
throw std::
|
|
257
|
+
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
123
258
|
|
|
124
259
|
ssize_t vectors_count = vectors_info.shape[0];
|
|
125
260
|
ssize_t vectors_dimensions = vectors_info.shape[1];
|
|
126
|
-
|
|
127
|
-
|
|
261
|
+
char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
|
|
262
|
+
if (vectors_dimensions != static_cast<ssize_t>(index.dimensions()))
|
|
263
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
128
264
|
|
|
129
|
-
py::array_t<label_t> labels_py({vectors_count, wanted});
|
|
130
|
-
py::array_t<distance_t> distances_py({vectors_count, wanted});
|
|
265
|
+
py::array_t<label_t> labels_py({vectors_count, static_cast<ssize_t>(wanted)});
|
|
266
|
+
py::array_t<distance_t> distances_py({vectors_count, static_cast<ssize_t>(wanted)});
|
|
131
267
|
py::array_t<Py_ssize_t> counts_py(vectors_count);
|
|
132
268
|
auto labels_py2d = labels_py.mutable_unchecked<2>();
|
|
133
269
|
auto distances_py2d = distances_py.mutable_unchecked<2>();
|
|
134
270
|
auto counts_py1d = counts_py.mutable_unchecked<1>();
|
|
135
271
|
|
|
136
|
-
char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
|
|
137
|
-
|
|
138
272
|
// https://docs.python.org/3/library/struct.html#format-characters
|
|
139
273
|
if (vectors_info.format == "e")
|
|
140
274
|
multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
@@ -155,7 +289,7 @@ static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssiz
|
|
|
155
289
|
index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
|
|
156
290
|
});
|
|
157
291
|
else
|
|
158
|
-
throw std::
|
|
292
|
+
throw std::invalid_argument("Incompatible scalars in the query matrix!");
|
|
159
293
|
|
|
160
294
|
py::tuple results(3);
|
|
161
295
|
results[0] = labels_py;
|
|
@@ -164,38 +298,116 @@ static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssiz
|
|
|
164
298
|
return results;
|
|
165
299
|
}
|
|
166
300
|
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
301
|
+
// clang-format off
|
|
302
|
+
template <typename index_at> void save_index(index_at const& index, std::string const& path) { index.save(path.c_str()); }
|
|
303
|
+
template <typename index_at> void load_index(index_at& index, std::string const& path) { index.load(path.c_str()); }
|
|
304
|
+
template <typename index_at> void view_index(index_at& index, std::string const& path) { index.view(path.c_str()); }
|
|
305
|
+
template <typename index_at> void clear_index(index_at& index) { index.clear(); }
|
|
306
|
+
// clang-format on
|
|
307
|
+
|
|
308
|
+
template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
|
|
309
|
+
if (begin == end)
|
|
310
|
+
return false;
|
|
311
|
+
element_at const* last = begin;
|
|
312
|
+
begin++;
|
|
313
|
+
for (; begin != end; ++begin, ++last) {
|
|
314
|
+
if (*begin == *last)
|
|
315
|
+
return true;
|
|
316
|
+
}
|
|
317
|
+
return false;
|
|
318
|
+
}
|
|
319
|
+
|
|
320
|
+
void validate_set(py::array_t<set_member_t> const& set) {
|
|
321
|
+
if (set.ndim() != 1)
|
|
322
|
+
throw std::invalid_argument("Set can't be multi-dimensional!");
|
|
323
|
+
if (set.strides(0) != sizeof(set_member_t))
|
|
324
|
+
throw std::invalid_argument("Set can't be strided!");
|
|
325
|
+
auto proxy = set.unchecked<1>();
|
|
326
|
+
set_member_t const* begin = &proxy(0);
|
|
327
|
+
set_member_t const* end = begin + proxy.size();
|
|
328
|
+
if (!std::is_sorted(begin, end))
|
|
329
|
+
throw std::invalid_argument("Set must be sorted!");
|
|
330
|
+
if (has_duplicates(begin, end))
|
|
331
|
+
throw std::invalid_argument("Set must be deduplicated!");
|
|
332
|
+
}
|
|
333
|
+
|
|
334
|
+
inline std::uint64_t hash_ror64(std::uint64_t v, int r) noexcept { return (v >> r) | (v << (64 - r)); }
|
|
335
|
+
|
|
336
|
+
inline std::uint64_t hash(std::uint64_t v) noexcept {
|
|
337
|
+
v ^= hash_ror64(v, 25) ^ hash_ror64(v, 50);
|
|
338
|
+
v *= 0xA24BAED4963EE407UL;
|
|
339
|
+
v ^= hash_ror64(v, 24) ^ hash_ror64(v, 49);
|
|
340
|
+
v *= 0x9FB21C651E98DF25UL;
|
|
341
|
+
return v ^ v >> 28;
|
|
342
|
+
}
|
|
343
|
+
|
|
344
|
+
template <typename scalar_at>
|
|
345
|
+
inline void hash_typed_buffer(hash_index_w_meta_t& index, py::buffer_info const& vector_info) noexcept {
|
|
346
|
+
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
347
|
+
ssize_t vector_dimensions = vector_info.shape[0];
|
|
348
|
+
ssize_t vector_stride = vector_info.strides[0];
|
|
349
|
+
std::memset(index.buffer_.data(), 0, index.words_ * sizeof(hash_word_t));
|
|
350
|
+
|
|
351
|
+
for (std::size_t i = 0; i != static_cast<std::size_t>(vector_dimensions); ++i) {
|
|
352
|
+
scalar_at scalar = *reinterpret_cast<scalar_at const*>(vector_data + i * vector_stride);
|
|
353
|
+
std::uint64_t scalar_hash = hash(scalar);
|
|
354
|
+
index.buffer_[scalar_hash % index.words_] |= hash_word_t(1) << (scalar_hash % bits_per_hash_word_k);
|
|
355
|
+
}
|
|
356
|
+
}
|
|
357
|
+
|
|
358
|
+
void hash_buffer(hash_index_w_meta_t& index, py::buffer vector) {
|
|
359
|
+
py::buffer_info info = vector.request();
|
|
360
|
+
if (info.ndim != 1)
|
|
361
|
+
throw std::invalid_argument("Array can't be multi-dimensional!");
|
|
362
|
+
|
|
363
|
+
// https://docs.python.org/3/library/struct.html#format-characters
|
|
364
|
+
if (info.format == "h" || info.format == "H")
|
|
365
|
+
return hash_typed_buffer<std::uint16_t>(index, info);
|
|
366
|
+
else if (info.format == "i" || info.format == "I" || info.format == "l" || info.format == "L")
|
|
367
|
+
return hash_typed_buffer<std::uint32_t>(index, info);
|
|
368
|
+
else if (info.format == "q" || info.format == "Q" || info.format == "n" || info.format == "N")
|
|
369
|
+
return hash_typed_buffer<std::uint64_t>(index, info);
|
|
370
|
+
else
|
|
371
|
+
throw std::invalid_argument("Array elements must be 16, 32, or 64 bit hashable integers!");
|
|
372
|
+
}
|
|
170
373
|
|
|
171
374
|
PYBIND11_MODULE(usearch, m) {
|
|
172
375
|
m.doc() = "Unum USearch Python bindings";
|
|
173
376
|
|
|
174
377
|
auto i = py::class_<native_index_t>(m, "Index");
|
|
175
378
|
|
|
176
|
-
i.def(py::init(&make_index),
|
|
177
|
-
py::kw_only(),
|
|
178
|
-
py::arg("
|
|
179
|
-
py::arg("capacity") = 0,
|
|
180
|
-
py::arg("dtype") = std::string("f32"),
|
|
181
|
-
py::arg("metric") = std::string("ip"),
|
|
182
|
-
py::arg("connectivity") = config_t::connectivity_default_k,
|
|
183
|
-
py::arg("expansion_add") = config_t::expansion_add_default_k,
|
|
184
|
-
py::arg("expansion_search") = config_t::expansion_search_default_k //
|
|
379
|
+
i.def(py::init(&make_index), //
|
|
380
|
+
py::kw_only(), //
|
|
381
|
+
py::arg("ndim") = 0, //
|
|
382
|
+
py::arg("capacity") = 0, //
|
|
383
|
+
py::arg("dtype") = std::string("f32"), //
|
|
384
|
+
py::arg("metric") = std::string("ip"), //
|
|
385
|
+
py::arg("connectivity") = config_t::connectivity_default_k, //
|
|
386
|
+
py::arg("expansion_add") = config_t::expansion_add_default_k, //
|
|
387
|
+
py::arg("expansion_search") = config_t::expansion_search_default_k, //
|
|
388
|
+
py::arg("metric_pointer") = 0 //
|
|
185
389
|
);
|
|
186
390
|
|
|
187
|
-
i.def(
|
|
188
|
-
"add", &
|
|
189
|
-
py::arg("labels"),
|
|
190
|
-
py::arg("vectors"),
|
|
191
|
-
py::kw_only(),
|
|
192
|
-
py::arg("copy") = true
|
|
391
|
+
i.def( //
|
|
392
|
+
"add", &add_many_to_index, //
|
|
393
|
+
py::arg("labels"), //
|
|
394
|
+
py::arg("vectors"), //
|
|
395
|
+
py::kw_only(), //
|
|
396
|
+
py::arg("copy") = true //
|
|
397
|
+
);
|
|
398
|
+
|
|
399
|
+
i.def( //
|
|
400
|
+
"add", &add_one_to_index, //
|
|
401
|
+
py::arg("label"), //
|
|
402
|
+
py::arg("vector"), //
|
|
403
|
+
py::kw_only(), //
|
|
404
|
+
py::arg("copy") = true //
|
|
193
405
|
);
|
|
194
406
|
|
|
195
|
-
i.def(
|
|
196
|
-
"search", &
|
|
197
|
-
py::arg("
|
|
198
|
-
py::arg("count") = 10
|
|
407
|
+
i.def( //
|
|
408
|
+
"search", &search_many_in_index, //
|
|
409
|
+
py::arg("query"), //
|
|
410
|
+
py::arg("count") = 10 //
|
|
199
411
|
);
|
|
200
412
|
|
|
201
413
|
i.def("__len__", &native_index_t::size);
|
|
@@ -204,7 +416,111 @@ PYBIND11_MODULE(usearch, m) {
|
|
|
204
416
|
i.def_property_readonly("connectivity", &native_index_t::connectivity);
|
|
205
417
|
i.def_property_readonly("capacity", &native_index_t::capacity);
|
|
206
418
|
|
|
207
|
-
i.def("save", &save_index
|
|
208
|
-
i.def("load", &load_index
|
|
209
|
-
i.def("view", &view_index
|
|
419
|
+
i.def("save", &save_index<native_index_t>, py::arg("path"));
|
|
420
|
+
i.def("load", &load_index<native_index_t>, py::arg("path"));
|
|
421
|
+
i.def("view", &view_index<native_index_t>, py::arg("path"));
|
|
422
|
+
i.def("clear", &clear_index<native_index_t>);
|
|
423
|
+
|
|
424
|
+
auto si = py::class_<sets_index_t>(m, "SetsIndex");
|
|
425
|
+
|
|
426
|
+
si.def( //
|
|
427
|
+
py::init(&make_sets_index), //
|
|
428
|
+
py::kw_only(), //
|
|
429
|
+
py::arg("capacity") = 0, //
|
|
430
|
+
py::arg("connectivity") = config_t::connectivity_default_k, //
|
|
431
|
+
py::arg("expansion_add") = config_t::expansion_add_default_k, //
|
|
432
|
+
py::arg("expansion_search") = config_t::expansion_search_default_k //
|
|
433
|
+
);
|
|
434
|
+
|
|
435
|
+
si.def( //
|
|
436
|
+
"add",
|
|
437
|
+
[](sets_index_t& index, label_t label, py::array_t<set_member_t> set, bool copy) {
|
|
438
|
+
validate_set(set);
|
|
439
|
+
if (index.size() + 1 >= index.capacity())
|
|
440
|
+
index.reserve(ceil2(index.size() + 1));
|
|
441
|
+
auto proxy = set.unchecked<1>();
|
|
442
|
+
auto view = set_view_t{proxy.data(0), static_cast<std::size_t>(proxy.shape(0))};
|
|
443
|
+
index.add(label, view, 0, copy);
|
|
444
|
+
}, //
|
|
445
|
+
py::arg("label"), //
|
|
446
|
+
py::arg("set"), //
|
|
447
|
+
py::kw_only(), //
|
|
448
|
+
py::arg("copy") = true //
|
|
449
|
+
);
|
|
450
|
+
|
|
451
|
+
si.def( //
|
|
452
|
+
"search",
|
|
453
|
+
[](sets_index_t& index, py::array_t<set_member_t> set, std::size_t count) -> py::array_t<label_t> {
|
|
454
|
+
validate_set(set);
|
|
455
|
+
auto proxy = set.unchecked<1>();
|
|
456
|
+
auto view = set_view_t{proxy.data(0), static_cast<std::size_t>(proxy.shape(0))};
|
|
457
|
+
auto labels_py = py::array_t<label_t>(py_shape_t{static_cast<ssize_t>(count)});
|
|
458
|
+
auto labels_proxy = labels_py.mutable_unchecked<1>();
|
|
459
|
+
auto found = index.search(view, count, &labels_proxy(0), nullptr, 0);
|
|
460
|
+
labels_py.resize(py_shape_t{static_cast<ssize_t>(found)});
|
|
461
|
+
return labels_py;
|
|
462
|
+
},
|
|
463
|
+
py::arg("set"), //
|
|
464
|
+
py::arg("count") = 10 //
|
|
465
|
+
);
|
|
466
|
+
|
|
467
|
+
si.def("__len__", &sets_index_t::size);
|
|
468
|
+
si.def_property_readonly("size", &sets_index_t::size);
|
|
469
|
+
si.def_property_readonly("connectivity", &sets_index_t::connectivity);
|
|
470
|
+
si.def_property_readonly("capacity", &sets_index_t::capacity);
|
|
471
|
+
|
|
472
|
+
si.def("save", &save_index<sets_index_t>, py::arg("path"));
|
|
473
|
+
si.def("load", &load_index<sets_index_t>, py::arg("path"));
|
|
474
|
+
si.def("view", &view_index<sets_index_t>, py::arg("path"));
|
|
475
|
+
si.def("clear", &clear_index<sets_index_t>);
|
|
476
|
+
|
|
477
|
+
auto hi = py::class_<hash_index_w_meta_t>(m, "HashIndex");
|
|
478
|
+
|
|
479
|
+
hi.def( //
|
|
480
|
+
py::init(&make_hash_index), //
|
|
481
|
+
py::kw_only(), //
|
|
482
|
+
py::arg("bits"), //
|
|
483
|
+
py::arg("capacity") = 0, //
|
|
484
|
+
py::arg("connectivity") = config_t::connectivity_default_k, //
|
|
485
|
+
py::arg("expansion_add") = config_t::expansion_add_default_k, //
|
|
486
|
+
py::arg("expansion_search") = config_t::expansion_search_default_k //
|
|
487
|
+
);
|
|
488
|
+
|
|
489
|
+
hi.def( //
|
|
490
|
+
"add",
|
|
491
|
+
[](hash_index_w_meta_t& index, label_t label, py::buffer array) {
|
|
492
|
+
if (index.size() + 1 >= index.capacity())
|
|
493
|
+
index.reserve(ceil2(index.size() + 1));
|
|
494
|
+
hash_buffer(index, array);
|
|
495
|
+
index.add(label, index.buffer(), 0);
|
|
496
|
+
}, //
|
|
497
|
+
py::arg("label"), //
|
|
498
|
+
py::arg("array") //
|
|
499
|
+
);
|
|
500
|
+
|
|
501
|
+
hi.def( //
|
|
502
|
+
"search",
|
|
503
|
+
[](hash_index_w_meta_t& index, py::buffer array, std::size_t count) -> py::array_t<label_t> {
|
|
504
|
+
if (index.size() + 1 >= index.capacity())
|
|
505
|
+
index.reserve(ceil2(index.size() + 1));
|
|
506
|
+
hash_buffer(index, array);
|
|
507
|
+
auto labels_py = py::array_t<label_t>(py_shape_t{static_cast<ssize_t>(count)});
|
|
508
|
+
auto labels_proxy = labels_py.mutable_unchecked<1>();
|
|
509
|
+
auto found = index.search(index.buffer(), count, &labels_proxy(0), nullptr, 0);
|
|
510
|
+
labels_py.resize(py_shape_t{static_cast<ssize_t>(found)});
|
|
511
|
+
return labels_py;
|
|
512
|
+
},
|
|
513
|
+
py::arg("array"), //
|
|
514
|
+
py::arg("count") = 10 //
|
|
515
|
+
);
|
|
516
|
+
|
|
517
|
+
hi.def("__len__", &hash_index_w_meta_t::size);
|
|
518
|
+
hi.def_property_readonly("size", &hash_index_w_meta_t::size);
|
|
519
|
+
hi.def_property_readonly("connectivity", &hash_index_w_meta_t::connectivity);
|
|
520
|
+
hi.def_property_readonly("capacity", &hash_index_w_meta_t::capacity);
|
|
521
|
+
|
|
522
|
+
hi.def("save", &save_index<hash_index_w_meta_t>, py::arg("path"));
|
|
523
|
+
hi.def("load", &load_index<hash_index_w_meta_t>, py::arg("path"));
|
|
524
|
+
hi.def("view", &view_index<hash_index_w_meta_t>, py::arg("path"));
|
|
525
|
+
hi.def("clear", &clear_index<hash_index_w_meta_t>);
|
|
210
526
|
}
|