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/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
- print('Welcome to basic kANN benchmark!')
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 construct_both() -> tuple:
68
-
69
- index_meta = faiss.IndexHNSWFlat(d, M)
70
- index_meta.hnsw.efSearch = efSearch
71
- index_meta.hnsw.efConstruction = efConstruction
72
- dt_meta = measure(lambda: index_meta.add(xb))
73
-
74
- index_unum = usearch.Index(
75
- dim=d,
76
- expansion_add=efConstruction,
77
- expansion_search=efSearch,
78
- connectivity=M,
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
- assert len(index_unum) == 0
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
- experiments = 10
88
- durations = [construct_both() for _ in range(experiments)]
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
- return index_from_name<native_index_t>(metric.c_str(), metric.size(), dimensions, accuracy, config);
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 add_to_index(native_index_t& index, py::buffer labels, py::buffer vectors, bool copy) {
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::runtime_error("Incompatible label type!");
155
+ throw std::invalid_argument("Incompatible label type!");
61
156
 
62
157
  if (labels_info.ndim != 1)
63
- 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!");
64
159
 
65
160
  if (vectors_info.ndim != 2)
66
- throw std::runtime_error("Expects a matrix of vectors to add!");
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::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!");
73
168
 
74
169
  if (labels_count != vectors_count)
75
- throw std::runtime_error("Number of labels and vectors must match!");
170
+ throw std::invalid_argument("Number of labels and vectors must match!");
76
171
 
77
- if (index.size() + vectors_count >= index.capacity()) {
78
- std::size_t next_capacity = ceil2(index.size() + vectors_count);
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::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_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 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);
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::runtime_error("Expects a matrix of vectors to add!");
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
- if (vectors_dimensions != index.dimensions())
127
- 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!");
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::runtime_error("Incompatible scalars in the query matrix!");
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
- static void save_index(native_index_t const& index, std::string const& path) { index.save(path.c_str()); }
168
- static void load_index(native_index_t& index, std::string const& path) { index.load(path.c_str()); }
169
- 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
+ }
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("dim"), //
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", &add_to_index, //
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", &search_in_index, //
197
- py::arg("vectors"), //
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, py::arg("path"));
208
- i.def("load", &load_index, py::arg("path"));
209
- 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>);
210
526
  }