deglib 0.2.2__tar.gz → 0.2.3__tar.gz

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.
Files changed (89) hide show
  1. {deglib-0.2.2/src/deglib.egg-info → deglib-0.2.3}/PKG-INFO +1 -1
  2. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/deglib.h +1 -0
  3. deglib-0.2.3/lib/deglib/include/deglib/search/searcher.h +454 -0
  4. {deglib-0.2.2 → deglib-0.2.3}/lib/test/CMakeLists.txt +1 -0
  5. {deglib-0.2.2 → deglib-0.2.3}/src/deg_cpp/deglib_cpp.cpp +180 -0
  6. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/__init__.py +4 -1
  7. deglib-0.2.3/src/deglib/search.py +208 -0
  8. {deglib-0.2.2 → deglib-0.2.3/src/deglib.egg-info}/PKG-INFO +1 -1
  9. {deglib-0.2.2 → deglib-0.2.3}/src/deglib.egg-info/SOURCES.txt +2 -0
  10. deglib-0.2.3/tests/test_searcher.py +162 -0
  11. deglib-0.2.2/src/deglib/search.py +0 -93
  12. {deglib-0.2.2 → deglib-0.2.3}/CMakeLists.txt +0 -0
  13. {deglib-0.2.2 → deglib-0.2.3}/LICENSE +0 -0
  14. {deglib-0.2.2 → deglib-0.2.3}/MANIFEST.in +0 -0
  15. {deglib-0.2.2 → deglib-0.2.3}/README.md +0 -0
  16. {deglib-0.2.2 → deglib-0.2.3}/lib/CMakeLists.txt +0 -0
  17. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/CMakeLists.txt +0 -0
  18. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/benchmark.h +0 -0
  19. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/build.h +0 -0
  20. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/dataset.h +0 -0
  21. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/file_io.h +0 -0
  22. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/logging.h +0 -0
  23. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/repository.h +0 -0
  24. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/stats.h +0 -0
  25. {deglib-0.2.2 → deglib-0.2.3}/lib/bench/include/stopwatch.h +0 -0
  26. {deglib-0.2.2 → deglib-0.2.3}/lib/cmake_modules/DetectCPUFeatures.cmake +0 -0
  27. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/CMakeLists.txt +0 -0
  28. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/analysis.h +0 -0
  29. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/builder.h +0 -0
  30. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/concurrent.h +0 -0
  31. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/config.h +0 -0
  32. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/evp_ip.h +0 -0
  33. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/fp16.h +0 -0
  34. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/fp16_ip.h +0 -0
  35. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/fp16_l2.h +0 -0
  36. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/fp32.h +0 -0
  37. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/fp32_ip.h +0 -0
  38. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/fp32_l2.h +0 -0
  39. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/int8.h +0 -0
  40. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/int8_ip.h +0 -0
  41. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/int8_l2.h +0 -0
  42. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/residual_mode.h +0 -0
  43. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/uint8.h +0 -0
  44. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/uint8_ip.h +0 -0
  45. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distance/uint8_l2.h +0 -0
  46. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/distances.h +0 -0
  47. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/filter.h +0 -0
  48. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/dynamic_graph.h +0 -0
  49. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/internal_graph.h +0 -0
  50. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/intrusive_list.h +0 -0
  51. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/mutable_graph.h +0 -0
  52. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/readonly_graph.h +0 -0
  53. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/sizebounded_graph.h +0 -0
  54. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph/visited_list_pool.h +0 -0
  55. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/graph.h +0 -0
  56. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/flas/fast_linear_assignment_sorter.h +0 -0
  57. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/flas/fast_linear_assignment_sorter_mt.h +0 -0
  58. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/flas/junker_volgenant_solver.h +0 -0
  59. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/pruning.h +0 -0
  60. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/quantization/evp_quantize.h +0 -0
  61. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/quantization/scalar_quantize.h +0 -0
  62. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization/transform.h +0 -0
  63. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/optimization.h +0 -0
  64. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/search.h +0 -0
  65. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/utils/cpu.h +0 -0
  66. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/utils/memory.h +0 -0
  67. {deglib-0.2.2 → deglib-0.2.3}/lib/deglib/include/deglib/utils/random.h +0 -0
  68. {deglib-0.2.2 → deglib-0.2.3}/lib/test/src/common/test_helpers.h +0 -0
  69. {deglib-0.2.2 → deglib-0.2.3}/pyproject.toml +0 -0
  70. {deglib-0.2.2 → deglib-0.2.3}/setup.cfg +0 -0
  71. {deglib-0.2.2 → deglib-0.2.3}/setup.py +0 -0
  72. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/analysis.py +0 -0
  73. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/builder.py +0 -0
  74. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/cpu.py +0 -0
  75. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/distances.py +0 -0
  76. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/graph.py +0 -0
  77. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/optimization.py +0 -0
  78. {deglib-0.2.2 → deglib-0.2.3}/src/deglib/utils.py +0 -0
  79. {deglib-0.2.2 → deglib-0.2.3}/src/deglib.egg-info/dependency_links.txt +0 -0
  80. {deglib-0.2.2 → deglib-0.2.3}/src/deglib.egg-info/requires.txt +0 -0
  81. {deglib-0.2.2 → deglib-0.2.3}/src/deglib.egg-info/top_level.txt +0 -0
  82. {deglib-0.2.2 → deglib-0.2.3}/tests/test_analysis.py +0 -0
  83. {deglib-0.2.2 → deglib-0.2.3}/tests/test_builder.py +0 -0
  84. {deglib-0.2.2 → deglib-0.2.3}/tests/test_filter.py +0 -0
  85. {deglib-0.2.2 → deglib-0.2.3}/tests/test_graph.py +0 -0
  86. {deglib-0.2.2 → deglib-0.2.3}/tests/test_presort.py +0 -0
  87. {deglib-0.2.2 → deglib-0.2.3}/tests/test_rerank.py +0 -0
  88. {deglib-0.2.2 → deglib-0.2.3}/tests/test_scalar_quantizer.py +0 -0
  89. {deglib-0.2.2 → deglib-0.2.3}/tests/test_transform.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: deglib
3
- Version: 0.2.2
3
+ Version: 0.2.3
4
4
  Summary: Python bindings for the Dynamic Exploration Graph library by Nico Hezel
5
5
  Author: Nico Hezel
6
6
  Author-email: Bruno Schilling <bruno.schilling@htw-berlin.de>
@@ -8,6 +8,7 @@
8
8
  #include "deglib/graph.h"
9
9
  #include "deglib/optimization.h"
10
10
  #include "deglib/search.h"
11
+ #include "deglib/search/searcher.h"
11
12
  #include "deglib/utils/memory.h"
12
13
  #include "deglib/utils/random.h"
13
14
 
@@ -0,0 +1,454 @@
1
+ #pragma once
2
+
3
+ #include "deglib/concurrent.h"
4
+ #include "deglib/distance/fp16.h"
5
+ #include "deglib/distances.h"
6
+ #include "deglib/filter.h"
7
+ #include "deglib/graph/internal_graph.h"
8
+ #include "deglib/optimization/quantization/evp_quantize.h"
9
+ #include "deglib/optimization/quantization/scalar_quantize.h"
10
+ #include "deglib/search.h"
11
+
12
+ #include <algorithm>
13
+ #include <cstddef>
14
+ #include <cstdint>
15
+ #include <limits>
16
+ #include <memory>
17
+ #include <optional>
18
+ #include <span>
19
+ #include <stdexcept>
20
+ #include <string>
21
+ #include <type_traits>
22
+ #include <utility>
23
+ #include <vector>
24
+
25
+ namespace deglib::search {
26
+
27
+ // ============================================================================
28
+ // Search Result Struct for Modern C++20 API
29
+ // ============================================================================
30
+
31
+ struct SearchResult {
32
+ std::vector<uint32_t> indices;
33
+ std::vector<float> distances;
34
+
35
+ size_t size() const noexcept { return indices.size(); }
36
+ bool empty() const noexcept { return indices.empty(); }
37
+ bool has_distances() const noexcept { return !distances.empty(); }
38
+ };
39
+
40
+ // ============================================================================
41
+ // Quantizer Concept Wrappers
42
+ // ============================================================================
43
+
44
+ struct NoQuantizer {
45
+ template <typename InT, typename OutByteT>
46
+ static inline void transform(const InT* in, uint32_t dim, OutByteT* out_bytes) {
47
+ std::memcpy(out_bytes, in, dim * sizeof(InT));
48
+ }
49
+ };
50
+
51
+ struct ScalarInt8Quantizer {
52
+ deglib::quantization::scalar::ScalarQuantizerInt8 q;
53
+ explicit ScalarInt8Quantizer(deglib::quantization::scalar::ScalarQuantizerInt8 quant) : q(quant) {}
54
+
55
+ template <typename InT, typename OutByteT>
56
+ inline void transform(const InT* in, uint32_t dim, OutByteT* out_bytes) const {
57
+ int8_t* dst = reinterpret_cast<int8_t*>(out_bytes);
58
+ if constexpr (std::is_same_v<InT, float>) {
59
+ for (uint32_t i = 0; i < dim; ++i) dst[i] = q.transform(in[i]);
60
+ } else if constexpr (std::is_same_v<InT, uint16_t>) {
61
+ for (uint32_t i = 0; i < dim; ++i) {
62
+ dst[i] = q.transform(deglib::distances::fp16::fp16_to_float(in[i]));
63
+ }
64
+ }
65
+ }
66
+ };
67
+
68
+ struct ScalarInt8PerDimQuantizer {
69
+ deglib::quantization::scalar::ScalarQuantizerInt8PerDim q;
70
+ explicit ScalarInt8PerDimQuantizer(deglib::quantization::scalar::ScalarQuantizerInt8PerDim quant) : q(std::move(quant)) {}
71
+
72
+ template <typename InT, typename OutByteT>
73
+ inline void transform(const InT* in, uint32_t dim, OutByteT* out_bytes) const {
74
+ int8_t* dst = reinterpret_cast<int8_t*>(out_bytes);
75
+ if constexpr (std::is_same_v<InT, float>) {
76
+ for (uint32_t i = 0; i < dim; ++i) dst[i] = q.transform(in[i], i);
77
+ } else if constexpr (std::is_same_v<InT, uint16_t>) {
78
+ for (uint32_t i = 0; i < dim; ++i) {
79
+ dst[i] = q.transform(deglib::distances::fp16::fp16_to_float(in[i]), i);
80
+ }
81
+ }
82
+ }
83
+ };
84
+
85
+ struct ScalarUint8Quantizer {
86
+ deglib::quantization::scalar::ScalarQuantizerUint8 q;
87
+ explicit ScalarUint8Quantizer(deglib::quantization::scalar::ScalarQuantizerUint8 quant) : q(quant) {}
88
+
89
+ template <typename InT, typename OutByteT>
90
+ inline void transform(const InT* in, uint32_t dim, OutByteT* out_bytes) const {
91
+ uint8_t* dst = reinterpret_cast<uint8_t*>(out_bytes);
92
+ if constexpr (std::is_same_v<InT, float>) {
93
+ for (uint32_t i = 0; i < dim; ++i) dst[i] = q.transform(in[i]);
94
+ } else if constexpr (std::is_same_v<InT, uint16_t>) {
95
+ for (uint32_t i = 0; i < dim; ++i) {
96
+ dst[i] = q.transform(deglib::distances::fp16::fp16_to_float(in[i]));
97
+ }
98
+ }
99
+ }
100
+ };
101
+
102
+ struct ScalarUint8PerDimQuantizer {
103
+ deglib::quantization::scalar::ScalarQuantizerUint8PerDim q;
104
+ explicit ScalarUint8PerDimQuantizer(deglib::quantization::scalar::ScalarQuantizerUint8PerDim quant) : q(std::move(quant)) {}
105
+
106
+ template <typename InT, typename OutByteT>
107
+ inline void transform(const InT* in, uint32_t dim, OutByteT* out_bytes) const {
108
+ uint8_t* dst = reinterpret_cast<uint8_t*>(out_bytes);
109
+ if constexpr (std::is_same_v<InT, float>) {
110
+ for (uint32_t i = 0; i < dim; ++i) dst[i] = q.transform(in[i], i);
111
+ } else if constexpr (std::is_same_v<InT, uint16_t>) {
112
+ for (uint32_t i = 0; i < dim; ++i) {
113
+ dst[i] = q.transform(deglib::distances::fp16::fp16_to_float(in[i]), i);
114
+ }
115
+ }
116
+ }
117
+ };
118
+
119
+ struct EVPQuantizer {
120
+ uint32_t non_zeros = 0;
121
+ explicit EVPQuantizer(uint32_t nz) : non_zeros(nz) {}
122
+
123
+ template <typename InT, typename OutByteT>
124
+ inline void transform(const InT* in, uint32_t dim, OutByteT* out_bytes) const {
125
+ if constexpr (std::is_same_v<InT, float>) {
126
+ auto evp = deglib::quantization::evp::quantize_single(in, dim, non_zeros);
127
+ std::memcpy(out_bytes, evp.data(), evp.size());
128
+ } else if constexpr (std::is_same_v<InT, uint16_t>) {
129
+ std::vector<float> f32(dim);
130
+ deglib::distances::fp16::fp16_to_floats(in, f32.data(), dim);
131
+ auto evp = deglib::quantization::evp::quantize_single(f32.data(), dim, non_zeros);
132
+ std::memcpy(out_bytes, evp.data(), evp.size());
133
+ }
134
+ }
135
+ };
136
+
137
+ // ============================================================================
138
+ // Refiner Concept Wrappers
139
+ // ============================================================================
140
+
141
+ struct NoRefiner {
142
+ static constexpr bool enabled = false;
143
+
144
+ template <typename QueryT>
145
+ static inline uint32_t rerank(
146
+ const QueryT*, uint32_t, const uint32_t*, size_t, uint32_t,
147
+ uint32_t*, float*, bool
148
+ ) {
149
+ return 0;
150
+ }
151
+ };
152
+
153
+ template <typename RefineDataT>
154
+ struct ExactRefiner {
155
+ static constexpr bool enabled = true;
156
+ deglib::distances::FloatSpace space;
157
+ const RefineDataT* base_vectors = nullptr;
158
+ size_t num_base_vectors = 0;
159
+
160
+ ExactRefiner(deglib::distances::FloatSpace sp, const RefineDataT* base, size_t count)
161
+ : space(std::move(sp)), base_vectors(base), num_base_vectors(count) {}
162
+
163
+ template <typename QueryT>
164
+ inline uint32_t rerank(
165
+ const QueryT* query, uint32_t dim, const uint32_t* candidate_indices, size_t num_cands,
166
+ uint32_t k, uint32_t* out_indices, float* out_distances, bool unsorted
167
+ ) const {
168
+ if (!base_vectors || num_cands == 0) return 0;
169
+
170
+ if constexpr (std::is_same_v<QueryT, RefineDataT>) {
171
+ auto reranked = deglib::search::rerank(
172
+ space, query, 1, base_vectors, num_base_vectors, candidate_indices, num_cands, k, 1
173
+ );
174
+ return populate_results(reranked, k, out_indices, out_distances, unsorted);
175
+ } else if constexpr (std::is_same_v<QueryT, float> && std::is_same_v<RefineDataT, uint16_t>) {
176
+ std::vector<uint16_t> q_fp16(dim);
177
+ deglib::distances::fp16::floats_to_fp16(query, q_fp16.data(), dim);
178
+ auto reranked = deglib::search::rerank(
179
+ space, q_fp16.data(), 1, base_vectors, num_base_vectors, candidate_indices, num_cands, k, 1
180
+ );
181
+ return populate_results(reranked, k, out_indices, out_distances, unsorted);
182
+ } else if constexpr (std::is_same_v<QueryT, uint16_t> && std::is_same_v<RefineDataT, float>) {
183
+ std::vector<float> q_f32(dim);
184
+ deglib::distances::fp16::fp16_to_floats(query, q_f32.data(), dim);
185
+ auto reranked = deglib::search::rerank(
186
+ space, q_f32.data(), 1, base_vectors, num_base_vectors, candidate_indices, num_cands, k, 1
187
+ );
188
+ return populate_results(reranked, k, out_indices, out_distances, unsorted);
189
+ }
190
+ return 0;
191
+ }
192
+
193
+ private:
194
+ static inline uint32_t populate_results(
195
+ std::vector<ResultSet>& reranked, uint32_t k, uint32_t* out_indices, float* out_distances, bool unsorted
196
+ ) {
197
+ if (reranked.empty()) return 0;
198
+ auto& heap = reranked[0];
199
+ const size_t top_n = std::min<size_t>(k, heap.size());
200
+
201
+ if (unsorted) {
202
+ for (size_t i = 0; i < top_n; ++i) {
203
+ out_indices[i] = heap[i].getIdentifier();
204
+ if (out_distances) out_distances[i] = heap[i].getDistance();
205
+ }
206
+ } else {
207
+ for (size_t i = top_n; i > 0; --i) {
208
+ const auto next = heap.top();
209
+ out_indices[i - 1] = next.getIdentifier();
210
+ if (out_distances) out_distances[i - 1] = next.getDistance();
211
+ heap.pop();
212
+ }
213
+ }
214
+ for (size_t i = top_n; i < k; ++i) {
215
+ out_indices[i] = std::numeric_limits<uint32_t>::max();
216
+ if (out_distances) out_distances[i] = std::numeric_limits<float>::max();
217
+ }
218
+ return static_cast<uint32_t>(top_n);
219
+ }
220
+ };
221
+
222
+ // ============================================================================
223
+ // Abstract Searcher Interface
224
+ // ============================================================================
225
+
226
+ class SearcherBase {
227
+ public:
228
+ virtual ~SearcherBase() = default;
229
+
230
+ virtual void set_query_arguments(float search_eps, float rerank_factor = 1.0f) = 0;
231
+ virtual void set_search_eps(float search_eps) = 0;
232
+ virtual float get_search_eps() const = 0;
233
+ virtual void set_rerank_factor(float rerank_factor) = 0;
234
+ virtual float get_rerank_factor() const = 0;
235
+
236
+ // --- Raw buffer API (Zero Overhead) ---
237
+ virtual uint32_t search_f32(const float* query, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, bool unsorted = false) const = 0;
238
+ virtual uint32_t search_f16(const uint16_t* query, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, bool unsorted = false) const = 0;
239
+
240
+ virtual void search_batch_f32(const float* queries, size_t n_queries, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, size_t threads = 1, bool unsorted = false) const = 0;
241
+ virtual void search_batch_f16(const uint16_t* queries, size_t n_queries, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, size_t threads = 1, bool unsorted = false) const = 0;
242
+
243
+ // --- Modern C++20 std::span and std::vector Convenience API ---
244
+ template <typename T>
245
+ uint32_t search(
246
+ std::span<const T> query,
247
+ uint32_t k,
248
+ std::span<uint32_t> out_indices,
249
+ std::span<float> out_distances = {},
250
+ bool unsorted = false
251
+ ) const {
252
+ if (out_indices.size() < k) {
253
+ throw std::invalid_argument("Searcher::search: out_indices span is smaller than k");
254
+ }
255
+ if (!out_distances.empty() && out_distances.size() < k) {
256
+ throw std::invalid_argument("Searcher::search: out_distances span is smaller than k");
257
+ }
258
+ float* d_ptr = out_distances.empty() ? nullptr : out_distances.data();
259
+ if constexpr (std::is_same_v<T, float>) {
260
+ return search_f32(query.data(), k, out_indices.data(), d_ptr, unsorted);
261
+ } else if constexpr (std::is_same_v<T, uint16_t>) {
262
+ return search_f16(query.data(), k, out_indices.data(), d_ptr, unsorted);
263
+ } else {
264
+ static_assert(sizeof(T) == 0, "Unsupported query type for search: must be float or uint16_t (fp16)");
265
+ }
266
+ }
267
+
268
+ template <typename T>
269
+ SearchResult search(std::span<const T> query, uint32_t k, bool return_distances = false, bool unsorted = false) const {
270
+ SearchResult res;
271
+ res.indices.resize(k);
272
+ if (return_distances) res.distances.resize(k);
273
+ uint32_t count = search<T>(
274
+ query, k, std::span<uint32_t>(res.indices),
275
+ return_distances ? std::span<float>(res.distances) : std::span<float>{},
276
+ unsorted
277
+ );
278
+ res.indices.resize(count);
279
+ if (return_distances) res.distances.resize(count);
280
+ return res;
281
+ }
282
+
283
+ template <typename T>
284
+ std::vector<SearchResult> search_batch(
285
+ std::span<const T> queries,
286
+ size_t n_queries,
287
+ uint32_t k,
288
+ size_t threads = 1,
289
+ bool return_distances = false,
290
+ bool unsorted = false
291
+ ) const {
292
+ const size_t dim = queries.size() / n_queries;
293
+ std::vector<SearchResult> results(n_queries);
294
+
295
+ std::vector<uint32_t> all_indices(n_queries * k);
296
+ std::vector<float> all_distances(return_distances ? (n_queries * k) : 0);
297
+
298
+ if constexpr (std::is_same_v<T, float>) {
299
+ search_batch_f32(queries.data(), n_queries, k, all_indices.data(), return_distances ? all_distances.data() : nullptr, threads, unsorted);
300
+ } else if constexpr (std::is_same_v<T, uint16_t>) {
301
+ search_batch_f16(queries.data(), n_queries, k, all_indices.data(), return_distances ? all_distances.data() : nullptr, threads, unsorted);
302
+ }
303
+
304
+ for (size_t q = 0; q < n_queries; ++q) {
305
+ results[q].indices.assign(all_indices.begin() + q * k, all_indices.begin() + (q + 1) * k);
306
+ if (return_distances) {
307
+ results[q].distances.assign(all_distances.begin() + q * k, all_distances.begin() + (q + 1) * k);
308
+ }
309
+ }
310
+ return results;
311
+ }
312
+ };
313
+
314
+ // ============================================================================
315
+ // Pure Templated Searcher Implementation (Zero-Dispatch / Zero-Branch Loop)
316
+ // ============================================================================
317
+
318
+ template <typename QuantT, typename RefinerT>
319
+ class SearcherImpl : public SearcherBase {
320
+ private:
321
+ const deglib::graph::InternalGraph* graph_ = nullptr;
322
+ QuantT quantizer_;
323
+ RefinerT refiner_;
324
+
325
+ float search_eps_ = 0.1f;
326
+ float rerank_factor_ = 1.0f;
327
+
328
+ public:
329
+ SearcherImpl(
330
+ const deglib::graph::InternalGraph& graph,
331
+ QuantT quantizer,
332
+ RefinerT refiner,
333
+ float search_eps = 0.1f,
334
+ float rerank_factor = 1.0f
335
+ )
336
+ : graph_(&graph),
337
+ quantizer_(std::move(quantizer)),
338
+ refiner_(std::move(refiner)),
339
+ search_eps_(search_eps),
340
+ rerank_factor_(rerank_factor) {}
341
+
342
+ void set_query_arguments(float search_eps, float rerank_factor = 1.0f) override {
343
+ search_eps_ = search_eps;
344
+ rerank_factor_ = rerank_factor;
345
+ }
346
+
347
+ void set_search_eps(float search_eps) override { search_eps_ = search_eps; }
348
+ float get_search_eps() const override { return search_eps_; }
349
+
350
+ void set_rerank_factor(float rerank_factor) override { rerank_factor_ = rerank_factor; }
351
+ float get_rerank_factor() const override { return rerank_factor_; }
352
+
353
+ template <typename QueryT>
354
+ inline uint32_t search_single_typed(
355
+ const QueryT* query, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, bool unsorted = false
356
+ ) const {
357
+ const uint32_t dim = graph_->getFeatureSpace().dim();
358
+ const uint32_t fetch_k = std::max(k, static_cast<uint32_t>(std::round(k * rerank_factor_)));
359
+ const size_t graph_feature_bytes = graph_->getFeatureSpace().get_data_size();
360
+
361
+ // 1. Static Query Transformation
362
+ std::vector<std::byte> query_bytes(graph_feature_bytes);
363
+ quantizer_.transform(query, dim, query_bytes.data());
364
+
365
+ // 2. Direct Graph Search
366
+ auto result = graph_->search(
367
+ std::span<const std::byte>(query_bytes.data(), graph_feature_bytes),
368
+ fetch_k, search_eps_, nullptr, 0
369
+ );
370
+ const size_t found_count = result.size();
371
+
372
+ // 3. Static Compile-Time Reranker Check
373
+ if constexpr (RefinerT::enabled) {
374
+ if (fetch_k > k) {
375
+ std::vector<uint32_t> candidate_indices(found_count);
376
+ for (size_t i = 0; i < found_count; ++i) {
377
+ candidate_indices[i] = graph_->getExternalLabel(result[i].getIdentifier());
378
+ }
379
+
380
+ uint32_t ref_count = refiner_.rerank(query, dim, candidate_indices.data(), found_count, k, out_indices, out_distances, unsorted);
381
+ if (ref_count > 0) {
382
+ return ref_count;
383
+ }
384
+ }
385
+ }
386
+
387
+ // Direct return without rerank
388
+ while (result.size() > k) {
389
+ result.pop();
390
+ }
391
+ const size_t top_n = result.size();
392
+ if (unsorted) {
393
+ for (size_t i = 0; i < top_n; ++i) {
394
+ out_indices[i] = graph_->getExternalLabel(result[i].getIdentifier());
395
+ if (out_distances) out_distances[i] = result[i].getDistance();
396
+ }
397
+ } else {
398
+ for (size_t i = top_n; i > 0; --i) {
399
+ const auto next = result.top();
400
+ out_indices[i - 1] = graph_->getExternalLabel(next.getIdentifier());
401
+ if (out_distances) out_distances[i - 1] = next.getDistance();
402
+ result.pop();
403
+ }
404
+ }
405
+ for (size_t i = top_n; i < k; ++i) {
406
+ out_indices[i] = std::numeric_limits<uint32_t>::max();
407
+ if (out_distances) out_distances[i] = std::numeric_limits<float>::max();
408
+ }
409
+ return static_cast<uint32_t>(top_n);
410
+ }
411
+
412
+ uint32_t search_f32(const float* query, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, bool unsorted = false) const override {
413
+ return search_single_typed<float>(query, k, out_indices, out_distances, unsorted);
414
+ }
415
+
416
+ uint32_t search_f16(const uint16_t* query, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, bool unsorted = false) const override {
417
+ return search_single_typed<uint16_t>(query, k, out_indices, out_distances, unsorted);
418
+ }
419
+
420
+ void search_batch_f32(const float* queries, size_t n_queries, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, size_t threads = 1, bool unsorted = false) const override {
421
+ const uint32_t dim = graph_->getFeatureSpace().dim();
422
+ deglib::concurrent::parallel_for(0, n_queries, threads, [&](size_t q, size_t) {
423
+ float* d_ptr = out_distances ? (out_distances + q * k) : nullptr;
424
+ search_f32(queries + q * dim, k, out_indices + q * k, d_ptr, unsorted);
425
+ });
426
+ }
427
+
428
+ void search_batch_f16(const uint16_t* queries, size_t n_queries, uint32_t k, uint32_t* out_indices, float* out_distances = nullptr, size_t threads = 1, bool unsorted = false) const override {
429
+ const uint32_t dim = graph_->getFeatureSpace().dim();
430
+ deglib::concurrent::parallel_for(0, n_queries, threads, [&](size_t q, size_t) {
431
+ float* d_ptr = out_distances ? (out_distances + q * k) : nullptr;
432
+ search_f16(queries + q * dim, k, out_indices + q * k, d_ptr, unsorted);
433
+ });
434
+ }
435
+ };
436
+
437
+ // ============================================================================
438
+ // Modern C++20 Factory Functions
439
+ // ============================================================================
440
+
441
+ template <typename QuantT = NoQuantizer, typename RefinerT = NoRefiner>
442
+ inline std::unique_ptr<SearcherBase> make_searcher(
443
+ const deglib::graph::InternalGraph& graph,
444
+ QuantT quantizer = NoQuantizer{},
445
+ RefinerT refiner = NoRefiner{},
446
+ float search_eps = 0.1f,
447
+ float rerank_factor = 1.0f
448
+ ) {
449
+ return std::make_unique<SearcherImpl<QuantT, RefinerT>>(
450
+ graph, std::move(quantizer), std::move(refiner), search_eps, rerank_factor
451
+ );
452
+ }
453
+
454
+ } // namespace deglib::search
@@ -55,6 +55,7 @@ add_deglib_test(test_fp16_l2 src/unit/distance/test_fp16_l2.cpp)
55
55
  add_deglib_test(test_fp16_inner_product src/unit/distance/test_fp16_inner_product.cpp)
56
56
  add_deglib_test(test_internal_graph src/unit/graph/test_internal_graph.cpp)
57
57
  add_deglib_test(test_search src/unit/test_search.cpp)
58
+ add_deglib_test(test_searcher src/unit/test_searcher.cpp)
58
59
  add_deglib_test(test_evp_quantize src/unit/optimization/quantization/test_evp_quantize.cpp)
59
60
  add_deglib_test(test_scalar_quantize src/unit/optimization/quantization/test_scalar_quantize.cpp)
60
61
  add_deglib_test(test_concurrent src/unit/test_concurrent.cpp)
@@ -737,6 +737,154 @@ py::object search_rerank(
737
737
  return result_indices;
738
738
  }
739
739
 
740
+ // ============================================================================
741
+ // Fast Searcher Wrapper
742
+ // ============================================================================
743
+
744
+ class SearcherPy {
745
+ private:
746
+ std::unique_ptr<deglib::search::SearcherBase> searcher_;
747
+ py::array base_vectors_holder_;
748
+
749
+ public:
750
+ SearcherPy(
751
+ const deglib::DynamicExplorationGraph& graph,
752
+ py::object quantizer = py::none(),
753
+ std::optional<deglib::distances::FloatSpace> rerank_space = std::nullopt,
754
+ std::optional<py::array> base_vectors = std::nullopt,
755
+ float search_eps = 0.1f,
756
+ float rerank_factor = 1.0f
757
+ ) {
758
+ const void* base_vectors_ptr = nullptr;
759
+ size_t num_base_vectors = 0;
760
+ bool is_fp16_base = false;
761
+
762
+ if (base_vectors.has_value() && !base_vectors->is_none()) {
763
+ base_vectors_holder_ = *base_vectors;
764
+ auto buf = base_vectors_holder_.request();
765
+ if (buf.ndim != 2) {
766
+ throw std::invalid_argument("base_vectors must be a 2D array");
767
+ }
768
+ base_vectors_ptr = buf.ptr;
769
+ num_base_vectors = buf.shape[0];
770
+ if (buf.itemsize == 2 && (buf.format == "H" || buf.format == "h" || buf.format == "e")) {
771
+ is_fp16_base = true;
772
+ }
773
+ }
774
+
775
+ auto make_searcher_for_quant = [&]<typename QuantT>(QuantT q) {
776
+ if (rerank_space.has_value() && base_vectors_ptr != nullptr) {
777
+ if (is_fp16_base) {
778
+ using RefinerT = deglib::search::ExactRefiner<uint16_t>;
779
+ searcher_ = std::make_unique<deglib::search::SearcherImpl<QuantT, RefinerT>>(
780
+ graph.internal(), std::move(q), RefinerT(*rerank_space, static_cast<const uint16_t*>(base_vectors_ptr), num_base_vectors), search_eps, rerank_factor
781
+ );
782
+ } else {
783
+ using RefinerT = deglib::search::ExactRefiner<float>;
784
+ searcher_ = std::make_unique<deglib::search::SearcherImpl<QuantT, RefinerT>>(
785
+ graph.internal(), std::move(q), RefinerT(*rerank_space, static_cast<const float*>(base_vectors_ptr), num_base_vectors), search_eps, rerank_factor
786
+ );
787
+ }
788
+ } else {
789
+ using RefinerT = deglib::search::NoRefiner;
790
+ searcher_ = std::make_unique<deglib::search::SearcherImpl<QuantT, RefinerT>>(
791
+ graph.internal(), std::move(q), RefinerT{}, search_eps, rerank_factor
792
+ );
793
+ }
794
+ };
795
+
796
+ if (quantizer.is_none()) {
797
+ make_searcher_for_quant(deglib::search::NoQuantizer{});
798
+ } else if (py::isinstance<deglib::quantization::scalar::ScalarQuantizerInt8>(quantizer)) {
799
+ auto q = py::cast<deglib::quantization::scalar::ScalarQuantizerInt8>(quantizer);
800
+ make_searcher_for_quant(deglib::search::ScalarInt8Quantizer(q));
801
+ } else if (py::isinstance<deglib::quantization::scalar::ScalarQuantizerInt8PerDim>(quantizer)) {
802
+ auto q = py::cast<deglib::quantization::scalar::ScalarQuantizerInt8PerDim>(quantizer);
803
+ make_searcher_for_quant(deglib::search::ScalarInt8PerDimQuantizer(q));
804
+ } else if (py::isinstance<deglib::quantization::scalar::ScalarQuantizerUint8>(quantizer)) {
805
+ auto q = py::cast<deglib::quantization::scalar::ScalarQuantizerUint8>(quantizer);
806
+ make_searcher_for_quant(deglib::search::ScalarUint8Quantizer(q));
807
+ } else if (py::isinstance<deglib::quantization::scalar::ScalarQuantizerUint8PerDim>(quantizer)) {
808
+ auto q = py::cast<deglib::quantization::scalar::ScalarQuantizerUint8PerDim>(quantizer);
809
+ make_searcher_for_quant(deglib::search::ScalarUint8PerDimQuantizer(q));
810
+ } else if (py::isinstance<py::int_>(quantizer)) {
811
+ uint32_t nz = py::cast<uint32_t>(quantizer);
812
+ make_searcher_for_quant(deglib::search::EVPQuantizer(nz));
813
+ } else {
814
+ throw std::invalid_argument("Unsupported quantizer type provided to create_searcher.");
815
+ }
816
+ }
817
+
818
+ void set_query_arguments(float search_eps, float rerank_factor = 1.0f) {
819
+ searcher_->set_query_arguments(search_eps, rerank_factor);
820
+ }
821
+
822
+ void set_search_eps(float search_eps) { searcher_->set_search_eps(search_eps); }
823
+ float get_search_eps() const { return searcher_->get_search_eps(); }
824
+
825
+ void set_rerank_factor(float rerank_factor) { searcher_->set_rerank_factor(rerank_factor); }
826
+ float get_rerank_factor() const { return searcher_->get_rerank_factor(); }
827
+
828
+ py::object search(py::array query, uint32_t k, bool return_distances = false, bool unsorted = false) {
829
+ auto buf = query.request();
830
+ if (buf.ndim != 1 && (buf.ndim != 2 || buf.shape[0] != 1)) {
831
+ throw std::invalid_argument("search query must be 1D vector (or 1xDim 2D array)");
832
+ }
833
+ py::array_t<uint32_t> result(static_cast<py::ssize_t>(k));
834
+ uint32_t* out_ptr = static_cast<uint32_t*>(result.request().ptr);
835
+
836
+ py::array_t<float> dist_result;
837
+ float* dist_ptr = nullptr;
838
+ if (return_distances) {
839
+ dist_result = py::array_t<float>(static_cast<py::ssize_t>(k));
840
+ dist_ptr = static_cast<float*>(dist_result.request().ptr);
841
+ }
842
+
843
+ uint32_t count = 0;
844
+ if (buf.itemsize == 2 && (buf.format == "H" || buf.format == "h" || buf.format == "e")) {
845
+ count = searcher_->search_f16(static_cast<const uint16_t*>(buf.ptr), k, out_ptr, dist_ptr, unsorted);
846
+ } else {
847
+ count = searcher_->search_f32(static_cast<const float*>(buf.ptr), k, out_ptr, dist_ptr, unsorted);
848
+ }
849
+
850
+ if (return_distances) {
851
+ return py::make_tuple(result, dist_result);
852
+ }
853
+ return result;
854
+ }
855
+
856
+ py::object search_batch(py::array queries, uint32_t k, size_t num_threads = 1, bool return_distances = false, bool unsorted = false) {
857
+ auto buf = queries.request();
858
+ if (buf.ndim != 2) {
859
+ throw std::invalid_argument("search_batch queries must be 2D array");
860
+ }
861
+ size_t n_queries = buf.shape[0];
862
+ auto result = py::array_t<uint32_t>({n_queries, size_t(k)});
863
+ uint32_t* out_ptr = static_cast<uint32_t*>(result.request().ptr);
864
+
865
+ py::array_t<float> dist_result;
866
+ float* dist_ptr = nullptr;
867
+ if (return_distances) {
868
+ dist_result = py::array_t<float>({n_queries, size_t(k)});
869
+ dist_ptr = static_cast<float*>(dist_result.request().ptr);
870
+ }
871
+
872
+ {
873
+ py::gil_scoped_release release;
874
+ if (buf.itemsize == 2 && (buf.format == "H" || buf.format == "h" || buf.format == "e")) {
875
+ searcher_->search_batch_f16(static_cast<const uint16_t*>(buf.ptr), n_queries, k, out_ptr, dist_ptr, num_threads, unsorted);
876
+ } else {
877
+ searcher_->search_batch_f32(static_cast<const float*>(buf.ptr), n_queries, k, out_ptr, dist_ptr, num_threads, unsorted);
878
+ }
879
+ }
880
+
881
+ if (return_distances) {
882
+ return py::make_tuple(result, dist_result);
883
+ }
884
+ return result;
885
+ }
886
+ };
887
+
740
888
  // ============================================================================
741
889
  // EVP Quantization Bindings
742
890
  // ============================================================================
@@ -1557,6 +1705,38 @@ PYBIND11_MODULE(deglib_cpp, m) {
1557
1705
  py::arg("k_top") = 0, py::arg("num_threads") = 0, py::arg("return_distances") = false, py::arg("unsorted") = false
1558
1706
  );
1559
1707
 
1708
+ py::class_<SearcherPy>(search_module, "Searcher")
1709
+ .def(
1710
+ py::init<
1711
+ const deglib::DynamicExplorationGraph&,
1712
+ py::object,
1713
+ std::optional<deglib::distances::FloatSpace>,
1714
+ std::optional<py::array>,
1715
+ float,
1716
+ float
1717
+ >(),
1718
+ py::arg("graph"),
1719
+ py::arg("quantizer") = py::none(),
1720
+ py::arg("rerank_space") = std::nullopt,
1721
+ py::arg("base_vectors") = std::nullopt,
1722
+ py::arg("search_eps") = 0.1f,
1723
+ py::arg("rerank_factor") = 1.0f,
1724
+ py::keep_alive<1, 2>()
1725
+ )
1726
+ .def("set_query_arguments", &SearcherPy::set_query_arguments, py::arg("search_eps"), py::arg("rerank_factor") = 1.0f)
1727
+ .def("set_search_eps", &SearcherPy::set_search_eps, py::arg("search_eps"))
1728
+ .def("get_search_eps", &SearcherPy::get_search_eps)
1729
+ .def("set_rerank_factor", &SearcherPy::set_rerank_factor, py::arg("rerank_factor"))
1730
+ .def("get_rerank_factor", &SearcherPy::get_rerank_factor)
1731
+ .def(
1732
+ "search", &SearcherPy::search, py::arg("query"), py::arg("k"), py::arg("return_distances") = false,
1733
+ py::arg("unsorted") = false
1734
+ )
1735
+ .def(
1736
+ "search_batch", &SearcherPy::search_batch, py::arg("queries"), py::arg("k"), py::arg("num_threads") = 1,
1737
+ py::arg("return_distances") = false, py::arg("unsorted") = false
1738
+ );
1739
+
1560
1740
  // graphs
1561
1741
  py::class_<deglib::DynamicExplorationGraph>(m, "DynamicExplorationGraph")
1562
1742
  .def(py::init<deglib::graph::InternalGraph&>())