usearch 0.22.3 → 1.0.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/.github/workflows/prerelease.yml +3 -3
- package/.releaserc +10 -2
- package/.vscode/launch.json +10 -10
- package/.vscode/settings.json +9 -1
- package/.vscode/tasks.json +4 -4
- package/CMakeLists.txt +6 -1
- package/Cargo.toml +2 -2
- package/README.md +35 -41
- package/build.rs +2 -2
- package/c/CMakeLists.txt +48 -0
- package/c/README.md +34 -0
- package/c/lib.cpp +83 -72
- package/c/test.c +20 -20
- package/c/usearch.h +132 -32
- package/conanfile.py +1 -1
- package/cpp/CMakeLists.txt +3 -1
- package/cpp/README.md +16 -16
- package/cpp/bench.cpp +50 -126
- package/cpp/test.cpp +164 -80
- package/docs/benchmarks.md +26 -5
- package/docs/c/reference.rst +3 -3
- package/docs/compilation.md +31 -49
- package/docs/cpp/index.rst +1 -1
- package/docs/cpp/reference.rst +4 -4
- package/golang/lib.go +275 -141
- package/golang/lib_test.go +105 -0
- package/include/usearch/index.hpp +1681 -1656
- package/include/usearch/index_dense.hpp +1400 -0
- package/include/usearch/index_plugins.hpp +1263 -0
- package/java/README.md +1 -1
- package/java/cloud/unum/usearch/Index.java +10 -10
- package/java/cloud/unum/usearch/cloud_unum_usearch_Index.cpp +34 -37
- package/java/test/IndexTest.java +1 -1
- package/javascript/README.md +6 -6
- package/javascript/lib.cpp +94 -55
- package/javascript/test.js +25 -11
- package/javascript/usearch.d.ts +87 -0
- package/javascript/usearch.js +2 -118
- package/objc/USearchObjective.mm +57 -31
- package/objc/include/USearchObjective.h +32 -22
- package/package.json +5 -3
- package/python/README.md +23 -14
- package/python/lib.cpp +468 -444
- package/python/scripts/bench.ipynb +3 -2
- package/python/scripts/bench.py +4 -4
- package/python/scripts/index_faiss.py +5 -5
- package/python/scripts/join.py +4 -4
- package/python/scripts/test.py +152 -61
- package/python/scripts/test_jit.py +10 -12
- package/python/scripts/test_sparse.py +32 -0
- package/python/usearch/client.py +26 -28
- package/python/usearch/eval.py +73 -23
- package/python/usearch/index.py +520 -228
- package/python/usearch/io.py +57 -25
- package/python/usearch/numba.py +4 -4
- package/python/usearch/server.py +10 -10
- package/rust/lib.cpp +22 -27
- package/rust/lib.hpp +11 -12
- package/rust/lib.rs +3 -3
- package/setup.py +5 -2
- package/simsimd/.github/workflows/prerelease.yml +73 -0
- package/simsimd/.github/workflows/release.yml +110 -0
- package/simsimd/.github/workflows/update_version.sh +3 -0
- package/simsimd/.releaserc +77 -0
- package/simsimd/.vscode/launch.json +27 -1
- package/simsimd/.vscode/settings.json +22 -1
- package/simsimd/.vscode/tasks.json +6 -0
- package/simsimd/CMakeLists.txt +3 -2
- package/simsimd/README.md +82 -46
- package/simsimd/bench.cxx +63 -24
- package/simsimd/include/simsimd/simsimd.h +88 -66
- package/simsimd/include/simsimd/simsimd_chem.h +126 -0
- package/simsimd/package.json +8 -0
- package/simsimd/pyproject.toml +42 -0
- package/simsimd/python/python.c +81 -0
- package/simsimd/python/run.sh +1 -0
- package/simsimd/python/test.c +109 -0
- package/simsimd/python/test.py +23 -0
- package/simsimd/setup.py +44 -0
- package/swift/Index+Sugar.swift +75 -23
- package/swift/README.md +26 -2
- package/swift/Test.swift +19 -7
- package/swift/USearch.swift +2 -2
- package/tsconfig.json +8 -0
- package/wasm/CMakeLists.txt +1 -0
- package/wasmer.toml +1 -1
- package/wolfram/CMakeLists.txt +1 -0
- package/wolfram/lib.cpp +11 -13
- package/golang/usearch.h +0 -9
- package/golang/usearch_par_test.go +0 -40
- package/golang/usearch_test.go +0 -82
- package/include/usearch/index_punned_dense.hpp +0 -1001
- package/include/usearch/index_punned_helpers.hpp +0 -676
package/python/lib.cpp
CHANGED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
/**
|
|
2
2
|
* @file python.cpp
|
|
3
|
-
* @author
|
|
3
|
+
* @author Ash Vardanian
|
|
4
4
|
* @brief Python bindings for Unum USearch.
|
|
5
5
|
* @date 2023-04-26
|
|
6
6
|
*
|
|
@@ -12,6 +12,7 @@
|
|
|
12
12
|
*
|
|
13
13
|
* @copyright Copyright (c) 2023
|
|
14
14
|
*/
|
|
15
|
+
#define __cpp_exceptions 1
|
|
15
16
|
#include <limits> // `std::numeric_limits`
|
|
16
17
|
#include <thread> // `std::thread`
|
|
17
18
|
|
|
@@ -21,7 +22,8 @@
|
|
|
21
22
|
#include <pybind11/pybind11.h>
|
|
22
23
|
#include <pybind11/stl.h>
|
|
23
24
|
|
|
24
|
-
#include <usearch/
|
|
25
|
+
#include <usearch/index_dense.hpp>
|
|
26
|
+
#include <usearch/index_plugins.hpp>
|
|
25
27
|
|
|
26
28
|
using namespace unum::usearch;
|
|
27
29
|
using namespace unum;
|
|
@@ -39,83 +41,104 @@ enum class metric_signature_t {
|
|
|
39
41
|
|
|
40
42
|
namespace py = pybind11;
|
|
41
43
|
using py_shape_t = py::array::ShapeContainer;
|
|
42
|
-
|
|
43
|
-
using
|
|
44
|
-
using
|
|
45
|
-
|
|
46
|
-
using
|
|
47
|
-
using
|
|
48
|
-
using
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
using native_t = dense_index_t;
|
|
44
|
+
#define key_t typename index_dense_t::key_t
|
|
45
|
+
using metric_t = metric_punned_t;
|
|
46
|
+
using distance_t = distance_punned_t;
|
|
47
|
+
|
|
48
|
+
using dense_add_result_t = typename index_dense_t::add_result_t;
|
|
49
|
+
using dense_search_result_t = typename index_dense_t::search_result_t;
|
|
50
|
+
using dense_labeling_result_t = typename index_dense_t::labeling_result_t;
|
|
51
|
+
|
|
52
|
+
struct dense_index_py_t : public index_dense_t {
|
|
53
|
+
using native_t = index_dense_t;
|
|
53
54
|
using native_t::add;
|
|
54
55
|
using native_t::capacity;
|
|
55
56
|
using native_t::reserve;
|
|
56
57
|
using native_t::search;
|
|
57
58
|
using native_t::size;
|
|
58
59
|
|
|
59
|
-
dense_index_py_t(native_t&& base) :
|
|
60
|
+
dense_index_py_t(native_t&& base) : index_dense_t(std::move(base)) {}
|
|
60
61
|
};
|
|
61
62
|
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
63
|
+
struct dense_indexes_py_t {
|
|
64
|
+
std::vector<std::shared_ptr<dense_index_py_t>> shards_;
|
|
65
|
+
|
|
66
|
+
void add(std::shared_ptr<dense_index_py_t> shard) { shards_.push_back(shard); }
|
|
67
|
+
std::size_t bytes_per_vector() const noexcept { return shards_.empty() ? 0 : shards_[0]->bytes_per_vector(); }
|
|
68
|
+
std::size_t scalar_words() const noexcept { return shards_.empty() ? 0 : shards_[0]->scalar_words(); }
|
|
69
|
+
index_limits_t limits() const noexcept { return {size(), std::numeric_limits<std::size_t>::max()}; }
|
|
70
|
+
|
|
71
|
+
void add_paths(std::vector<std::string> const& paths, bool view = true, std::size_t threads = 0) {
|
|
72
|
+
if (!threads)
|
|
73
|
+
threads = std::thread::hardware_concurrency();
|
|
74
|
+
|
|
75
|
+
shards_.reserve(shards_.size() + paths.size());
|
|
76
|
+
std::mutex shards_mutex;
|
|
77
|
+
executor_default_t{threads}.execute_bulk(paths.size(), [&](std::size_t, std::size_t task_idx) {
|
|
78
|
+
index_dense_t index = index_dense_t::make(paths[task_idx].c_str(), view);
|
|
79
|
+
if (!index)
|
|
80
|
+
return;
|
|
81
|
+
auto shared_index = std::make_shared<dense_index_py_t>(std::move(index));
|
|
82
|
+
std::unique_lock<std::mutex> lock(shards_mutex);
|
|
83
|
+
shards_.push_back(shared_index);
|
|
84
|
+
if (PyErr_CheckSignals() != 0)
|
|
85
|
+
throw py::error_already_set();
|
|
86
|
+
});
|
|
87
|
+
}
|
|
73
88
|
|
|
74
|
-
|
|
89
|
+
std::size_t size() const noexcept {
|
|
90
|
+
std::size_t result = 0;
|
|
91
|
+
for (auto const& shard : shards_)
|
|
92
|
+
result += shard->size();
|
|
93
|
+
return result;
|
|
94
|
+
}
|
|
75
95
|
};
|
|
76
96
|
|
|
77
97
|
template <typename scalar_at>
|
|
78
|
-
metric_t typed_udf(
|
|
79
|
-
metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr,
|
|
98
|
+
metric_t typed_udf( //
|
|
99
|
+
metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, //
|
|
100
|
+
scalar_kind_t scalar_kind, std::size_t dimensions) {
|
|
80
101
|
//
|
|
81
|
-
metric_t
|
|
82
|
-
|
|
83
|
-
|
|
102
|
+
using stl_function_t = metric_t::stl_function_t;
|
|
103
|
+
stl_function_t stl_function;
|
|
104
|
+
std::size_t scalar_words =
|
|
105
|
+
divide_round_up(dimensions * bits_per_scalar(scalar_kind), bits_per_scalar_word(scalar_kind));
|
|
84
106
|
switch (signature) {
|
|
85
107
|
case metric_signature_t::array_array_k:
|
|
86
|
-
|
|
87
|
-
using metric_raw_t =
|
|
108
|
+
stl_function = [metric_uintptr](byte_t const* a, byte_t const* b) -> distance_t {
|
|
109
|
+
using metric_raw_t = distance_punned_t (*)(scalar_at const*, scalar_at const*);
|
|
88
110
|
metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
|
|
89
|
-
return metric_ptr((scalar_at const*)a
|
|
111
|
+
return metric_ptr((scalar_at const*)a, (scalar_at const*)b);
|
|
90
112
|
};
|
|
91
113
|
break;
|
|
92
114
|
case metric_signature_t::array_array_size_k:
|
|
93
|
-
|
|
94
|
-
using metric_raw_t =
|
|
115
|
+
stl_function = [metric_uintptr, scalar_words](byte_t const* a, byte_t const* b) -> distance_t {
|
|
116
|
+
using metric_raw_t = distance_punned_t (*)(scalar_at const*, scalar_at const*, size_t);
|
|
95
117
|
metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
|
|
96
|
-
return metric_ptr((scalar_at const*)a
|
|
118
|
+
return metric_ptr((scalar_at const*)a, (scalar_at const*)b, scalar_words);
|
|
97
119
|
};
|
|
98
120
|
break;
|
|
99
121
|
case metric_signature_t::array_size_array_size_k:
|
|
100
|
-
|
|
101
|
-
using metric_raw_t =
|
|
122
|
+
stl_function = [metric_uintptr, scalar_words](byte_t const* a, byte_t const* b) -> distance_t {
|
|
123
|
+
using metric_raw_t = distance_punned_t (*)(scalar_at const*, size_t, scalar_at const*, size_t);
|
|
102
124
|
metric_raw_t metric_ptr = reinterpret_cast<metric_raw_t>(metric_uintptr);
|
|
103
|
-
return metric_ptr(
|
|
104
|
-
(scalar_at const*)a.data(), a.size() / sizeof(scalar_at), //
|
|
105
|
-
(scalar_at const*)b.data(), b.size() / sizeof(scalar_at));
|
|
125
|
+
return metric_ptr((scalar_at const*)a, scalar_words, (scalar_at const*)b, scalar_words);
|
|
106
126
|
};
|
|
107
127
|
break;
|
|
108
128
|
}
|
|
109
|
-
return
|
|
129
|
+
return metric_t(stl_function, dimensions, kind, scalar_kind);
|
|
110
130
|
}
|
|
111
131
|
|
|
112
|
-
metric_t udf(
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
case scalar_kind_t::
|
|
118
|
-
case scalar_kind_t::
|
|
132
|
+
metric_t udf( //
|
|
133
|
+
metric_kind_t kind, metric_signature_t signature, std::uintptr_t metric_uintptr, //
|
|
134
|
+
scalar_kind_t scalar_kind, std::size_t dimensions) {
|
|
135
|
+
|
|
136
|
+
switch (scalar_kind) {
|
|
137
|
+
case scalar_kind_t::b1x8_k: return typed_udf<b1x8_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
138
|
+
case scalar_kind_t::i8_k: return typed_udf<i8_bits_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
139
|
+
case scalar_kind_t::f16_k: return typed_udf<f16_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
140
|
+
case scalar_kind_t::f32_k: return typed_udf<f32_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
141
|
+
case scalar_kind_t::f64_k: return typed_udf<f64_t>(kind, signature, metric_uintptr, scalar_kind, dimensions);
|
|
119
142
|
default: return {};
|
|
120
143
|
}
|
|
121
144
|
}
|
|
@@ -123,37 +146,19 @@ metric_t udf(metric_kind_t kind, metric_signature_t signature, std::uintptr_t me
|
|
|
123
146
|
static dense_index_py_t make_index( //
|
|
124
147
|
std::size_t dimensions, //
|
|
125
148
|
scalar_kind_t scalar_kind, //
|
|
126
|
-
metric_kind_t metric_kind, //
|
|
127
149
|
std::size_t connectivity, //
|
|
128
150
|
std::size_t expansion_add, //
|
|
129
151
|
std::size_t expansion_search, //
|
|
152
|
+
metric_kind_t metric_kind, //
|
|
130
153
|
metric_signature_t metric_signature, //
|
|
131
|
-
std::uintptr_t metric_uintptr
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
if (metric_uintptr)
|
|
141
|
-
return dense_index_t::make( //
|
|
142
|
-
dimensions, udf(metric_kind, metric_signature, metric_uintptr, scalar_kind), config, scalar_kind,
|
|
143
|
-
expansion_add, expansion_search);
|
|
144
|
-
else
|
|
145
|
-
return dense_index_t::make(dimensions, metric_kind, config, scalar_kind, expansion_add, expansion_search);
|
|
146
|
-
}
|
|
147
|
-
|
|
148
|
-
static std::unique_ptr<sparse_index_py_t> make_sparse_index( //
|
|
149
|
-
std::size_t connectivity, //
|
|
150
|
-
std::size_t expansion_add, //
|
|
151
|
-
std::size_t expansion_search //
|
|
152
|
-
) {
|
|
153
|
-
index_config_t config;
|
|
154
|
-
config.connectivity = connectivity;
|
|
155
|
-
|
|
156
|
-
return std::unique_ptr<sparse_index_py_t>(new sparse_index_py_t(sparse_index_t(config)));
|
|
154
|
+
std::uintptr_t metric_uintptr) {
|
|
155
|
+
|
|
156
|
+
index_dense_config_t config(connectivity, expansion_add, expansion_search);
|
|
157
|
+
metric_t metric = //
|
|
158
|
+
metric_uintptr //
|
|
159
|
+
? udf(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
|
|
160
|
+
: metric_t(dimensions, metric_kind, scalar_kind);
|
|
161
|
+
return index_dense_t::make(metric, config);
|
|
157
162
|
}
|
|
158
163
|
|
|
159
164
|
scalar_kind_t numpy_string_to_kind(std::string const& name) {
|
|
@@ -161,162 +166,96 @@ scalar_kind_t numpy_string_to_kind(std::string const& name) {
|
|
|
161
166
|
if (name == "B" || name == "<B" || name == "u1" || name == "|u1")
|
|
162
167
|
return scalar_kind_t::b1x8_k;
|
|
163
168
|
else if (name == "b" || name == "<b" || name == "i1" || name == "|i1")
|
|
164
|
-
return scalar_kind_t::
|
|
169
|
+
return scalar_kind_t::i8_k;
|
|
165
170
|
else if (name == "e" || name == "<e" || name == "f2" || name == "<f2")
|
|
166
171
|
return scalar_kind_t::f16_k;
|
|
167
172
|
else if (name == "f" || name == "<f" || name == "f4" || name == "<f4")
|
|
168
173
|
return scalar_kind_t::f32_k;
|
|
169
|
-
else if (name == "d" || name == "<d" || name == "
|
|
174
|
+
else if (name == "d" || name == "<d" || name == "i8" || name == "<i8")
|
|
170
175
|
return scalar_kind_t::f64_k;
|
|
171
176
|
else
|
|
172
177
|
return scalar_kind_t::unknown_k;
|
|
173
178
|
}
|
|
174
179
|
|
|
175
|
-
static void add_one_to_index(dense_index_py_t& index, label_t label, py::buffer vector, bool copy, std::size_t) {
|
|
176
|
-
|
|
177
|
-
py::buffer_info vector_info = vector.request();
|
|
178
|
-
if (vector_info.ndim != 1)
|
|
179
|
-
throw std::invalid_argument("Expects a vector, not a higher-rank tensor!");
|
|
180
|
-
|
|
181
|
-
Py_ssize_t vector_dimensions = vector_info.shape[0];
|
|
182
|
-
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
183
|
-
if (vector_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
|
|
184
|
-
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
185
|
-
|
|
186
|
-
if (!index.reserve(ceil2(index.size() + 1)))
|
|
187
|
-
throw std::invalid_argument("Out of memory!");
|
|
188
|
-
|
|
189
|
-
add_config_t config;
|
|
190
|
-
config.store_vector = copy;
|
|
191
|
-
|
|
192
|
-
switch (numpy_string_to_kind(vector_info.format)) {
|
|
193
|
-
case scalar_kind_t::b1x8_k: index.add(label, (b1x8_t const*)(vector_data), config).error.raise(); break;
|
|
194
|
-
case scalar_kind_t::f8_k: index.add(label, (f8_bits_t const*)(vector_data), config).error.raise(); break;
|
|
195
|
-
case scalar_kind_t::f16_k: index.add(label, (f16_t const*)(vector_data), config).error.raise(); break;
|
|
196
|
-
case scalar_kind_t::f32_k: index.add(label, (f32_t const*)(vector_data), config).error.raise(); break;
|
|
197
|
-
case scalar_kind_t::f64_k: index.add(label, (f64_t const*)(vector_data), config).error.raise(); break;
|
|
198
|
-
case scalar_kind_t::unknown_k:
|
|
199
|
-
throw std::invalid_argument("Incompatible scalars in the vector: " + vector_info.format);
|
|
200
|
-
}
|
|
201
|
-
}
|
|
202
|
-
|
|
203
180
|
template <typename scalar_at>
|
|
204
|
-
static void add_typed_to_index(
|
|
205
|
-
dense_index_py_t& index,
|
|
206
|
-
py::buffer_info const&
|
|
207
|
-
bool
|
|
181
|
+
static void add_typed_to_index( //
|
|
182
|
+
dense_index_py_t& index, //
|
|
183
|
+
py::buffer_info const& keys_info, py::buffer_info const& vectors_info, //
|
|
184
|
+
bool force_copy, std::size_t threads) {
|
|
208
185
|
|
|
209
186
|
Py_ssize_t vectors_count = vectors_info.shape[0];
|
|
210
|
-
|
|
211
|
-
|
|
187
|
+
byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);
|
|
188
|
+
byte_t const* keys_data = reinterpret_cast<byte_t const*>(keys_info.ptr);
|
|
212
189
|
|
|
213
190
|
executor_default_t{threads}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
214
|
-
|
|
215
|
-
config.
|
|
191
|
+
index_dense_update_config_t config;
|
|
192
|
+
config.force_vector_copy = force_copy;
|
|
216
193
|
config.thread = thread_idx;
|
|
217
|
-
|
|
194
|
+
key_t key = *reinterpret_cast<key_t const*>(keys_data + task_idx * keys_info.strides[0]);
|
|
218
195
|
scalar_at const* vector = reinterpret_cast<scalar_at const*>(vectors_data + task_idx * vectors_info.strides[0]);
|
|
219
|
-
index.add(
|
|
196
|
+
dense_add_result_t result = index.add(key, vector, config);
|
|
197
|
+
result.error.raise();
|
|
220
198
|
if (PyErr_CheckSignals() != 0)
|
|
221
199
|
throw py::error_already_set();
|
|
222
200
|
});
|
|
223
201
|
}
|
|
224
202
|
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
203
|
+
template <typename index_at>
|
|
204
|
+
static void add_many_to_index( //
|
|
205
|
+
index_at& index, py::buffer keys, py::buffer vectors, //
|
|
206
|
+
bool force_copy, std::size_t threads) {
|
|
228
207
|
|
|
229
|
-
py::buffer_info
|
|
208
|
+
py::buffer_info keys_info = keys.request();
|
|
230
209
|
py::buffer_info vectors_info = vectors.request();
|
|
231
210
|
|
|
232
|
-
if (
|
|
233
|
-
throw std::invalid_argument("Incompatible
|
|
211
|
+
if (keys_info.itemsize != sizeof(key_t))
|
|
212
|
+
throw std::invalid_argument("Incompatible key type!");
|
|
234
213
|
|
|
235
|
-
if (
|
|
214
|
+
if (keys_info.ndim != 1)
|
|
236
215
|
throw std::invalid_argument("Labels must be placed in a single-dimensional array!");
|
|
237
216
|
|
|
238
217
|
if (vectors_info.ndim != 2)
|
|
239
218
|
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
240
219
|
|
|
241
|
-
Py_ssize_t
|
|
220
|
+
Py_ssize_t keys_count = keys_info.shape[0];
|
|
242
221
|
Py_ssize_t vectors_count = vectors_info.shape[0];
|
|
243
222
|
Py_ssize_t vectors_dimensions = vectors_info.shape[1];
|
|
244
223
|
if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
|
|
245
224
|
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
246
225
|
|
|
247
|
-
if (
|
|
248
|
-
throw std::invalid_argument("Number of
|
|
226
|
+
if (keys_count != vectors_count)
|
|
227
|
+
throw std::invalid_argument("Number of keys and vectors must match!");
|
|
249
228
|
|
|
250
229
|
if (!threads)
|
|
251
230
|
threads = std::thread::hardware_concurrency();
|
|
252
231
|
if (!index.reserve(index_limits_t(ceil2(index.size() + vectors_count), threads)))
|
|
253
232
|
throw std::invalid_argument("Out of memory!");
|
|
254
233
|
|
|
234
|
+
// clang-format off
|
|
255
235
|
switch (numpy_string_to_kind(vectors_info.format)) {
|
|
256
|
-
case scalar_kind_t::b1x8_k: add_typed_to_index<b1x8_t>(index,
|
|
257
|
-
case scalar_kind_t::
|
|
258
|
-
case scalar_kind_t::f16_k: add_typed_to_index<f16_t>(index,
|
|
259
|
-
case scalar_kind_t::f32_k: add_typed_to_index<f32_t>(index,
|
|
260
|
-
case scalar_kind_t::f64_k: add_typed_to_index<f64_t>(index,
|
|
261
|
-
|
|
262
|
-
throw std::invalid_argument("Incompatible scalars in the vectors matrix: " + vectors_info.format);
|
|
236
|
+
case scalar_kind_t::b1x8_k: add_typed_to_index<b1x8_t>(index, keys_info, vectors_info, force_copy, threads); break;
|
|
237
|
+
case scalar_kind_t::i8_k: add_typed_to_index<i8_bits_t>(index, keys_info, vectors_info, force_copy, threads); break;
|
|
238
|
+
case scalar_kind_t::f16_k: add_typed_to_index<f16_t>(index, keys_info, vectors_info, force_copy, threads); break;
|
|
239
|
+
case scalar_kind_t::f32_k: add_typed_to_index<f32_t>(index, keys_info, vectors_info, force_copy, threads); break;
|
|
240
|
+
case scalar_kind_t::f64_k: add_typed_to_index<f64_t>(index, keys_info, vectors_info, force_copy, threads); break;
|
|
241
|
+
default: throw std::invalid_argument("Incompatible scalars in the vectors matrix: " + vectors_info.format);
|
|
263
242
|
}
|
|
264
|
-
|
|
265
|
-
|
|
266
|
-
static py::tuple search_one_in_index(dense_index_py_t& index, py::buffer vector, std::size_t wanted, bool exact) {
|
|
267
|
-
|
|
268
|
-
py::buffer_info vector_info = vector.request();
|
|
269
|
-
Py_ssize_t vector_dimensions = vector_info.shape[0];
|
|
270
|
-
char const* vector_data = reinterpret_cast<char const*>(vector_info.ptr);
|
|
271
|
-
if (vector_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
|
|
272
|
-
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
273
|
-
|
|
274
|
-
py::array_t<label_t> labels_py(static_cast<Py_ssize_t>(wanted));
|
|
275
|
-
py::array_t<distance_t> distances_py(static_cast<Py_ssize_t>(wanted));
|
|
276
|
-
std::size_t count{};
|
|
277
|
-
auto labels_py1d = labels_py.template mutable_unchecked<1>();
|
|
278
|
-
auto distances_py1d = distances_py.template mutable_unchecked<1>();
|
|
279
|
-
|
|
280
|
-
search_config_t config;
|
|
281
|
-
config.exact = exact;
|
|
282
|
-
|
|
283
|
-
auto raise_and_dump = [&](dense_search_result_t result) {
|
|
284
|
-
result.error.raise();
|
|
285
|
-
count = result.dump_to(&labels_py1d(0), &distances_py1d(0));
|
|
286
|
-
};
|
|
287
|
-
|
|
288
|
-
switch (numpy_string_to_kind(vector_info.format)) {
|
|
289
|
-
case scalar_kind_t::b1x8_k: raise_and_dump(index.search((b1x8_t const*)(vector_data), wanted, config)); break;
|
|
290
|
-
case scalar_kind_t::f8_k: raise_and_dump(index.search((f8_bits_t const*)(vector_data), wanted, config)); break;
|
|
291
|
-
case scalar_kind_t::f16_k: raise_and_dump(index.search((f16_t const*)(vector_data), wanted, config)); break;
|
|
292
|
-
case scalar_kind_t::f32_k: raise_and_dump(index.search((f32_t const*)(vector_data), wanted, config)); break;
|
|
293
|
-
case scalar_kind_t::f64_k: raise_and_dump(index.search((f64_t const*)(vector_data), wanted, config)); break;
|
|
294
|
-
case scalar_kind_t::unknown_k:
|
|
295
|
-
throw std::invalid_argument("Incompatible scalars in the query vector: " + vector_info.format);
|
|
296
|
-
}
|
|
297
|
-
|
|
298
|
-
labels_py.resize(py_shape_t{static_cast<Py_ssize_t>(count)});
|
|
299
|
-
distances_py.resize(py_shape_t{static_cast<Py_ssize_t>(count)});
|
|
300
|
-
|
|
301
|
-
py::tuple results(3);
|
|
302
|
-
results[0] = labels_py;
|
|
303
|
-
results[1] = distances_py;
|
|
304
|
-
results[2] = static_cast<Py_ssize_t>(count);
|
|
305
|
-
return results;
|
|
243
|
+
// clang-format on
|
|
306
244
|
}
|
|
307
245
|
|
|
308
246
|
template <typename scalar_at>
|
|
309
247
|
static void search_typed( //
|
|
310
248
|
dense_index_py_t& index, py::buffer_info& vectors_info, //
|
|
311
249
|
std::size_t wanted, bool exact, std::size_t threads, //
|
|
312
|
-
py::array_t<
|
|
250
|
+
py::array_t<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
|
|
251
|
+
std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances) {
|
|
313
252
|
|
|
314
|
-
auto
|
|
253
|
+
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
315
254
|
auto distances_py2d = distances_py.template mutable_unchecked<2>();
|
|
316
255
|
auto counts_py1d = counts_py.template mutable_unchecked<1>();
|
|
317
256
|
|
|
318
257
|
Py_ssize_t vectors_count = vectors_info.shape[0];
|
|
319
|
-
|
|
258
|
+
byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);
|
|
320
259
|
|
|
321
260
|
if (!threads)
|
|
322
261
|
threads = std::thread::hardware_concurrency();
|
|
@@ -324,19 +263,77 @@ static void search_typed( //
|
|
|
324
263
|
throw std::invalid_argument("Out of memory!");
|
|
325
264
|
|
|
326
265
|
executor_default_t{threads}.execute_bulk(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {
|
|
327
|
-
|
|
266
|
+
index_search_config_t config;
|
|
328
267
|
config.thread = thread_idx;
|
|
329
268
|
config.exact = exact;
|
|
330
269
|
scalar_at const* vector = (scalar_at const*)(vectors_data + task_idx * vectors_info.strides[0]);
|
|
331
270
|
dense_search_result_t result = index.search(vector, wanted, config);
|
|
332
271
|
result.error.raise();
|
|
333
272
|
counts_py1d(task_idx) =
|
|
334
|
-
static_cast<Py_ssize_t>(result.dump_to(&
|
|
273
|
+
static_cast<Py_ssize_t>(result.dump_to(&keys_py2d(task_idx, 0), &distances_py2d(task_idx, 0)));
|
|
274
|
+
|
|
275
|
+
stats_visited_members += result.visited_members;
|
|
276
|
+
stats_computed_distances += result.computed_distances;
|
|
335
277
|
if (PyErr_CheckSignals() != 0)
|
|
336
278
|
throw py::error_already_set();
|
|
337
279
|
});
|
|
338
280
|
}
|
|
339
281
|
|
|
282
|
+
template <typename scalar_at>
|
|
283
|
+
static void search_typed( //
|
|
284
|
+
dense_indexes_py_t& indexes, py::buffer_info& vectors_info, //
|
|
285
|
+
std::size_t wanted, bool exact, std::size_t threads, //
|
|
286
|
+
py::array_t<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,
|
|
287
|
+
std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances) {
|
|
288
|
+
|
|
289
|
+
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
290
|
+
auto distances_py2d = distances_py.template mutable_unchecked<2>();
|
|
291
|
+
auto counts_py1d = counts_py.template mutable_unchecked<1>();
|
|
292
|
+
|
|
293
|
+
Py_ssize_t vectors_count = vectors_info.shape[0];
|
|
294
|
+
byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);
|
|
295
|
+
for (std::size_t vector_idx = 0; vector_idx != static_cast<std::size_t>(vectors_count); ++vector_idx)
|
|
296
|
+
counts_py1d(vector_idx) = 0;
|
|
297
|
+
|
|
298
|
+
if (!threads)
|
|
299
|
+
threads = std::thread::hardware_concurrency();
|
|
300
|
+
|
|
301
|
+
std::vector<std::mutex> query_mutexes(static_cast<std::size_t>(vectors_count));
|
|
302
|
+
executor_default_t{threads}.execute_bulk(indexes.shards_.size(), [&](std::size_t, std::size_t task_idx) {
|
|
303
|
+
dense_index_py_t& index = *indexes.shards_[task_idx].get();
|
|
304
|
+
|
|
305
|
+
index_limits_t limits;
|
|
306
|
+
limits.members = index.size();
|
|
307
|
+
limits.threads_add = 0;
|
|
308
|
+
limits.threads_search = 1;
|
|
309
|
+
if (!index.reserve(limits))
|
|
310
|
+
throw std::bad_alloc();
|
|
311
|
+
|
|
312
|
+
index_search_config_t config;
|
|
313
|
+
config.thread = 0;
|
|
314
|
+
config.exact = exact;
|
|
315
|
+
|
|
316
|
+
for (std::size_t vector_idx = 0; vector_idx != static_cast<std::size_t>(vectors_count); ++vector_idx) {
|
|
317
|
+
scalar_at const* vector = (scalar_at const*)(vectors_data + vector_idx * vectors_info.strides[0]);
|
|
318
|
+
dense_search_result_t result = index.search(vector, wanted, config);
|
|
319
|
+
result.error.raise();
|
|
320
|
+
{
|
|
321
|
+
std::unique_lock<std::mutex> lock(query_mutexes[vector_idx]);
|
|
322
|
+
counts_py1d(vector_idx) = static_cast<Py_ssize_t>(result.merge_into( //
|
|
323
|
+
&keys_py2d(vector_idx, 0), //
|
|
324
|
+
&distances_py2d(vector_idx, 0), //
|
|
325
|
+
static_cast<std::size_t>(counts_py1d(vector_idx)), //
|
|
326
|
+
wanted));
|
|
327
|
+
}
|
|
328
|
+
|
|
329
|
+
stats_visited_members += result.visited_members;
|
|
330
|
+
stats_computed_distances += result.computed_distances;
|
|
331
|
+
if (PyErr_CheckSignals() != 0)
|
|
332
|
+
throw py::error_already_set();
|
|
333
|
+
}
|
|
334
|
+
});
|
|
335
|
+
}
|
|
336
|
+
|
|
340
337
|
/**
|
|
341
338
|
* @param vectors Matrix of vectors to search for.
|
|
342
339
|
* @param wanted Number of matches per request.
|
|
@@ -346,18 +343,17 @@ static void search_typed( //
|
|
|
346
343
|
* 2. matrix of distances,
|
|
347
344
|
* 3. array with match counts.
|
|
348
345
|
*/
|
|
346
|
+
template <typename index_at>
|
|
349
347
|
static py::tuple search_many_in_index( //
|
|
350
|
-
|
|
348
|
+
index_at& index, py::buffer vectors, std::size_t wanted, bool exact, std::size_t threads) {
|
|
351
349
|
|
|
352
350
|
if (wanted == 0)
|
|
353
|
-
return py::tuple(
|
|
351
|
+
return py::tuple(5);
|
|
354
352
|
|
|
355
353
|
if (index.limits().threads_search < threads)
|
|
356
354
|
throw std::invalid_argument("Can't use that many threads!");
|
|
357
355
|
|
|
358
356
|
py::buffer_info vectors_info = vectors.request();
|
|
359
|
-
if (vectors_info.ndim == 1)
|
|
360
|
-
return search_one_in_index(index, vectors, wanted, exact);
|
|
361
357
|
if (vectors_info.ndim != 2)
|
|
362
358
|
throw std::invalid_argument("Expects a matrix of vectors to add!");
|
|
363
359
|
|
|
@@ -366,65 +362,193 @@ static py::tuple search_many_in_index( //
|
|
|
366
362
|
if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))
|
|
367
363
|
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
368
364
|
|
|
369
|
-
py::array_t<
|
|
370
|
-
py::array_t<distance_t>
|
|
371
|
-
py::array_t<Py_ssize_t>
|
|
365
|
+
py::array_t<key_t> keys_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
|
|
366
|
+
py::array_t<distance_t> distances_py({vectors_count, static_cast<Py_ssize_t>(wanted)});
|
|
367
|
+
py::array_t<Py_ssize_t> counts_py(vectors_count);
|
|
368
|
+
std::atomic<std::size_t> stats_visited_members(0);
|
|
369
|
+
std::atomic<std::size_t> stats_computed_distances(0);
|
|
372
370
|
|
|
371
|
+
// clang-format off
|
|
373
372
|
switch (numpy_string_to_kind(vectors_info.format)) {
|
|
374
|
-
case scalar_kind_t::b1x8_k: search_typed<b1x8_t>(index, vectors_info, wanted, exact, threads,
|
|
375
|
-
case scalar_kind_t::
|
|
376
|
-
case scalar_kind_t::f16_k: search_typed<f16_t>(index, vectors_info, wanted, exact, threads,
|
|
377
|
-
case scalar_kind_t::f32_k: search_typed<f32_t>(index, vectors_info, wanted, exact, threads,
|
|
378
|
-
case scalar_kind_t::f64_k: search_typed<f64_t>(index, vectors_info, wanted, exact, threads,
|
|
379
|
-
|
|
380
|
-
throw std::invalid_argument("Incompatible scalars in the query matrix: " + vectors_info.format);
|
|
373
|
+
case scalar_kind_t::b1x8_k: search_typed<b1x8_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances); break;
|
|
374
|
+
case scalar_kind_t::i8_k: search_typed<i8_bits_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances); break;
|
|
375
|
+
case scalar_kind_t::f16_k: search_typed<f16_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances); break;
|
|
376
|
+
case scalar_kind_t::f32_k: search_typed<f32_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances); break;
|
|
377
|
+
case scalar_kind_t::f64_k: search_typed<f64_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances); break;
|
|
378
|
+
default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + vectors_info.format);
|
|
381
379
|
}
|
|
380
|
+
// clang-format on
|
|
382
381
|
|
|
383
|
-
py::tuple results(
|
|
384
|
-
results[0] =
|
|
385
|
-
results[1] =
|
|
386
|
-
results[2] =
|
|
382
|
+
py::tuple results(5);
|
|
383
|
+
results[0] = keys_py;
|
|
384
|
+
results[1] = distances_py;
|
|
385
|
+
results[2] = counts_py;
|
|
386
|
+
results[3] = stats_visited_members.load();
|
|
387
|
+
results[4] = stats_computed_distances.load();
|
|
387
388
|
return results;
|
|
388
389
|
}
|
|
389
390
|
|
|
390
|
-
|
|
391
|
+
template <typename scalar_at>
|
|
392
|
+
static void search_typed_brute_force( //
|
|
393
|
+
py::buffer_info& dataset_info, py::buffer_info& queries_info, //
|
|
394
|
+
std::size_t wanted, std::size_t threads, metric_t const& metric, //
|
|
395
|
+
py::array_t<key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py) {
|
|
396
|
+
|
|
397
|
+
auto keys_py2d = keys_py.template mutable_unchecked<2>();
|
|
398
|
+
auto distances_py2d = distances_py.template mutable_unchecked<2>();
|
|
399
|
+
auto counts_py1d = counts_py.template mutable_unchecked<1>();
|
|
400
|
+
|
|
401
|
+
std::size_t dataset_count = static_cast<std::size_t>(dataset_info.shape[0]);
|
|
402
|
+
std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);
|
|
403
|
+
std::size_t dimensions = static_cast<std::size_t>(dataset_info.shape[1]);
|
|
404
|
+
|
|
405
|
+
byte_t const* dataset_data = reinterpret_cast<byte_t const*>(dataset_info.ptr);
|
|
406
|
+
byte_t const* queries_data = reinterpret_cast<byte_t const*>(queries_info.ptr);
|
|
407
|
+
for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)
|
|
408
|
+
counts_py1d(query_idx) = 0;
|
|
409
|
+
|
|
410
|
+
if (!threads)
|
|
411
|
+
threads = std::thread::hardware_concurrency();
|
|
412
|
+
|
|
413
|
+
std::size_t tasks_count = static_cast<std::size_t>(dataset_count * queries_count);
|
|
414
|
+
visits_bitset_t query_mutexes(static_cast<std::size_t>(queries_count));
|
|
415
|
+
if (!query_mutexes)
|
|
416
|
+
throw std::bad_alloc();
|
|
417
|
+
|
|
418
|
+
executor_default_t{threads}.execute_bulk(tasks_count, [&](std::size_t, std::size_t task_idx) {
|
|
419
|
+
//
|
|
420
|
+
std::size_t dataset_idx = task_idx / queries_count;
|
|
421
|
+
std::size_t query_idx = task_idx % queries_count;
|
|
422
|
+
|
|
423
|
+
byte_t const* dataset = dataset_data + dataset_idx * dataset_info.strides[0];
|
|
424
|
+
byte_t const* query = queries_data + query_idx * queries_info.strides[0];
|
|
425
|
+
distance_t distance = metric(dataset, query);
|
|
426
|
+
|
|
427
|
+
{
|
|
428
|
+
auto lock = query_mutexes.lock(query_idx);
|
|
429
|
+
key_t* keys = &keys_py2d(query_idx, 0);
|
|
430
|
+
distance_t* distances = &distances_py2d(query_idx, 0);
|
|
431
|
+
std::size_t& matches = reinterpret_cast<std::size_t&>(counts_py1d(query_idx));
|
|
432
|
+
if (matches == wanted)
|
|
433
|
+
if (distances[wanted - 1] <= distance)
|
|
434
|
+
return;
|
|
435
|
+
|
|
436
|
+
std::size_t offset = std::lower_bound(distances, distances + matches, distance) - distances;
|
|
437
|
+
|
|
438
|
+
std::size_t count_worse = matches - offset - (wanted == matches);
|
|
439
|
+
std::memmove(keys + offset + 1, keys + offset, count_worse * sizeof(key_t));
|
|
440
|
+
std::memmove(distances + offset + 1, distances + offset, count_worse * sizeof(distance_t));
|
|
441
|
+
keys[offset] = static_cast<key_t>(dataset_idx);
|
|
442
|
+
distances[offset] = distance;
|
|
443
|
+
matches += matches != wanted;
|
|
444
|
+
}
|
|
445
|
+
|
|
446
|
+
if (PyErr_CheckSignals() != 0)
|
|
447
|
+
throw py::error_already_set();
|
|
448
|
+
});
|
|
449
|
+
}
|
|
450
|
+
|
|
451
|
+
static py::tuple search_many_brute_force( //
|
|
452
|
+
py::buffer dataset, py::buffer queries, //
|
|
453
|
+
std::size_t wanted, std::size_t threads, //
|
|
454
|
+
metric_kind_t metric_kind, //
|
|
455
|
+
metric_signature_t metric_signature, //
|
|
456
|
+
std::uintptr_t metric_uintptr) {
|
|
457
|
+
|
|
458
|
+
if (wanted == 0)
|
|
459
|
+
return py::tuple(5);
|
|
460
|
+
|
|
461
|
+
py::buffer_info dataset_info = dataset.request();
|
|
462
|
+
py::buffer_info queries_info = queries.request();
|
|
463
|
+
if (dataset_info.ndim != 2 || queries_info.ndim != 2)
|
|
464
|
+
throw std::invalid_argument("Expects a matrix of dataset to add!");
|
|
465
|
+
|
|
466
|
+
Py_ssize_t dataset_count = dataset_info.shape[0];
|
|
467
|
+
Py_ssize_t dataset_dimensions = dataset_info.shape[1];
|
|
468
|
+
Py_ssize_t queries_count = queries_info.shape[0];
|
|
469
|
+
Py_ssize_t queries_dimensions = queries_info.shape[1];
|
|
470
|
+
if (dataset_dimensions != queries_dimensions)
|
|
471
|
+
throw std::invalid_argument("The number of vector dimensions doesn't match!");
|
|
472
|
+
|
|
473
|
+
scalar_kind_t dataset_kind = numpy_string_to_kind(dataset_info.format);
|
|
474
|
+
scalar_kind_t queries_kind = numpy_string_to_kind(queries_info.format);
|
|
475
|
+
if (dataset_kind != queries_kind)
|
|
476
|
+
throw std::invalid_argument("The types of vectors don't match!");
|
|
477
|
+
|
|
478
|
+
py::array_t<key_t> keys_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
|
|
479
|
+
py::array_t<distance_t> distances_py({dataset_count, static_cast<Py_ssize_t>(wanted)});
|
|
480
|
+
py::array_t<Py_ssize_t> counts_py(dataset_count);
|
|
481
|
+
|
|
482
|
+
std::size_t dimensions = static_cast<std::size_t>(queries_dimensions);
|
|
483
|
+
metric_t metric = //
|
|
484
|
+
metric_uintptr //
|
|
485
|
+
? udf(metric_kind, metric_signature, metric_uintptr, queries_kind, dimensions)
|
|
486
|
+
: metric_t(dimensions, metric_kind, queries_kind);
|
|
487
|
+
|
|
488
|
+
// clang-format off
|
|
489
|
+
switch (dataset_kind) {
|
|
490
|
+
case scalar_kind_t::b1x8_k: search_typed_brute_force<b1x8_t>(dataset_info, queries_info, wanted, threads, metric, keys_py, distances_py, counts_py); break;
|
|
491
|
+
case scalar_kind_t::i8_k: search_typed_brute_force<i8_bits_t>(dataset_info, queries_info, wanted, threads, metric, keys_py, distances_py, counts_py); break;
|
|
492
|
+
case scalar_kind_t::f16_k: search_typed_brute_force<f16_t>(dataset_info, queries_info, wanted, threads, metric, keys_py, distances_py, counts_py); break;
|
|
493
|
+
case scalar_kind_t::f32_k: search_typed_brute_force<f32_t>(dataset_info, queries_info, wanted, threads, metric, keys_py, distances_py, counts_py); break;
|
|
494
|
+
case scalar_kind_t::f64_k: search_typed_brute_force<f64_t>(dataset_info, queries_info, wanted, threads, metric, keys_py, distances_py, counts_py); break;
|
|
495
|
+
default: throw std::invalid_argument("Incompatible vector types: " + dataset_info.format);
|
|
496
|
+
}
|
|
497
|
+
// clang-format on
|
|
498
|
+
|
|
499
|
+
py::tuple results(5);
|
|
500
|
+
results[0] = keys_py;
|
|
501
|
+
results[1] = distances_py;
|
|
502
|
+
results[2] = counts_py;
|
|
503
|
+
results[3] = 0;
|
|
504
|
+
results[4] = static_cast<std::size_t>(dataset_count * queries_count);
|
|
505
|
+
return results;
|
|
506
|
+
}
|
|
507
|
+
|
|
508
|
+
static std::unordered_map<key_t, key_t> join_index( //
|
|
391
509
|
dense_index_py_t const& a, dense_index_py_t const& b, //
|
|
392
510
|
std::size_t max_proposals, bool exact) {
|
|
393
511
|
|
|
394
|
-
std::unordered_map<
|
|
512
|
+
std::unordered_map<key_t, key_t> a_to_b;
|
|
513
|
+
dummy_label_to_label_mapping_t b_to_a;
|
|
395
514
|
a_to_b.reserve((std::min)(a.size(), b.size()));
|
|
396
515
|
|
|
397
|
-
|
|
398
|
-
join_config_t config;
|
|
399
|
-
|
|
516
|
+
index_join_config_t config;
|
|
400
517
|
config.max_proposals = max_proposals;
|
|
401
518
|
config.exact = exact;
|
|
402
519
|
config.expansion = (std::max)(a.expansion_search(), b.expansion_search());
|
|
403
520
|
std::size_t threads = (std::min)(a.limits().threads(), b.limits().threads());
|
|
404
521
|
executor_default_t executor{threads};
|
|
405
|
-
join_result_t result =
|
|
406
|
-
a, b, config, //
|
|
407
|
-
a_to_b, //
|
|
408
|
-
dummy_label_to_label_mapping_t{}, //
|
|
409
|
-
executor);
|
|
522
|
+
join_result_t result = a.join(b, config, a_to_b, b_to_a, executor);
|
|
410
523
|
result.error.raise();
|
|
524
|
+
|
|
411
525
|
return a_to_b;
|
|
412
526
|
}
|
|
413
527
|
|
|
414
528
|
static dense_index_py_t copy_index(dense_index_py_t const& index) {
|
|
415
529
|
|
|
416
530
|
using copy_result_t = typename dense_index_py_t::copy_result_t;
|
|
417
|
-
|
|
418
|
-
|
|
419
|
-
copy_result_t result = index.copy();
|
|
531
|
+
index_copy_config_t config;
|
|
532
|
+
copy_result_t result = index.copy(config);
|
|
420
533
|
result.error.raise();
|
|
421
534
|
return std::move(result.index);
|
|
422
535
|
}
|
|
423
536
|
|
|
537
|
+
static void compact_index(dense_index_py_t& index, std::size_t threads) {
|
|
538
|
+
|
|
539
|
+
if (!threads)
|
|
540
|
+
threads = std::thread::hardware_concurrency();
|
|
541
|
+
if (!index.reserve(index_limits_t(index.size(), threads)))
|
|
542
|
+
throw std::invalid_argument("Out of memory!");
|
|
543
|
+
|
|
544
|
+
index.compact(executor_default_t{threads});
|
|
545
|
+
}
|
|
546
|
+
|
|
424
547
|
// clang-format off
|
|
425
548
|
template <typename index_at> void save_index(index_at const& index, std::string const& path) { index.save(path.c_str()).error.raise(); }
|
|
426
549
|
template <typename index_at> void load_index(index_at& index, std::string const& path) { index.load(path.c_str()).error.raise(); }
|
|
427
550
|
template <typename index_at> void view_index(index_at& index, std::string const& path) { index.view(path.c_str()).error.raise(); }
|
|
551
|
+
template <typename index_at> void reset_index(index_at& index) { index.reset(); }
|
|
428
552
|
template <typename index_at> void clear_index(index_at& index) { index.clear(); }
|
|
429
553
|
template <typename index_at> std::size_t max_level(index_at const &index) { return index.max_level(); }
|
|
430
554
|
template <typename index_at> typename index_at::stats_t compute_stats(index_at const &index) { return index.stats(); }
|
|
@@ -432,41 +556,39 @@ template <typename index_at> typename index_at::stats_t compute_level_stats(inde
|
|
|
432
556
|
// clang-format on
|
|
433
557
|
|
|
434
558
|
template <typename internal_at, typename external_at = internal_at, typename index_at = void>
|
|
435
|
-
py::object get_typed_member(index_at const& index,
|
|
436
|
-
|
|
437
|
-
py::array_t<external_at> result_py(static_cast<Py_ssize_t>(result_slots));
|
|
559
|
+
py::object get_typed_member(index_at const& index, key_t key) {
|
|
560
|
+
py::array_t<external_at> result_py(static_cast<Py_ssize_t>(index.scalar_words()));
|
|
438
561
|
auto result_py1d = result_py.template mutable_unchecked<1>();
|
|
439
|
-
if (!index.get(
|
|
562
|
+
if (!index.get(key, (internal_at*)&result_py1d(0)))
|
|
440
563
|
return py::none();
|
|
441
564
|
return py::object(result_py);
|
|
442
565
|
}
|
|
443
566
|
|
|
444
|
-
template <typename index_at> py::object get_member(index_at const& index,
|
|
567
|
+
template <typename index_at> py::object get_member(index_at const& index, key_t key, scalar_kind_t scalar_kind) {
|
|
445
568
|
if (scalar_kind == scalar_kind_t::f32_k)
|
|
446
|
-
return get_typed_member<f32_t>(index,
|
|
569
|
+
return get_typed_member<f32_t>(index, key);
|
|
447
570
|
else if (scalar_kind == scalar_kind_t::f64_k)
|
|
448
|
-
return get_typed_member<f64_t>(index,
|
|
571
|
+
return get_typed_member<f64_t>(index, key);
|
|
449
572
|
else if (scalar_kind == scalar_kind_t::f16_k)
|
|
450
|
-
return get_typed_member<f16_t, std::uint16_t>(index,
|
|
451
|
-
else if (scalar_kind == scalar_kind_t::
|
|
452
|
-
return get_typed_member<
|
|
573
|
+
return get_typed_member<f16_t, std::uint16_t>(index, key);
|
|
574
|
+
else if (scalar_kind == scalar_kind_t::i8_k)
|
|
575
|
+
return get_typed_member<i8_bits_t, std::int8_t>(index, key);
|
|
453
576
|
else if (scalar_kind == scalar_kind_t::b1x8_k)
|
|
454
|
-
return get_typed_member<b1x8_t, std::uint8_t>(index,
|
|
577
|
+
return get_typed_member<b1x8_t, std::uint8_t>(index, key);
|
|
455
578
|
else
|
|
456
579
|
throw std::invalid_argument("Incompatible scalars in the query matrix!");
|
|
457
580
|
}
|
|
458
581
|
|
|
459
|
-
template <typename index_at>
|
|
460
|
-
py::array_t<label_t> get_labels(index_at const& index, std::size_t offset, std::size_t limit) {
|
|
582
|
+
template <typename index_at> py::array_t<key_t> get_keys(index_at const& index, std::size_t offset, std::size_t limit) {
|
|
461
583
|
limit = std::min(index.size(), limit);
|
|
462
|
-
py::array_t<
|
|
584
|
+
py::array_t<key_t> result_py(static_cast<Py_ssize_t>(limit));
|
|
463
585
|
auto result_py1d = result_py.template mutable_unchecked<1>();
|
|
464
|
-
index.
|
|
586
|
+
index.export_keys(&result_py1d(0), offset, limit);
|
|
465
587
|
return result_py;
|
|
466
588
|
}
|
|
467
589
|
|
|
468
|
-
template <typename index_at> py::array_t<
|
|
469
|
-
return
|
|
590
|
+
template <typename index_at> py::array_t<key_t> get_all_keys(index_at const& index) {
|
|
591
|
+
return get_keys(index, 0, index.size());
|
|
470
592
|
}
|
|
471
593
|
|
|
472
594
|
template <typename element_at> bool has_duplicates(element_at const* begin, element_at const* end) {
|
|
@@ -481,97 +603,6 @@ template <typename element_at> bool has_duplicates(element_at const* begin, elem
|
|
|
481
603
|
return false;
|
|
482
604
|
}
|
|
483
605
|
|
|
484
|
-
void validate_set(py::array_t<set_member_t> const& set) {
|
|
485
|
-
if (set.ndim() != 1)
|
|
486
|
-
throw std::invalid_argument("Set can't be multi-dimensional!");
|
|
487
|
-
if (set.strides(0) != sizeof(set_member_t))
|
|
488
|
-
throw std::invalid_argument("Set can't be strided!");
|
|
489
|
-
auto proxy = set.unchecked<1>();
|
|
490
|
-
set_member_t const* begin = &proxy(0);
|
|
491
|
-
set_member_t const* end = begin + proxy.size();
|
|
492
|
-
if (!std::is_sorted(begin, end))
|
|
493
|
-
throw std::invalid_argument("Set must be sorted!");
|
|
494
|
-
if (has_duplicates(begin, end))
|
|
495
|
-
throw std::invalid_argument("Set must be deduplicated!");
|
|
496
|
-
}
|
|
497
|
-
#if 0
|
|
498
|
-
|
|
499
|
-
inline std::size_t hash_ror64(std::size_t v, int r) noexcept { return (v >> r) | (v << (64 - r)); }
|
|
500
|
-
|
|
501
|
-
inline std::size_t hash(std::size_t v) noexcept {
|
|
502
|
-
v ^= hash_ror64(v, 25) ^ hash_ror64(v, 50);
|
|
503
|
-
v *= 0xA24BAED4963EE407UL;
|
|
504
|
-
v ^= hash_ror64(v, 24) ^ hash_ror64(v, 49);
|
|
505
|
-
v *= 0x9FB21C651E98DF25UL;
|
|
506
|
-
return v ^ v >> 28;
|
|
507
|
-
}
|
|
508
|
-
|
|
509
|
-
template <typename scalar_at>
|
|
510
|
-
void hash_typed_row( //
|
|
511
|
-
byte_t const* scalars, std::size_t scalar_stride, std::size_t scalars_count, //
|
|
512
|
-
b1x8_t* hashes, std::size_t bits) noexcept {
|
|
513
|
-
|
|
514
|
-
std::size_t words = divide_round_up<8>(bits);
|
|
515
|
-
for (std::size_t i = 0; i != scalars_count; ++i) {
|
|
516
|
-
scalar_at scalar;
|
|
517
|
-
std::memcpy(&scalar, scalars + i * scalar_stride, sizeof(scalar_at));
|
|
518
|
-
std::size_t scalar_hash = hash(scalar);
|
|
519
|
-
hashes[scalar_hash & (words - 1)] |= std::uint8_t(1) << (scalar_hash % 8);
|
|
520
|
-
}
|
|
521
|
-
}
|
|
522
|
-
|
|
523
|
-
template <typename scalar_at>
|
|
524
|
-
py::array_t<bits_numpy_word_t> hash_typed_buffer(py::buffer_info const& info, std::size_t bits) {
|
|
525
|
-
byte_t const* data = reinterpret_cast<byte_t const*>(info.ptr);
|
|
526
|
-
Py_ssize_t elements_per_row = info.shape[0];
|
|
527
|
-
Py_ssize_t elements_stride = info.strides[0];
|
|
528
|
-
Py_ssize_t words_per_row = divide_round_up<std::size_t>(bits, 8);
|
|
529
|
-
|
|
530
|
-
if (info.ndim == 2) {
|
|
531
|
-
Py_ssize_t rows_count = info.shape[1];
|
|
532
|
-
Py_ssize_t rows_stride = info.strides[1];
|
|
533
|
-
auto hashes_shape = py_shape_t{rows_count, words_per_row};
|
|
534
|
-
auto hashes_py = py::array_t<bits_numpy_word_t>(hashes_shape);
|
|
535
|
-
auto hashes_proxy = hashes_py.template mutable_unchecked<2>();
|
|
536
|
-
|
|
537
|
-
for (Py_ssize_t i = 0; i != rows_count; ++i)
|
|
538
|
-
hash_typed_row( //
|
|
539
|
-
data + i * rows_stride, hashes_proxy.stride(1), static_cast<std::size_t>(elements_per_row), //
|
|
540
|
-
&hashes_proxy(i, 0), bits);
|
|
541
|
-
|
|
542
|
-
return hashes_py;
|
|
543
|
-
} else {
|
|
544
|
-
auto hashes_shape = py_shape_t{words_per_row};
|
|
545
|
-
auto hashes_py = py::array_t<bits_numpy_word_t>(hashes_shape);
|
|
546
|
-
auto hashes_proxy = hashes_py.template mutable_unchecked<1>();
|
|
547
|
-
|
|
548
|
-
hash_typed_row( //
|
|
549
|
-
data, hashes_info.strides[0], static_cast<std::size_t>(elements_per_row), //
|
|
550
|
-
&hashes_proxy(0), bits);
|
|
551
|
-
|
|
552
|
-
return hashes_py;
|
|
553
|
-
}
|
|
554
|
-
}
|
|
555
|
-
|
|
556
|
-
py::array_t<bits_numpy_word_t> hash_buffer(py::buffer vector, std::size_t bits) {
|
|
557
|
-
if (std::bitset<64>(bits).count() == 1)
|
|
558
|
-
throw std::invalid_argument("Number of bits must be a power of two!");
|
|
559
|
-
py::buffer_info info = vector.request();
|
|
560
|
-
if (info.ndim > 2)
|
|
561
|
-
throw std::invalid_argument("Only one or two-dimensional arrays are supported!");
|
|
562
|
-
|
|
563
|
-
// https://docs.python.org/3/library/struct.html#format-characters
|
|
564
|
-
if (info.format == "h" || info.format == "H")
|
|
565
|
-
return hash_typed_buffer<std::uint16_t>(info, bits);
|
|
566
|
-
else if (info.format == "i" || info.format == "I" || info.format == "l" || info.format == "L")
|
|
567
|
-
return hash_typed_buffer<std::uint32_t>(info, bits);
|
|
568
|
-
else if (info.format == "q" || info.format == "Q" || info.format == "n" || info.format == "N")
|
|
569
|
-
return hash_typed_buffer<std::uint64_t>(info, bits);
|
|
570
|
-
else
|
|
571
|
-
throw std::invalid_argument("Array elements must be 16, 32, or 64 bit hashable integers!");
|
|
572
|
-
}
|
|
573
|
-
#endif
|
|
574
|
-
|
|
575
606
|
PYBIND11_MODULE(compiled, m) {
|
|
576
607
|
m.doc() = "Smaller & Faster Single-File Vector Search Engine from Unum";
|
|
577
608
|
|
|
@@ -590,92 +621,109 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
590
621
|
|
|
591
622
|
py::enum_<metric_kind_t>(m, "MetricKind")
|
|
592
623
|
.value("Unknown", metric_kind_t::unknown_k)
|
|
593
|
-
|
|
624
|
+
|
|
625
|
+
.value("IP", metric_kind_t::cos_k)
|
|
594
626
|
.value("Cos", metric_kind_t::cos_k)
|
|
595
627
|
.value("L2sq", metric_kind_t::l2sq_k)
|
|
628
|
+
|
|
596
629
|
.value("Haversine", metric_kind_t::haversine_k)
|
|
597
630
|
.value("Pearson", metric_kind_t::pearson_k)
|
|
598
631
|
.value("Jaccard", metric_kind_t::jaccard_k)
|
|
599
632
|
.value("Hamming", metric_kind_t::hamming_k)
|
|
600
633
|
.value("Tanimoto", metric_kind_t::tanimoto_k)
|
|
601
|
-
.value("Sorensen", metric_kind_t::sorensen_k)
|
|
634
|
+
.value("Sorensen", metric_kind_t::sorensen_k)
|
|
635
|
+
|
|
636
|
+
.value("Cosine", metric_kind_t::cos_k)
|
|
637
|
+
.value("InnerProduct", metric_kind_t::cos_k);
|
|
602
638
|
|
|
603
639
|
py::enum_<scalar_kind_t>(m, "ScalarKind")
|
|
640
|
+
.value("Unknown", scalar_kind_t::unknown_k)
|
|
641
|
+
.value("B1", scalar_kind_t::b1x8_k)
|
|
642
|
+
.value("U40", scalar_kind_t::u40_k)
|
|
643
|
+
.value("UUID", scalar_kind_t::uuid_k)
|
|
604
644
|
.value("F64", scalar_kind_t::f64_k)
|
|
605
645
|
.value("F32", scalar_kind_t::f32_k)
|
|
606
646
|
.value("F16", scalar_kind_t::f16_k)
|
|
607
647
|
.value("F8", scalar_kind_t::f8_k)
|
|
608
|
-
.value("
|
|
609
|
-
|
|
610
|
-
|
|
611
|
-
|
|
612
|
-
|
|
613
|
-
|
|
614
|
-
|
|
615
|
-
|
|
616
|
-
|
|
617
|
-
|
|
618
|
-
|
|
619
|
-
|
|
620
|
-
|
|
648
|
+
.value("U64", scalar_kind_t::u64_k)
|
|
649
|
+
.value("U32", scalar_kind_t::u32_k)
|
|
650
|
+
.value("U16", scalar_kind_t::u16_k)
|
|
651
|
+
.value("U8", scalar_kind_t::u8_k)
|
|
652
|
+
.value("I64", scalar_kind_t::i64_k)
|
|
653
|
+
.value("I32", scalar_kind_t::i32_k)
|
|
654
|
+
.value("I16", scalar_kind_t::i16_k)
|
|
655
|
+
.value("I8", scalar_kind_t::i8_k);
|
|
656
|
+
|
|
657
|
+
m.def("index_dense_metadata", [](std::string const& path) -> py::dict {
|
|
658
|
+
index_dense_metadata_result_t meta = index_dense_metadata(path.c_str());
|
|
659
|
+
meta.error.raise();
|
|
660
|
+
index_dense_head_t const& head = meta.head;
|
|
661
|
+
|
|
662
|
+
py::dict result;
|
|
663
|
+
result["matrix_included"] = !meta.config.exclude_vectors;
|
|
664
|
+
result["matrix_uses_64_bit_dimensions"] = meta.config.use_64_bit_dimensions;
|
|
665
|
+
|
|
666
|
+
result["version"] = std::to_string(head.version_major) + "." + //
|
|
667
|
+
std::to_string(head.version_minor) + "." + //
|
|
668
|
+
std::to_string(head.version_patch);
|
|
669
|
+
|
|
670
|
+
result["kind_metric"] = metric_kind_t(head.kind_metric);
|
|
671
|
+
result["kind_scalar"] = scalar_kind_t(head.kind_scalar);
|
|
672
|
+
result["kind_key"] = scalar_kind_t(head.kind_key);
|
|
673
|
+
result["kind_compressed_slot"] = scalar_kind_t(head.kind_compressed_slot);
|
|
674
|
+
|
|
675
|
+
result["count_present"] = std::uint64_t(head.count_present);
|
|
676
|
+
result["count_deleted"] = std::uint64_t(head.count_deleted);
|
|
677
|
+
result["dimensions"] = std::uint64_t(head.dimensions);
|
|
678
|
+
|
|
679
|
+
return result;
|
|
621
680
|
});
|
|
622
|
-
|
|
623
|
-
|
|
624
|
-
|
|
625
|
-
|
|
626
|
-
|
|
627
|
-
|
|
628
|
-
|
|
629
|
-
|
|
630
|
-
|
|
631
|
-
|
|
632
|
-
|
|
633
|
-
|
|
634
|
-
|
|
635
|
-
auto i = py::class_<dense_index_py_t>(m, "Index");
|
|
681
|
+
|
|
682
|
+
m.def("exact_search", &search_many_brute_force, //
|
|
683
|
+
py::arg("dataset"), //
|
|
684
|
+
py::arg("queries"), //
|
|
685
|
+
py::arg("count") = 10, //
|
|
686
|
+
py::kw_only(), //
|
|
687
|
+
py::arg("threads") = 0, //
|
|
688
|
+
py::arg("metric_kind") = metric_kind_t::cos_k, //
|
|
689
|
+
py::arg("metric_signature") = metric_signature_t::array_array_k, //
|
|
690
|
+
py::arg("metric_pointer") = 0 //
|
|
691
|
+
);
|
|
692
|
+
|
|
693
|
+
auto i = py::class_<dense_index_py_t, std::shared_ptr<dense_index_py_t>>(m, "Index");
|
|
636
694
|
|
|
637
695
|
i.def(py::init(&make_index), //
|
|
638
696
|
py::kw_only(), //
|
|
639
697
|
py::arg("ndim") = 0, //
|
|
640
698
|
py::arg("dtype") = scalar_kind_t::f32_k, //
|
|
641
|
-
py::arg("metric") = metric_kind_t::ip_k, //
|
|
642
699
|
py::arg("connectivity") = default_connectivity(), //
|
|
643
700
|
py::arg("expansion_add") = default_expansion_add(), //
|
|
644
701
|
py::arg("expansion_search") = default_expansion_search(), //
|
|
702
|
+
py::arg("metric_kind") = metric_kind_t::cos_k, //
|
|
645
703
|
py::arg("metric_signature") = metric_signature_t::array_array_k, //
|
|
646
|
-
py::arg("metric_pointer") = 0
|
|
647
|
-
py::arg("tune") = false //
|
|
648
|
-
);
|
|
649
|
-
|
|
650
|
-
i.def( //
|
|
651
|
-
"add", &add_many_to_index, //
|
|
652
|
-
py::arg("labels"), //
|
|
653
|
-
py::arg("vectors"), //
|
|
654
|
-
py::kw_only(), //
|
|
655
|
-
py::arg("copy") = true, //
|
|
656
|
-
py::arg("threads") = 0 //
|
|
704
|
+
py::arg("metric_pointer") = 0 //
|
|
657
705
|
);
|
|
658
706
|
|
|
659
|
-
i.def(
|
|
660
|
-
"add", &
|
|
661
|
-
py::arg("
|
|
662
|
-
py::arg("
|
|
663
|
-
py::kw_only(),
|
|
664
|
-
py::arg("copy") = true,
|
|
665
|
-
py::arg("threads") = 0
|
|
707
|
+
i.def( //
|
|
708
|
+
"add", &add_many_to_index<dense_index_py_t>, //
|
|
709
|
+
py::arg("keys"), //
|
|
710
|
+
py::arg("vectors"), //
|
|
711
|
+
py::kw_only(), //
|
|
712
|
+
py::arg("copy") = true, //
|
|
713
|
+
py::arg("threads") = 0 //
|
|
666
714
|
);
|
|
667
715
|
|
|
668
|
-
i.def(
|
|
669
|
-
"search", &search_many_in_index
|
|
670
|
-
py::arg("query"),
|
|
671
|
-
py::arg("count") = 10,
|
|
672
|
-
py::arg("exact") = false,
|
|
673
|
-
py::arg("threads") = 0
|
|
716
|
+
i.def( //
|
|
717
|
+
"search", &search_many_in_index<dense_index_py_t>, //
|
|
718
|
+
py::arg("query"), //
|
|
719
|
+
py::arg("count") = 10, //
|
|
720
|
+
py::arg("exact") = false, //
|
|
721
|
+
py::arg("threads") = 0 //
|
|
674
722
|
);
|
|
675
723
|
|
|
676
724
|
i.def(
|
|
677
725
|
"rename",
|
|
678
|
-
[](dense_index_py_t& index,
|
|
726
|
+
[](dense_index_py_t& index, key_t from, key_t to) -> bool {
|
|
679
727
|
dense_labeling_result_t result = index.rename(from, to);
|
|
680
728
|
result.error.raise();
|
|
681
729
|
return result.completed;
|
|
@@ -684,8 +732,8 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
684
732
|
|
|
685
733
|
i.def(
|
|
686
734
|
"remove",
|
|
687
|
-
[](dense_index_py_t& index,
|
|
688
|
-
dense_labeling_result_t result = index.remove(
|
|
735
|
+
[](dense_index_py_t& index, key_t key, bool compact, std::size_t threads) -> bool {
|
|
736
|
+
dense_labeling_result_t result = index.remove(key);
|
|
689
737
|
result.error.raise();
|
|
690
738
|
if (!compact)
|
|
691
739
|
return result.completed;
|
|
@@ -695,16 +743,15 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
695
743
|
if (!index.reserve(index_limits_t(index.size(), threads)))
|
|
696
744
|
throw std::invalid_argument("Out of memory!");
|
|
697
745
|
|
|
698
|
-
index.
|
|
746
|
+
index.isolate(executor_default_t{threads});
|
|
699
747
|
return result.completed;
|
|
700
748
|
},
|
|
701
|
-
py::arg("
|
|
749
|
+
py::arg("key"), py::arg("compact"), py::arg("threads"));
|
|
702
750
|
|
|
703
751
|
i.def(
|
|
704
752
|
"remove",
|
|
705
|
-
[](dense_index_py_t& index, std::vector<
|
|
706
|
-
|
|
707
|
-
dense_labeling_result_t result = index.remove(labels.begin(), labels.end());
|
|
753
|
+
[](dense_index_py_t& index, std::vector<key_t> const& keys, bool compact, std::size_t threads) -> std::size_t {
|
|
754
|
+
dense_labeling_result_t result = index.remove(keys.begin(), keys.end());
|
|
708
755
|
result.error.raise();
|
|
709
756
|
if (!compact)
|
|
710
757
|
return result.completed;
|
|
@@ -714,18 +761,19 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
714
761
|
if (!index.reserve(index_limits_t(index.size(), threads)))
|
|
715
762
|
throw std::invalid_argument("Out of memory!");
|
|
716
763
|
|
|
717
|
-
index.
|
|
764
|
+
index.isolate(executor_default_t{threads});
|
|
718
765
|
return result.completed;
|
|
719
766
|
},
|
|
720
|
-
py::arg("
|
|
767
|
+
py::arg("key"), py::arg("compact"), py::arg("threads"));
|
|
721
768
|
|
|
722
769
|
i.def("__len__", &dense_index_py_t::size);
|
|
723
770
|
i.def_property_readonly("size", &dense_index_py_t::size);
|
|
724
|
-
i.def_property_readonly("ndim", &dense_index_py_t::dimensions);
|
|
725
771
|
i.def_property_readonly("connectivity", &dense_index_py_t::connectivity);
|
|
726
772
|
i.def_property_readonly("capacity", &dense_index_py_t::capacity);
|
|
773
|
+
i.def_property_readonly("ndim",
|
|
774
|
+
[](dense_index_py_t const& index) -> std::size_t { return index.metric().dimensions(); });
|
|
727
775
|
i.def_property_readonly( //
|
|
728
|
-
"dtype", [](dense_index_py_t const& index) -> scalar_kind_t { return index.
|
|
776
|
+
"dtype", [](dense_index_py_t const& index) -> scalar_kind_t { return index.scalar_kind(); });
|
|
729
777
|
i.def_property_readonly( //
|
|
730
778
|
"memory_usage", [](dense_index_py_t const& index) -> std::size_t { return index.memory_usage(); },
|
|
731
779
|
py::call_guard<py::gil_scoped_release>());
|
|
@@ -733,17 +781,36 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
733
781
|
i.def_property("expansion_add", &dense_index_py_t::expansion_add, &dense_index_py_t::change_expansion_add);
|
|
734
782
|
i.def_property("expansion_search", &dense_index_py_t::expansion_search, &dense_index_py_t::change_expansion_search);
|
|
735
783
|
|
|
736
|
-
i.
|
|
737
|
-
|
|
784
|
+
i.def(
|
|
785
|
+
"change_metric",
|
|
786
|
+
[](dense_index_py_t& index, metric_kind_t metric_kind, metric_signature_t metric_signature,
|
|
787
|
+
std::uintptr_t metric_uintptr) {
|
|
788
|
+
scalar_kind_t scalar_kind = index.scalar_kind();
|
|
789
|
+
std::size_t dimensions = index.dimensions();
|
|
790
|
+
metric_t metric = //
|
|
791
|
+
metric_uintptr //
|
|
792
|
+
? udf(metric_kind, metric_signature, metric_uintptr, scalar_kind, dimensions)
|
|
793
|
+
: metric_t(dimensions, metric_kind, scalar_kind);
|
|
794
|
+
index.change_metric(std::move(metric));
|
|
795
|
+
},
|
|
796
|
+
py::arg("metric_kind") = metric_kind_t::cos_k, //
|
|
797
|
+
py::arg("metric_signature") = metric_signature_t::array_array_k, //
|
|
798
|
+
py::arg("metric_pointer") = 0 //
|
|
799
|
+
);
|
|
800
|
+
|
|
801
|
+
i.def_property_readonly("keys", &get_all_keys<dense_index_py_t>);
|
|
802
|
+
i.def("get_keys", &get_keys<dense_index_py_t>, py::arg("offset") = 0,
|
|
738
803
|
py::arg("limit") = std::numeric_limits<std::size_t>::max());
|
|
739
804
|
i.def("__contains__", &dense_index_py_t::contains);
|
|
740
|
-
i.def("__getitem__", &get_member<dense_index_py_t>, py::arg("
|
|
805
|
+
i.def("__getitem__", &get_member<dense_index_py_t>, py::arg("key"), py::arg("dtype") = scalar_kind_t::f32_k);
|
|
741
806
|
|
|
742
807
|
i.def("save", &save_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
|
|
743
808
|
i.def("load", &load_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
|
|
744
809
|
i.def("view", &view_index<dense_index_py_t>, py::arg("path"), py::call_guard<py::gil_scoped_release>());
|
|
810
|
+
i.def("reset", &reset_index<dense_index_py_t>, py::call_guard<py::gil_scoped_release>());
|
|
745
811
|
i.def("clear", &clear_index<dense_index_py_t>, py::call_guard<py::gil_scoped_release>());
|
|
746
812
|
i.def("copy", ©_index, py::call_guard<py::gil_scoped_release>());
|
|
813
|
+
i.def("compact", &compact_index, py::call_guard<py::gil_scoped_release>());
|
|
747
814
|
i.def("join", &join_index, py::arg("other"), py::arg("max_proposals") = 0, py::arg("exact") = false,
|
|
748
815
|
py::call_guard<py::gil_scoped_release>());
|
|
749
816
|
|
|
@@ -758,60 +825,17 @@ PYBIND11_MODULE(compiled, m) {
|
|
|
758
825
|
i.def_property_readonly("levels_stats", &compute_stats<dense_index_py_t>);
|
|
759
826
|
i.def("level_stats", &compute_level_stats<dense_index_py_t>, py::arg("level"));
|
|
760
827
|
|
|
761
|
-
auto
|
|
762
|
-
|
|
763
|
-
|
|
764
|
-
|
|
765
|
-
|
|
766
|
-
|
|
767
|
-
|
|
768
|
-
|
|
769
|
-
|
|
770
|
-
|
|
771
|
-
|
|
772
|
-
"
|
|
773
|
-
[](sparse_index_py_t& index, label_t label, py::array_t<set_member_t> set, bool copy) {
|
|
774
|
-
validate_set(set);
|
|
775
|
-
if (index.size() + 1 >= index.capacity())
|
|
776
|
-
index.reserve(ceil2(index.size() + 1));
|
|
777
|
-
auto proxy = set.unchecked<1>();
|
|
778
|
-
auto view = set_view_t{&proxy(0), static_cast<std::size_t>(proxy.shape(0))};
|
|
779
|
-
|
|
780
|
-
add_config_t config;
|
|
781
|
-
config.store_vector = copy;
|
|
782
|
-
index.add(label, view, config).error.raise();
|
|
783
|
-
}, //
|
|
784
|
-
py::arg("label"), //
|
|
785
|
-
py::arg("set"), //
|
|
786
|
-
py::kw_only(), //
|
|
787
|
-
py::arg("copy") = true //
|
|
828
|
+
auto is = py::class_<dense_indexes_py_t>(m, "Indexes");
|
|
829
|
+
is.def(py::init());
|
|
830
|
+
is.def("__len__", &dense_indexes_py_t::size);
|
|
831
|
+
is.def("add", &dense_indexes_py_t::add);
|
|
832
|
+
is.def("add_paths", &dense_indexes_py_t::add_paths, py::arg("paths"), py::arg("view") = true,
|
|
833
|
+
py::arg("threads") = 0);
|
|
834
|
+
is.def( //
|
|
835
|
+
"search", &search_many_in_index<dense_indexes_py_t>, //
|
|
836
|
+
py::arg("query"), //
|
|
837
|
+
py::arg("count") = 10, //
|
|
838
|
+
py::arg("exact") = false, //
|
|
839
|
+
py::arg("threads") = 0 //
|
|
788
840
|
);
|
|
789
|
-
|
|
790
|
-
si.def( //
|
|
791
|
-
"search",
|
|
792
|
-
[](sparse_index_py_t& index, py::array_t<set_member_t> set, std::size_t count) -> py::array_t<label_t> {
|
|
793
|
-
validate_set(set);
|
|
794
|
-
auto proxy = set.unchecked<1>();
|
|
795
|
-
auto view = set_view_t{&proxy(0), static_cast<std::size_t>(proxy.shape(0))};
|
|
796
|
-
auto labels_py = py::array_t<label_t>(py_shape_t{static_cast<Py_ssize_t>(count)});
|
|
797
|
-
auto labels_proxy = labels_py.template mutable_unchecked<1>();
|
|
798
|
-
auto result = index.search(view, count);
|
|
799
|
-
result.error.raise();
|
|
800
|
-
auto found = result.dump_to(&labels_proxy(0));
|
|
801
|
-
labels_py.resize(py_shape_t{static_cast<Py_ssize_t>(found)});
|
|
802
|
-
return labels_py;
|
|
803
|
-
},
|
|
804
|
-
py::arg("set"), //
|
|
805
|
-
py::arg("count") = 10 //
|
|
806
|
-
);
|
|
807
|
-
|
|
808
|
-
si.def("__len__", &sparse_index_py_t::size);
|
|
809
|
-
si.def_property_readonly("size", &sparse_index_py_t::size);
|
|
810
|
-
si.def_property_readonly("connectivity", &sparse_index_py_t::connectivity);
|
|
811
|
-
si.def_property_readonly("capacity", &sparse_index_py_t::capacity);
|
|
812
|
-
|
|
813
|
-
si.def("save", &save_index<sparse_index_py_t>, py::arg("path"));
|
|
814
|
-
si.def("load", &load_index<sparse_index_py_t>, py::arg("path"));
|
|
815
|
-
si.def("view", &view_index<sparse_index_py_t>, py::arg("path"));
|
|
816
|
-
si.def("clear", &sparse_index_py_t::clear);
|
|
817
841
|
}
|