usearch 0.1.10 → 0.2.2

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/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, //
@@ -56,33 +86,91 @@ static native_index_t make_index( //
56
86
  return index_from_name<native_index_t>(metric.c_str(), metric.size(), dimensions, accuracy, config);
57
87
  }
58
88
 
59
- static void add_to_index(native_index_t& index, py::buffer labels, py::buffer vectors, bool copy) {
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));
122
+ }
123
+
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_bits_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) {
60
150
 
61
151
  py::buffer_info labels_info = labels.request();
62
152
  py::buffer_info vectors_info = vectors.request();
63
153
 
64
154
  if (labels_info.format != py::format_descriptor<label_t>::format())
65
- throw std::runtime_error("Incompatible label type!");
155
+ throw std::invalid_argument("Incompatible label type!");
66
156
 
67
157
  if (labels_info.ndim != 1)
68
- throw std::runtime_error("Labels must be placed in a single-dimensional array!");
158
+ throw std::invalid_argument("Labels must be placed in a single-dimensional array!");
69
159
 
70
160
  if (vectors_info.ndim != 2)
71
- throw std::runtime_error("Expects a matrix of vectors to add!");
161
+ throw std::invalid_argument("Expects a matrix of vectors to add!");
72
162
 
73
163
  ssize_t labels_count = labels_info.shape[0];
74
164
  ssize_t vectors_count = vectors_info.shape[0];
75
165
  ssize_t vectors_dimensions = vectors_info.shape[1];
76
- if (vectors_dimensions != index.dimensions())
77
- throw std::runtime_error("The number of vector dimensions doesn't match!");
166
+ if (vectors_dimensions != static_cast<ssize_t>(index.dimensions()))
167
+ throw std::invalid_argument("The number of vector dimensions doesn't match!");
78
168
 
79
169
  if (labels_count != vectors_count)
80
- throw std::runtime_error("Number of labels and vectors must match!");
170
+ throw std::invalid_argument("Number of labels and vectors must match!");
81
171
 
82
- if (index.size() + vectors_count >= index.capacity()) {
83
- std::size_t next_capacity = ceil2(index.size() + vectors_count);
84
- index.reserve(next_capacity);
85
- }
172
+ if (index.size() + vectors_count >= index.capacity())
173
+ index.reserve(ceil2(index.size() + vectors_count));
86
174
 
87
175
  char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
88
176
  char const* labels_data = reinterpret_cast<char const*>(labels_info.ptr);
@@ -91,8 +179,8 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
91
179
  if (vectors_info.format == "e")
92
180
  multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
93
181
  label_t label = *reinterpret_cast<label_t const*>(labels_data + task_idx * labels_info.strides[0]);
94
- f16_converted_t const* vector =
95
- reinterpret_cast<f16_converted_t const*>(vectors_data + task_idx * vectors_info.strides[0]);
182
+ f16_bits_t const* vector =
183
+ reinterpret_cast<f16_bits_t const*>(vectors_data + task_idx * vectors_info.strides[0]);
96
184
  index.add(label, vector, thread_idx, copy);
97
185
  });
98
186
  else if (vectors_info.format == "f")
@@ -108,11 +196,48 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
108
196
  index.add(label, vector, thread_idx, copy);
109
197
  });
110
198
  else
111
- throw std::runtime_error("Incompatible scalars in the vectors matrix!");
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_bits_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;
112
237
  }
113
238
 
114
239
  /**
115
- * @param vectors Matrix of vectors to search for..
240
+ * @param vectors Matrix of vectors to search for.
116
241
  * @param wanted Number of matches per request.
117
242
  *
118
243
  * @return Tuple with:
@@ -120,30 +245,34 @@ static void add_to_index(native_index_t& index, py::buffer labels, py::buffer ve
120
245
  * 2. matrix of distances,
121
246
  * 3. array with match counts.
122
247
  */
123
- static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssize_t wanted) {
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);
124
252
 
125
253
  py::buffer_info vectors_info = vectors.request();
254
+ if (vectors_info.ndim == 1)
255
+ return search_one_in_index(index, vectors, wanted);
126
256
  if (vectors_info.ndim != 2)
127
- throw std::runtime_error("Expects a matrix of vectors to add!");
257
+ throw std::invalid_argument("Expects a matrix of vectors to add!");
128
258
 
129
259
  ssize_t vectors_count = vectors_info.shape[0];
130
260
  ssize_t vectors_dimensions = vectors_info.shape[1];
131
- if (vectors_dimensions != index.dimensions())
132
- throw std::runtime_error("The number of vector dimensions doesn't match!");
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!");
133
264
 
134
- py::array_t<label_t> labels_py({vectors_count, wanted});
135
- 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)});
136
267
  py::array_t<Py_ssize_t> counts_py(vectors_count);
137
268
  auto labels_py2d = labels_py.mutable_unchecked<2>();
138
269
  auto distances_py2d = distances_py.mutable_unchecked<2>();
139
270
  auto counts_py1d = counts_py.mutable_unchecked<1>();
140
271
 
141
- char const* vectors_data = reinterpret_cast<char const*>(vectors_info.ptr);
142
-
143
272
  // https://docs.python.org/3/library/struct.html#format-characters
144
273
  if (vectors_info.format == "e")
145
274
  multithreaded(index.concurrency(), vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
146
- f16_converted_t const* vector = (f16_converted_t const*)(vectors_data + task_idx * vectors_info.strides[0]);
275
+ f16_bits_t const* vector = (f16_bits_t const*)(vectors_data + task_idx * vectors_info.strides[0]);
147
276
  counts_py1d(task_idx) = static_cast<Py_ssize_t>(
148
277
  index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
149
278
  });
@@ -160,7 +289,7 @@ static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssiz
160
289
  index.search(vector, wanted, &labels_py2d(task_idx, 0), &distances_py2d(task_idx, 0), thread_idx));
161
290
  });
162
291
  else
163
- throw std::runtime_error("Incompatible scalars in the query matrix!");
292
+ throw std::invalid_argument("Incompatible scalars in the query matrix!");
164
293
 
165
294
  py::tuple results(3);
166
295
  results[0] = labels_py;
@@ -169,18 +298,87 @@ static py::tuple search_in_index(native_index_t& index, py::buffer vectors, ssiz
169
298
  return results;
170
299
  }
171
300
 
172
- static void save_index(native_index_t const& index, std::string const& path) { index.save(path.c_str()); }
173
- static void load_index(native_index_t& index, std::string const& path) { index.load(path.c_str()); }
174
- static void view_index(native_index_t& index, std::string const& path) { index.view(path.c_str()); }
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
+ }
175
333
 
176
- PYBIND11_MODULE(usearch, m) {
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
+ }
373
+
374
+ PYBIND11_MODULE(index, m) {
177
375
  m.doc() = "Unum USearch Python bindings";
178
376
 
179
377
  auto i = py::class_<native_index_t>(m, "Index");
180
378
 
181
379
  i.def(py::init(&make_index), //
182
380
  py::kw_only(), //
183
- py::arg("ndim"), //
381
+ py::arg("ndim") = 0, //
184
382
  py::arg("capacity") = 0, //
185
383
  py::arg("dtype") = std::string("f32"), //
186
384
  py::arg("metric") = std::string("ip"), //
@@ -190,18 +388,26 @@ PYBIND11_MODULE(usearch, m) {
190
388
  py::arg("metric_pointer") = 0 //
191
389
  );
192
390
 
193
- i.def( //
194
- "add", &add_to_index, //
195
- py::arg("labels"), //
196
- py::arg("vectors"), //
197
- py::kw_only(), //
198
- 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 //
199
405
  );
200
406
 
201
- i.def( //
202
- "search", &search_in_index, //
203
- py::arg("vectors"), //
204
- py::arg("count") = 10 //
407
+ i.def( //
408
+ "search", &search_many_in_index, //
409
+ py::arg("query"), //
410
+ py::arg("count") = 10 //
205
411
  );
206
412
 
207
413
  i.def("__len__", &native_index_t::size);
@@ -210,7 +416,111 @@ PYBIND11_MODULE(usearch, m) {
210
416
  i.def_property_readonly("connectivity", &native_index_t::connectivity);
211
417
  i.def_property_readonly("capacity", &native_index_t::capacity);
212
418
 
213
- i.def("save", &save_index, py::arg("path"));
214
- i.def("load", &load_index, py::arg("path"));
215
- i.def("view", &view_index, py::arg("path"));
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>);
216
526
  }
package/python/test.py CHANGED
@@ -1,7 +1,6 @@
1
1
  import numpy as np
2
2
 
3
- import usearch
4
-
3
+ from usearch.index import Index, SetsIndex
5
4
  from numba import cfunc, types, carray
6
5
 
7
6
  # Showcases how to use Numba to JIT-compile similarity measures for USearch.
@@ -29,13 +28,13 @@ vectors = np.random.uniform(
29
28
  0, 0.3, (count_vectors, count_dimensions)).astype(np.float32)
30
29
  labels = np.array(range(count_vectors), dtype=np.longlong)
31
30
 
32
- index_udf = usearch.Index(ndim=96, metric_pointer=python_dot.address)
31
+ index_udf = Index(ndim=96, metric_pointer=python_dot.address)
33
32
  index_udf.add(labels, vectors, copy=True)
34
33
  results = index_udf.search(vectors, 10)
35
34
  print('found', results[0].shape, results[1].shape, results[2].shape)
36
35
 
37
36
 
38
- index = usearch.Index(ndim=96)
37
+ index = Index(ndim=96)
39
38
  index.add(labels, vectors, copy=True)
40
39
  assert index.ndim == count_dimensions
41
40
  assert index.capacity >= count_vectors
@@ -43,3 +42,10 @@ assert len(index) == count_vectors
43
42
 
44
43
  results = index.search(vectors, 10)
45
44
  print('found', results[0].shape, results[1].shape, results[2].shape)
45
+
46
+
47
+ index_sets = SetsIndex()
48
+ index_sets.add(10, np.array([10, 12, 15], dtype=np.uint32))
49
+ index_sets.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
50
+ results = index_sets.search(np.array([12, 15], dtype=np.uint32), 20)
51
+ print(results)
File without changes
package/rust/lib.cpp CHANGED
@@ -24,9 +24,20 @@ Matches Index::search_in_thread(rust::Slice<float const> vector, size_t count, s
24
24
  return results;
25
25
  }
26
26
 
27
- void Index::add(uint32_t label, rust::Slice<float const> vector) const { return add_in_thread(label, vector, 0); }
27
+ void Index::add(uint32_t label, rust::Slice<float const> vector) const {
28
+ index_->add(label, vector.data()); //
29
+ }
30
+
28
31
  Matches Index::search(rust::Slice<float const> vector, size_t count) const {
29
- return search_in_thread(vector, count, 0);
32
+ Matches results;
33
+ results.labels.reserve(count);
34
+ results.distances.reserve(count);
35
+ for (size_t i = 0; i != count; ++i)
36
+ results.labels.push_back(0), results.distances.push_back(0);
37
+ results.count = index_->search(vector.data(), count, results.labels.data(), results.distances.data());
38
+ results.labels.truncate(results.count);
39
+ results.distances.truncate(results.count);
40
+ return results;
30
41
  }
31
42
 
32
43
  void Index::reserve(size_t capacity) const { index_->reserve(capacity); }
package/setup.py CHANGED
@@ -9,11 +9,15 @@ compile_args = ['-std=c++11', '-O3']
9
9
  if sys.platform == 'darwin':
10
10
  compile_args.append('-mmacosx-version-min=10.13')
11
11
 
12
+ if sys.platform == 'win32':
13
+ compile_args = ['/std:c++14', '/O2']
14
+
15
+
12
16
  ext_modules = [
13
17
  Pybind11Extension(
14
- 'usearch',
18
+ 'usearch.index',
15
19
  ['python/lib.cpp'],
16
- extra_compile_args=compile_args,
20
+ extra_compile_args=compile_args
17
21
  ),
18
22
  ]
19
23
 
@@ -30,7 +34,8 @@ setup(
30
34
 
31
35
  name=__lib_name__,
32
36
  version=__version__,
33
- packages=find_packages(),
37
+ packages=['usearch'],
38
+ package_dir={'usearch': 'python/usearch'},
34
39
 
35
40
  description='Smaller & Faster Single-File Vector Search Engine from Unum',
36
41
  long_description=long_description,
@@ -255,6 +255,21 @@ inline static simsimd_f32_t simsimd_dot_f32x4avx2(simsimd_f32_t const* a, simsim
255
255
  #endif
256
256
  }
257
257
 
258
+ inline static simsimd_f32_t simsimd_dot_f16x16avx512(simsimd_f16_t const* a, simsimd_f16_t const* b, size_t d) {
259
+ #if defined(__AVX512F__)
260
+ __m512 ab_vec = _mm512_set1_ps(0);
261
+ for (size_t i = 0; i != d; i += 16)
262
+ ab_vec = _mm512_fmadd_ps( //
263
+ _mm512_cvtxph_ps(_mm256_loadu_ph(a + i)), //
264
+ _mm512_cvtxph_ps(_mm256_loadu_ph(b + i)), //
265
+ ab_vec);
266
+ return _mm512_reduce_add_ps(ab_vec);
267
+ #else
268
+ (void)a, (void)b, (void)d;
269
+ return 0;
270
+ #endif
271
+ }
272
+
258
273
  inline static int32_t simsimd_dot_i8x16avx2(int8_t const* a, int8_t const* b, size_t d) {
259
274
  #if defined(__AVX2__)
260
275
  __m256i ab_vec = _mm256_set1_epi32(0);