Skip to content

Commit ba91a70

Browse files
NikolaosPapailiouNikolaosPapailiou
authored
Merge pull request #144 from TileDB-Inc/npapa/fix-empty-results
Fix empty results
2 parents 3e0d06b + 58111d8 commit ba91a70

File tree

22 files changed

+897
-534
lines changed

22 files changed

+897
-534
lines changed

‎apis/python/src/tiledb/vector_search/module.cc‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@
1111
namespace py = pybind11;
1212
using Ctx = tiledb::Context;
1313

14-
bool global_debug = false;
15-
1614
bool enable_stats = false;
1715
std::vector<json> core_stats;
1816

@@ -126,7 +124,8 @@ static void declare_qv_query_heap_infinite_ram(py::module& m, const std::string&
126124
size_t k_nn,
127125
size_t nthreads) -> py::tuple { //std::pair<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
128126

129-
auto r = detail::ivf::qv_query_heap_infinite_ram(
127+
// auto r = detail::ivf::qv_query_heap_infinite_ram(
128+
auto r = detail::ivf::query_infinite_ram(
130129
parts,
131130
centroids,
132131
query_vectors,
@@ -178,7 +177,7 @@ static void declare_nuv_query_heap_infinite_ram(py::module& m, const std::string
178177
std::vector<Id_Type>& ids,
179178
size_t nprobe,
180179
size_t k_nn,
181-
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
180+
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> { // TODO change return type
182181

183182
auto r = detail::ivf::nuv_query_heap_infinite_ram_reg_blocked(
184183
parts,
@@ -205,7 +204,7 @@ static void declare_nuv_query_heap_finite_ram(py::module& m, const std::string&
205204
size_t nprobe,
206205
size_t k_nn,
207206
size_t upper_bound,
208-
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> { // TODO change return type
207+
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> { // TODO change return type
209208

210209
auto r = detail::ivf::nuv_query_heap_finite_ram_reg_blocked<T, Id_Type>(
211210
ctx,
@@ -280,7 +279,7 @@ static void declare_ivf_index_tdb(py::module& m, const std::string& suffix) {
280279
}, py::keep_alive<1,2>());
281280
}
282281

283-
template <class T=float, class U=size_t>
282+
template <class T=float, class U=uint64_t>
284283
static void declareFixedMinPairHeap(py::module& mod) {
285284
using PyFixedMinPairHeap = py::class_<fixed_min_pair_heap<T, U>>;
286285
PyFixedMinPairHeap cls(mod, "FixedMinPairHeap", py::buffer_protocol());
@@ -357,7 +356,7 @@ void declareStdVector(py::module& m, const std::string& suffix) {
357356
});
358357
}
359358

360-
template <typename T, typename indices_type = size_t>
359+
template <typename T, typename indices_type = uint64_t>
361360
void declarePartitionIvfIndex(py::module& m, const std::string& suffix) {
362361
m.def(("partition_ivf_index_" + suffix).c_str(),
363362
[](ColMajorMatrix<float>& centroids,
@@ -401,8 +400,8 @@ static void declare_vq_query_heap(py::module& m, const std::string& suffix) {
401400
ColMajorMatrix<float>& query_vectors,
402401
const std::vector<uint64_t> &ids,
403402
int k,
404-
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
405-
auto r = detail::flat::vq_query_heap<tdbColMajorMatrix<T>, ColMajorMatrix<float>, uint64_t>(data, query_vectors, ids, k, nthreads);
403+
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
404+
auto r = detail::flat::vq_query_heap(data, query_vectors, ids, k, nthreads);
406405
return r;
407406
});
408407
}
@@ -414,8 +413,8 @@ static void declare_vq_query_heap_pyarray(py::module& m, const std::string& suff
414413
ColMajorMatrix<float>& query_vectors,
415414
const std::vector<uint64_t> &ids,
416415
int k,
417-
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
418-
auto r = detail::flat::vq_query_heap<ColMajorMatrix<T>, ColMajorMatrix<float>, uint64_t>(data, query_vectors, ids, k, nthreads);
416+
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
417+
auto r = detail::flat::vq_query_heap(data, query_vectors, ids, k, nthreads);
419418
return r;
420419
});
421420
}
@@ -497,7 +496,7 @@ PYBIND11_MODULE(_tiledbvspy, m) {
497496
[](ColMajorMatrix<float>& data,
498497
ColMajorMatrix<float>& query_vectors,
499498
int k,
500-
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
499+
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
501500
auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
502501
return r;
503502
});
@@ -506,13 +505,13 @@ PYBIND11_MODULE(_tiledbvspy, m) {
506505
[](tdbColMajorMatrix<uint8_t>& data,
507506
ColMajorMatrix<float>& query_vectors,
508507
int k,
509-
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<size_t>> {
508+
size_t nthreads) -> std::tuple<ColMajorMatrix<float>, ColMajorMatrix<uint64_t>> {
510509
auto r = detail::flat::vq_query_heap(data, query_vectors, k, nthreads);
511510
return r;
512511
});
513512

514513
m.def("validate_top_k_u64",
515-
[](const ColMajorMatrix<size_t>& top_k,
514+
[](const ColMajorMatrix<uint64_t>& top_k,
516515
const ColMajorMatrix<int32_t>& ground_truth) -> bool {
517516
return validate_top_k(top_k, ground_truth);
518517
});
@@ -535,9 +534,11 @@ PYBIND11_MODULE(_tiledbvspy, m) {
535534
return json{core_stats}.dump();
536535
});
537536

537+
#if 0
538538
m.def("set_debug", [](bool debug) {
539539
global_debug = debug;
540540
});
541+
#endif
541542

542543
declare_vq_query_heap<uint8_t>(m, "u8");
543544
declare_vq_query_heap<float>(m, "f32");

‎src/include/detail/flat/qv.h‎

Lines changed: 49 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -73,10 +73,13 @@ namespace detail::flat {
7373
*/
7474
template <class DB, class Q>
7575
[[deprecated]] auto qv_query_heap_0(
76-
DB& db, const Q& q, int k_nn, unsigned int nthreads) {
76+
const DB& db, const Q& q, int k_nn, unsigned int nthreads) {
7777
scoped_timer _{tdb_func__};
7878

79-
ColMajorMatrix<size_t> top_k(k_nn, size(q));
79+
using id_type = size_t;
80+
using score_type = float;
81+
82+
ColMajorMatrix<id_type> top_k(k_nn, size(q));
8083

8184
auto par = stdx::execution::indexed_parallel_policy{nthreads};
8285
stdx::range_for_each(
@@ -85,7 +88,7 @@ template <class DB, class Q>
8588
size_t size_db = size(db);
8689

8790
// @todo can we do this more efficiently?
88-
Vector<float> scores(size_db);
91+
Vector<score_type> scores(size_db);
8992

9093
for (size_t i = 0; i < size_db; ++i) {
9194
scores[i] = L2(q_vec, db[i]);
@@ -113,40 +116,44 @@ template <class DB, class Q>
113116
* @return A matrix of size k x #queries containing the top k results for each
114117
* query.
115118
*/
116-
template <class T, class DB, class Q, class Index>
119+
template <class T, class DB, class Q, class ID>
117120
auto qv_query_heap(
118121
T,
119-
DB& db,
120-
Q& q,
121-
const std::vector<Index>& ids,
122+
const DB& db,
123+
const Q& q,
124+
const ID& ids,
122125
int k_nn,
123126
unsigned nthreads);
124127

125128
template <class DB, class Q>
126-
auto qv_query_heap(DB& db, Q& q, int k_nn, unsigned nthreads) {
129+
auto qv_query_heap(const DB& db, const Q& q, int k_nn, unsigned nthreads) {
127130
return qv_query_heap(
128-
without_ids{}, db, q, std::vector<size_t>{}, k_nn, nthreads);
131+
without_ids{}, db, q, std::vector<uint64_t>{}, k_nn, nthreads);
129132
}
130133

131-
template <class DB, class Q, class Index>
134+
template <class DB, class Q, class ID>
132135
auto qv_query_heap(
133-
DB& db, Q& q, const std::vector<Index>& ids, int k_nn, unsigned nthreads) {
136+
const DB& db, const Q& q, const ID& ids, int k_nn, unsigned nthreads) {
134137
return qv_query_heap(with_ids{}, db, q, ids, k_nn, nthreads);
135138
}
136139

137140
// @todo Add to out of core
138-
template <class T, class DB, class Q, class Index>
141+
template <class T, class DB, class Q, class ID>
139142
auto qv_query_heap(
140143
T,
141-
DB& db,
142-
Q& query,
143-
const std::vector<Index>& ids,
144+
const DB& db,
145+
const Q& query,
146+
const ID& ids,
144147
int k_nn,
145148
unsigned nthreads) {
146149
scoped_timer _{tdb_func__};
147150

148-
auto top_k = ColMajorMatrix<size_t>(k_nn, query.num_cols());
149-
auto top_k_scores = ColMajorMatrix<float>(k_nn, query.num_cols());
151+
// using feature_type = typename std::remove_reference_t<decltype(db)>::value_type;
152+
using id_type = typename std::remove_reference_t<decltype(ids)>::value_type;
153+
using score_type = float;
154+
155+
auto top_k = ColMajorMatrix<id_type>(k_nn, query.num_cols());
156+
auto top_k_scores = ColMajorMatrix<score_type>(k_nn, query.num_cols());
150157

151158
// Have to do explicit asynchronous threading here, as the current parallel
152159
// algorithms have iterator-based interaces, and the `Matrix` class does not
@@ -159,7 +166,7 @@ auto qv_query_heap(
159166
std::move(par),
160167
query,
161168
[&, size_db](auto&& q_vec, auto&& n = 0, auto&& j = 0) {
162-
fixed_min_pair_heap<float, size_t> min_scores(k_nn);
169+
fixed_min_pair_heap<score_type, id_type> min_scores(k_nn);
163170

164171
for (size_t i = 0; i < size_db; ++i) {
165172
auto score = L2(q_vec, db[i]);
@@ -190,35 +197,40 @@ auto qv_query_heap(
190197
* @return A matrix of size k x #queries containing the top k results for each
191198
* query.
192199
*/
193-
template <class T, class DB, class Q, class Index>
200+
template <class T, class DB, class Q, class ID>
194201
auto qv_query_heap_tiled(
195202
T,
196203
DB& db,
197-
Q& q,
198-
const std::vector<Index>& ids,
204+
const Q& q,
205+
const ID& ids,
199206
int k_nn,
200207
unsigned nthreads);
201208

202209
template <class DB, class Q>
203-
auto qv_query_heap_tiled(DB& db, Q& q, int k_nn, unsigned nthreads) {
210+
auto qv_query_heap_tiled(DB& db, const Q& q, int k_nn, unsigned nthreads) {
204211
return qv_query_heap_tiled(
205-
without_ids{}, db, q, std::vector<size_t>{}, k_nn, nthreads);
212+
without_ids{}, db, q, std::vector<uint64_t>{}, k_nn, nthreads);
206213
}
207214

208-
template <class DB, class Q, class Index>
215+
template <class DB, class Q, class ID>
209216
auto qv_query_heap_tiled(
210-
DB& db, Q& q, const std::vector<Index>& ids, int k_nn, unsigned nthreads) {
217+
DB& db, Q& q, const ID& ids, int k_nn, unsigned nthreads) {
211218
return qv_query_heap_tiled(with_ids{}, db, q, ids, k_nn, nthreads);
212219
}
213220

214-
template <class T, class DB, class Q, class Index>
221+
template <class T, class DB, class Q, class ID>
215222
auto qv_query_heap_tiled(
216223
T,
217224
DB& db,
218-
Q& query,
219-
[[maybe_unused]] const std::vector<Index>& ids,
225+
const Q& query,
226+
[[maybe_unused]] const ID& ids,
220227
int k_nn,
221228
unsigned nthreads) {
229+
230+
// using feature_type = typename std::remove_reference_t<decltype(db)>::value_type;
231+
using id_type = typename std::remove_reference_t<decltype(ids)>::value_type;
232+
using score_type = float;
233+
222234
if constexpr (is_loadable_v<decltype(db)>) {
223235
db.load();
224236
}
@@ -236,8 +248,8 @@ auto qv_query_heap_tiled(
236248
std::vector<std::future<void>> futs;
237249
futs.reserve(nthreads);
238250

239-
auto min_scores = std::vector<fixed_min_pair_heap<float, size_t>>(
240-
size(query), fixed_min_pair_heap<float, size_t>(k_nn));
251+
auto min_scores = std::vector<fixed_min_pair_heap<score_type, id_type>>(
252+
size(query), fixed_min_pair_heap<score_type, id_type>(k_nn));
241253

242254
// @todo: Use range::for_each
243255
for (size_t n = 0; n < nthreads; ++n) {
@@ -251,8 +263,8 @@ auto qv_query_heap_tiled(
251263
auto len = 2 * ((stop - start) / 2);
252264
auto end = start + len;
253265

254-
// auto min_scores0 = fixed_min_pair_heap<float, size_t> (k);
255-
// auto min_scores1 = fixed_min_pair_heap<float, size_t> (k);
266+
// auto min_scores0 = fixed_min_pair_heap<score_type, id_type> (k);
267+
// auto min_scores1 = fixed_min_pair_heap<score_type, id_type> (k);
256268

257269
for (auto j = start; j != end; j += 2) {
258270
auto j0 = j + 0;
@@ -365,15 +377,17 @@ template <class DB, class Q>
365377
auto qv_partition(const DB& db, const Q& q, unsigned nthreads) {
366378
scoped_timer _{tdb_func__};
367379

380+
// Just need a single vector -- creating an index, not ids, so hardcoded size_t is okay to use here
381+
using id_type = size_t;
382+
using score_type = float;
368383
auto size_db = size(db);
369384

370-
// Just need a single vector
371-
std::vector<size_t> top_k(q.num_cols());
385+
std::vector<id_type> top_k(q.num_cols());
372386

373387
auto par = stdx::execution::indexed_parallel_policy{(size_t)nthreads};
374388
stdx::range_for_each(
375389
std::move(par), q, [&, size_db](auto&& qvec, auto&& n = 0, auto&& j = 0) {
376-
float min_score = std::numeric_limits<float>::max();
390+
score_type min_score = std::numeric_limits<score_type>::max();
377391
size_t idx = 0;
378392

379393
for (size_t i = 0; i < size_db; ++i) {

0 commit comments

Comments
 (0)