mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Introduce early stop facilities for IVF, attempt 2 (#5160)
Summary: # Introduce early-stop controls for IVF search This PR adds optional early-stop controls to IVF search through `SearchParametersIVF`. IVF search looks in a sequence of buckets. Normally it keeps looking until it has checked all requested buckets. This PR adds a few ways to stop earlier: - stop after enough candidate vectors have been checked, - stop FastScan after enough buckets have been checked, - stop range search after several nearby buckets produce no results. All new fields default to disabled, so existing searches keep their current behavior unless callers opt in. ## New SearchParametersIVF fields Default values: | Field | Default | Effect | |---|---:|---| | `ensure_topk_full` | `false` | disabled; keep the existing hard `max_codes` behavior | | `max_lists_num` | `0` | disabled; FastScan is bounded only by `nprobe` | | `max_empty_result_buckets` | `0` | disabled; range search scans all requested probes | `ensure_topk_full` For k-NN search, `max_codes` limits how many candidate vectors are scanned. With `ensure_topk_full=false`, the search stops when the budget is reached. With `ensure_topk_full=true`, very small budgets become less aggressive: `max_codes` is treated as at least `k` post-filter scans. This helps when `max_codes < k`, or when an `IDSelector` filters out many candidates before a distance is computed. In the generic IVF path this is supported for `parallel_mode` 0 and 3. In FastScan k-NN it is supported by the per-query implementations 10 and 11. `max_lists_num` FastScan k-NN can also stop after visiting a limited number of inverted lists. `0` means disabled. With `ensure_topk_full`, this list budget is treated as at least `k` lists. This is only used by the FastScan per-query path (`search_implem_10`, also used for implementation 11). When FastScan is left on `implem = 0`, setting an early-stop k-NN option makes the dispatcher choose the per-query implementation automatically. If an incompatible FastScan implementation is selected explicitly, Faiss throws instead of ignoring the option. `max_empty_result_buckets` This is for range search. It stops probing a query after that many consecutive probed lists add no in-radius results. The counter resets as soon as a list adds at least one result. `0` means disabled. Setting this trades recall for less work: once nearby buckets repeatedly produce no hits, farther buckets are less likely to help. The generic IVF range path supports this in `parallel_mode == 0`. FastScan range search also supports it through implementation 10. With `implem = 0`, FastScan selects implementation 10 automatically when this option is set; explicit incompatible implementations throw. ## Selector-aware scan accounting This PR adds `InvertedListScannerStats`, returned by inverted-list scanner methods. It currently reports: - `scan_cnt`: number of distances actually computed, - `nheap_updates`: number of result updates. The important behavior change is that `IndexIVFStats::ndis` now uses `scan_cnt`. With an `IDSelector`, candidates rejected by the selector are not counted as computed distances. Without a selector, this matches the previous count because every scanned code reaches distance computation. This same post-filter count is used by `ensure_topk_full`, so a restrictive selector no longer consumes the search budget with candidates that were never eligible for the result heap. ## FastScan behavior FastScan k-NN now applies `max_codes`, `max_lists_num`, and `ensure_topk_full` in the per-query implementation. The code keeps separate counters per query, so multi-query batches have independent budgets. For `max_codes`, FastScan uses the result handler's scanned-row count. This means filtered-out rows do not consume the code budget. FastScan range search rejects k-NN-only options (`max_codes`, `max_lists_num`, and `ensure_topk_full`). It accepts `max_empty_result_buckets` when the selected implementation can preserve per-query probe order. ## API exposure The new fields are exposed through: - C++ `SearchParametersIVF`, - Python bindings, - the C API. The statistics is tracked via `InvertedListScanner::scan_codes(..., ResultHandler& handler)` call. In order to make it work inside `IndexIVF::search()`, it has been decided to transition it into this particular override of `InvertedListScanner::scan_codes()`. As a result, non-`ResultHandler` version of `scan_codes()` has been marked as non-virtual, so the custom `InvertedListScanner` implementations should override `ResultHandler`-based versions only. Another consequence is that `IndexIVFPQFastScanScanner` and `IndexIVFRaBitQFastScanScanner` accept only `HeapResultHandler` types of `ResultHandler`. These are the compromises that had to be made. ## Tests The tests cover: - default parameters preserving baseline k-NN and range-search results, - `ensure_topk_full` filling more result slots under tight budgets, - selector-aware `IndexIVFStats::ndis`, - range-search early exit and counter reset behavior, - unsupported parallel modes and FastScan implementations, - FastScan k-NN and range-search early-stop behavior, - Python and C API field exposure. Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5160 Reviewed By: limqiying Differential Revision: D103044128 Pulled By: mnorris11 fbshipit-source-id: 80f1241c13ffc8173a3da0f2974fa9b6dd318890
This commit is contained in:
Alexandr Guzhva
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
812010fe19
commit
799bf3aaec
@@ -161,6 +161,10 @@ endif()
|
||||
add_executable(example_c EXCLUDE_FROM_ALL example_c.c)
|
||||
target_link_libraries(example_c PRIVATE faiss_c)
|
||||
|
||||
add_executable(example_search_params_ivf_c EXCLUDE_FROM_ALL
|
||||
example_search_params_ivf_c.c)
|
||||
target_link_libraries(example_search_params_ivf_c PRIVATE faiss_c)
|
||||
|
||||
if(FAISS_ENABLE_GPU)
|
||||
add_subdirectory(gpu)
|
||||
endif()
|
||||
|
||||
@@ -54,6 +54,15 @@ DEFINE_SETTER(SearchParametersIVF, size_t, nprobe)
|
||||
DEFINE_GETTER(SearchParametersIVF, size_t, max_codes)
|
||||
DEFINE_SETTER(SearchParametersIVF, size_t, max_codes)
|
||||
|
||||
DEFINE_GETTER(SearchParametersIVF, size_t, max_lists_num)
|
||||
DEFINE_SETTER(SearchParametersIVF, size_t, max_lists_num)
|
||||
|
||||
DEFINE_GETTER(SearchParametersIVF, int, ensure_topk_full)
|
||||
DEFINE_SETTER_STATIC(SearchParametersIVF, bool, int, ensure_topk_full)
|
||||
|
||||
DEFINE_GETTER(SearchParametersIVF, size_t, max_empty_result_buckets)
|
||||
DEFINE_SETTER(SearchParametersIVF, size_t, max_empty_result_buckets)
|
||||
|
||||
/// IndexIVF definitions
|
||||
|
||||
DEFINE_DESTRUCTOR(IndexIVF)
|
||||
|
||||
@@ -33,6 +33,12 @@ int faiss_SearchParametersIVF_new_with(
|
||||
FAISS_DECLARE_GETTER(SearchParametersIVF, const FaissIDSelector*, sel)
|
||||
FAISS_DECLARE_GETTER_SETTER(SearchParametersIVF, size_t, nprobe)
|
||||
FAISS_DECLARE_GETTER_SETTER(SearchParametersIVF, size_t, max_codes)
|
||||
FAISS_DECLARE_GETTER_SETTER(SearchParametersIVF, size_t, max_lists_num)
|
||||
FAISS_DECLARE_GETTER_SETTER(SearchParametersIVF, int, ensure_topk_full)
|
||||
FAISS_DECLARE_GETTER_SETTER(
|
||||
SearchParametersIVF,
|
||||
size_t,
|
||||
max_empty_result_buckets)
|
||||
|
||||
/** Index based on a inverted file (IVF)
|
||||
*
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
/*
|
||||
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
// -*- c -*-
|
||||
|
||||
// Smoke test for the three new fields on FaissSearchParametersIVF:
|
||||
// - max_lists_num
|
||||
// - ensure_topk_full
|
||||
// - max_empty_result_buckets
|
||||
// Verifies that the C API symbols link, that getters/setters round-trip,
|
||||
// and that the default values match the C++ struct defaults.
|
||||
|
||||
#include <assert.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
||||
#include "IndexIVF_c.h"
|
||||
#include "error_c.h"
|
||||
|
||||
#define FAISS_TRY(C) \
|
||||
{ \
|
||||
if (C) { \
|
||||
fprintf(stderr, "%s", faiss_get_last_error()); \
|
||||
exit(-1); \
|
||||
} \
|
||||
}
|
||||
|
||||
int main(void) {
|
||||
FaissSearchParametersIVF* sp = NULL;
|
||||
FAISS_TRY(faiss_SearchParametersIVF_new(&sp));
|
||||
|
||||
// Defaults: nprobe=1, everything else 0/false.
|
||||
assert(faiss_SearchParametersIVF_nprobe(sp) == 1);
|
||||
assert(faiss_SearchParametersIVF_max_codes(sp) == 0);
|
||||
assert(faiss_SearchParametersIVF_max_lists_num(sp) == 0);
|
||||
assert(faiss_SearchParametersIVF_ensure_topk_full(sp) == 0);
|
||||
assert(faiss_SearchParametersIVF_max_empty_result_buckets(sp) == 0);
|
||||
|
||||
// Roundtrip setters.
|
||||
faiss_SearchParametersIVF_set_max_lists_num(sp, 4);
|
||||
assert(faiss_SearchParametersIVF_max_lists_num(sp) == 4);
|
||||
|
||||
faiss_SearchParametersIVF_set_ensure_topk_full(sp, 1);
|
||||
assert(faiss_SearchParametersIVF_ensure_topk_full(sp) != 0);
|
||||
|
||||
faiss_SearchParametersIVF_set_ensure_topk_full(sp, 0);
|
||||
assert(faiss_SearchParametersIVF_ensure_topk_full(sp) == 0);
|
||||
|
||||
faiss_SearchParametersIVF_set_max_empty_result_buckets(sp, 3);
|
||||
assert(faiss_SearchParametersIVF_max_empty_result_buckets(sp) == 3);
|
||||
|
||||
faiss_SearchParametersIVF_free(sp);
|
||||
|
||||
printf("example_search_params_ivf_c: OK\n");
|
||||
return 0;
|
||||
}
|
||||
+74
-11
@@ -420,6 +420,9 @@ void IndexIVF::search_preassigned(
|
||||
|
||||
const idx_t unlimited_list_size = std::numeric_limits<idx_t>::max();
|
||||
idx_t cur_max_codes = params ? params->max_codes : this->max_codes;
|
||||
const bool ensure_topk_full = params ? params->ensure_topk_full : false;
|
||||
idx_t effective_max_codes = cur_max_codes;
|
||||
|
||||
IDSelector* sel = params ? params->sel : nullptr;
|
||||
const IDSelectorRange* selr = dynamic_cast<const IDSelectorRange*>(sel);
|
||||
if (selr) {
|
||||
@@ -454,9 +457,17 @@ void IndexIVF::search_preassigned(
|
||||
cur_max_codes == 0 || pmode == 0 || pmode == 3,
|
||||
"max_codes supported only for parallel_mode = 0 or 3");
|
||||
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
!ensure_topk_full || pmode == 0 || pmode == 3,
|
||||
"ensure_topk_full supported only for parallel_mode = 0 or 3");
|
||||
|
||||
if (cur_max_codes == 0) {
|
||||
cur_max_codes = unlimited_list_size;
|
||||
}
|
||||
// Budget used by the probe loop below. ensure_topk_full makes a small
|
||||
// max_codes budget large enough to give k post-filter candidates a chance.
|
||||
effective_max_codes =
|
||||
ensure_topk_full ? std::max(cur_max_codes, k) : cur_max_codes;
|
||||
|
||||
[[maybe_unused]] bool do_parallel = omp_get_max_threads() >= 2 &&
|
||||
(pmode == 0 ? false
|
||||
@@ -549,7 +560,6 @@ void IndexIVF::search_preassigned(
|
||||
nlistv++;
|
||||
if (invlists->use_iterator) {
|
||||
size_t list_size = 0;
|
||||
|
||||
std::unique_ptr<InvertedListsIterator> it(
|
||||
invlists->get_iterator(key, inverted_list_context));
|
||||
|
||||
@@ -588,10 +598,25 @@ void IndexIVF::search_preassigned(
|
||||
ids += jmin;
|
||||
}
|
||||
|
||||
nheap += scanner->scan_codes(
|
||||
list_size, codes, ids, simi, idxi, k);
|
||||
|
||||
return list_size;
|
||||
size_t old_scan_cnt = 0;
|
||||
size_t old_heap_updates = 0;
|
||||
if (metric_type == METRIC_INNER_PRODUCT) {
|
||||
HeapResultHandler<HeapForIP, false> handler(
|
||||
k, simi, idxi);
|
||||
old_scan_cnt = handler.stats.scan_cnt;
|
||||
old_heap_updates = handler.stats.nheap_updates;
|
||||
scanner->scan_codes(list_size, codes, ids, handler);
|
||||
nheap += handler.stats.nheap_updates - old_heap_updates;
|
||||
return handler.stats.scan_cnt - old_scan_cnt;
|
||||
} else {
|
||||
HeapResultHandler<HeapForL2, false> handler(
|
||||
k, simi, idxi);
|
||||
old_scan_cnt = handler.stats.scan_cnt;
|
||||
old_heap_updates = handler.stats.nheap_updates;
|
||||
scanner->scan_codes(list_size, codes, ids, handler);
|
||||
nheap += handler.stats.nheap_updates - old_heap_updates;
|
||||
return handler.stats.scan_cnt - old_scan_cnt;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -617,13 +642,23 @@ void IndexIVF::search_preassigned(
|
||||
|
||||
// loop over probes
|
||||
for (idx_t ik = 0; ik < cur_nprobe; ik++) {
|
||||
// For soft budgets, scan whole lists so
|
||||
// IDSelector-filtered rows do not consume the
|
||||
// remaining code budget.
|
||||
const idx_t list_size_max = ensure_topk_full
|
||||
? unlimited_list_size
|
||||
: effective_max_codes - nscan;
|
||||
nscan += scan_one_list(
|
||||
keys[i * cur_nprobe + ik],
|
||||
coarse_dis[i * cur_nprobe + ik],
|
||||
simi,
|
||||
idxi,
|
||||
cur_max_codes - nscan);
|
||||
if (nscan >= cur_max_codes) {
|
||||
list_size_max);
|
||||
|
||||
// Early-stop check: apply max_codes after each
|
||||
// list. nscan is the number of distances
|
||||
// actually computed.
|
||||
if (nscan >= effective_max_codes) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -792,6 +827,9 @@ void IndexIVF::range_search_preassigned(
|
||||
FAISS_THROW_IF_NOT(cur_nprobe > 0);
|
||||
|
||||
idx_t cur_max_codes = params ? params->max_codes : this->max_codes;
|
||||
// Range-search early-stop budget. 0 disables the empty-bucket stop.
|
||||
const size_t max_empty_result_buckets =
|
||||
params ? params->max_empty_result_buckets : 0;
|
||||
IDSelector* sel = params ? params->sel : nullptr;
|
||||
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
@@ -799,6 +837,10 @@ void IndexIVF::range_search_preassigned(
|
||||
(cur_max_codes == 0 && store_pairs == false),
|
||||
"iterable inverted lists don't support max_codes and store_pairs");
|
||||
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
max_empty_result_buckets == 0 || parallel_mode == 0,
|
||||
"max_empty_result_buckets supported only for parallel_mode = 0");
|
||||
|
||||
size_t nlistv = 0, ndis = 0;
|
||||
|
||||
std::exception_ptr ex;
|
||||
@@ -846,24 +888,26 @@ void IndexIVF::range_search_preassigned(
|
||||
return;
|
||||
}
|
||||
|
||||
size_t list_size = 0;
|
||||
scanner->set_list(key, coarse_dis[i * cur_nprobe + ik]);
|
||||
const size_t scan_cnt0 = qres.stats.scan_cnt;
|
||||
if (invlists->use_iterator) {
|
||||
size_t list_size = 0;
|
||||
std::unique_ptr<InvertedListsIterator> it(
|
||||
invlists->get_iterator(key, inverted_list_context));
|
||||
|
||||
scanner->iterate_codes_range(
|
||||
it.get(), radius, qres, list_size);
|
||||
qres.stats.scan_cnt += list_size;
|
||||
} else {
|
||||
InvertedLists::ScopedCodes scodes(invlists, key);
|
||||
InvertedLists::ScopedIds ids(invlists, key);
|
||||
list_size = invlists->list_size(key);
|
||||
size_t list_size = invlists->list_size(key);
|
||||
|
||||
scanner->scan_codes_range(
|
||||
list_size, scodes.get(), ids.get(), radius, qres);
|
||||
}
|
||||
nlistv++;
|
||||
ndis += list_size;
|
||||
ndis += qres.stats.scan_cnt - scan_cnt0;
|
||||
};
|
||||
|
||||
if (parallel_mode == 0) {
|
||||
@@ -871,11 +915,23 @@ void IndexIVF::range_search_preassigned(
|
||||
for (idx_t i = 0; i < nx; i++) {
|
||||
try {
|
||||
scanner->set_query(x + i * d);
|
||||
|
||||
RangeQueryResult& qres = pres.new_result(i);
|
||||
|
||||
// Stop after enough consecutive probes add no range
|
||||
// results. A hit resets the counter.
|
||||
size_t prev_nres = qres.nres;
|
||||
size_t ndup = 0;
|
||||
for (idx_t ik = 0; ik < cur_nprobe; ik++) {
|
||||
scan_list_func(i, ik, qres);
|
||||
if (max_empty_result_buckets > 0) {
|
||||
// Early-stop check: stop range search after
|
||||
// enough consecutive empty probes.
|
||||
ndup = (qres.nres == prev_nres) ? ndup + 1 : 0;
|
||||
if (ndup >= max_empty_result_buckets) {
|
||||
break;
|
||||
}
|
||||
prev_nres = qres.nres;
|
||||
}
|
||||
}
|
||||
} catch (...) {
|
||||
omp_capture_exception(ex);
|
||||
@@ -1460,10 +1516,14 @@ void InvertedListScanner::scan_codes_range(
|
||||
using C = CMax<float, idx_t>;
|
||||
RangeResultHandler<C, false> handler(&res, radius);
|
||||
scan_codes(list_size, codes, ids, handler);
|
||||
res.stats.scan_cnt += handler.stats.scan_cnt;
|
||||
res.stats.nheap_updates += handler.stats.nheap_updates;
|
||||
} else {
|
||||
using C = CMin<float, idx_t>;
|
||||
RangeResultHandler<C, false> handler(&res, radius);
|
||||
scan_codes(list_size, codes, ids, handler);
|
||||
res.stats.scan_cnt += handler.stats.scan_cnt;
|
||||
res.stats.nheap_updates += handler.stats.nheap_updates;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1472,6 +1532,7 @@ void InvertedListScanner::iterate_codes_range(
|
||||
float radius,
|
||||
RangeQueryResult& res,
|
||||
size_t& list_size) const {
|
||||
size_t nup = 0;
|
||||
list_size = 0;
|
||||
for (; it->is_available(); it->next()) {
|
||||
auto id_and_codes = it->get_id_and_codes();
|
||||
@@ -1481,9 +1542,11 @@ void InvertedListScanner::iterate_codes_range(
|
||||
: dis > radius; // TODO templatize to remove this test
|
||||
if (keep) {
|
||||
res.add(dis, id_and_codes.first);
|
||||
nup++;
|
||||
}
|
||||
list_size++;
|
||||
}
|
||||
res.stats.nheap_updates += nup;
|
||||
}
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
+20
-1
@@ -68,6 +68,25 @@ struct Level1Quantizer {
|
||||
struct SearchParametersIVF : SearchParameters {
|
||||
size_t nprobe = 1; ///< number of probes at query time
|
||||
size_t max_codes = 0; ///< max nb of codes to visit to do a query
|
||||
|
||||
/// FastScan k-NN only: maximum number of inverted lists to visit.
|
||||
/// 0 means unlimited, i.e. bounded only by nprobe. When set together
|
||||
/// with max_codes, either budget may stop the scan. With
|
||||
/// ensure_topk_full, this limit is treated as at least k lists.
|
||||
size_t max_lists_num = 0;
|
||||
|
||||
/// For k-NN search, make small early-stop budgets less aggressive:
|
||||
/// max_codes is treated as at least k post-IDSelector scans. Supported
|
||||
/// by generic IVF in parallel_mode 0 and 3, and by FastScan k-NN
|
||||
/// implementations 10 and 11.
|
||||
bool ensure_topk_full = false;
|
||||
|
||||
/// Range-search only: stop after this many consecutive probed lists add
|
||||
/// no in-radius results. 0 disables the heuristic. This trades recall
|
||||
/// for less work. Supported in parallel_mode 0; FastScan range search
|
||||
/// uses implementation 10 for this option.
|
||||
size_t max_empty_result_buckets = 0;
|
||||
|
||||
SearchParameters* quantizer_params = nullptr;
|
||||
/// context object to pass to InvertedLists
|
||||
void* inverted_list_context = nullptr;
|
||||
@@ -514,7 +533,7 @@ struct InvertedListScanner {
|
||||
* @param k heap size
|
||||
* @return number of heap updates performed
|
||||
*/
|
||||
virtual size_t scan_codes(
|
||||
size_t scan_codes(
|
||||
size_t n,
|
||||
const uint8_t* codes,
|
||||
const idx_t* ids,
|
||||
|
||||
+127
-7
@@ -368,7 +368,11 @@ void IndexIVFFastScan::search_preassigned(
|
||||
IndexIVFStats* stats) const {
|
||||
size_t cur_nprobe = this->nprobe;
|
||||
if (params) {
|
||||
FAISS_THROW_IF_NOT(params->max_codes == 0);
|
||||
// Range-search-only option.
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
params->max_empty_result_buckets == 0,
|
||||
"max_empty_result_buckets is a range-search knob and is "
|
||||
"not honored by fastscan knn search");
|
||||
cur_nprobe = params->nprobe;
|
||||
}
|
||||
|
||||
@@ -395,6 +399,18 @@ void IndexIVFFastScan::range_search(
|
||||
params = dynamic_cast<const IVFSearchParameters*>(params_in);
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
params, "IndexIVFFastScan params have incorrect type");
|
||||
// k-NN-only options.
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
params->max_lists_num == 0,
|
||||
"max_lists_num is a knn knob and is not honored by "
|
||||
"fastscan range search");
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
!params->ensure_topk_full,
|
||||
"ensure_topk_full is a knn knob and is not honored by "
|
||||
"fastscan range search");
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
params->max_codes == 0,
|
||||
"max_codes is not honored by fastscan range search");
|
||||
cur_nprobe = params->nprobe;
|
||||
}
|
||||
FastScanDistancePostProcessing empty_context{};
|
||||
@@ -526,8 +542,14 @@ void IndexIVFFastScan::search_dispatch_implem(
|
||||
// actual implementation used
|
||||
int impl = implem;
|
||||
|
||||
// Early-stop k-NN options require the per-query implementations.
|
||||
const bool any_early_term_knob = params &&
|
||||
(params->max_codes != 0 || params->max_lists_num != 0 ||
|
||||
params->ensure_topk_full);
|
||||
|
||||
if (impl == 0) {
|
||||
if (bbs == 32) {
|
||||
// Auto-select the per-query path when early-stop budgets are used.
|
||||
if (bbs == 32 && !any_early_term_knob) {
|
||||
impl = 12;
|
||||
} else {
|
||||
impl = 10;
|
||||
@@ -544,6 +566,15 @@ void IndexIVFFastScan::search_dispatch_implem(
|
||||
impl -= 100;
|
||||
}
|
||||
|
||||
if (any_early_term_knob) {
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
impl == 10 || impl == 11,
|
||||
"max_codes / max_lists_num / ensure_topk_full are only "
|
||||
"supported by IndexIVFFastScan implem 10/11; set "
|
||||
"index.implem = 10 (or 11 for k>20) explicitly, or leave it "
|
||||
"at the default 0");
|
||||
}
|
||||
|
||||
CoarseQuantizedWithBuffer cq(cq_in);
|
||||
cq.nprobe = cur_nprobe;
|
||||
|
||||
@@ -597,6 +628,7 @@ void IndexIVFFastScan::search_dispatch_implem(
|
||||
search_implem_10(
|
||||
n,
|
||||
x,
|
||||
k,
|
||||
*handler,
|
||||
cq,
|
||||
&ndis,
|
||||
@@ -658,6 +690,7 @@ void IndexIVFFastScan::search_dispatch_implem(
|
||||
search_implem_10(
|
||||
i1 - i0,
|
||||
x + i0 * d,
|
||||
k,
|
||||
*handler,
|
||||
cq_i,
|
||||
&ndis,
|
||||
@@ -695,11 +728,23 @@ void IndexIVFFastScan::range_search_dispatch_implem(
|
||||
if (n == 0) {
|
||||
return;
|
||||
}
|
||||
// FastScan range early-stop budget: enabled only for ordered per-query
|
||||
// scanning below.
|
||||
const bool use_empty_result_early_exit =
|
||||
params && params->max_empty_result_buckets != 0;
|
||||
const int pmode = this->parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT;
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
!use_empty_result_early_exit || pmode == 0,
|
||||
"max_empty_result_buckets supported only for parallel_mode = 0");
|
||||
|
||||
// actual implementation used
|
||||
int impl = implem;
|
||||
|
||||
if (impl == 0) {
|
||||
if (bbs == 32) {
|
||||
if (use_empty_result_early_exit) {
|
||||
// Empty-bucket early stop needs per-query probe order.
|
||||
impl = 10;
|
||||
} else if (bbs == 32) {
|
||||
impl = 12;
|
||||
} else {
|
||||
impl = 10;
|
||||
@@ -715,6 +760,11 @@ void IndexIVFFastScan::range_search_dispatch_implem(
|
||||
impl -= 100;
|
||||
}
|
||||
|
||||
FAISS_THROW_IF_NOT_MSG(
|
||||
!use_empty_result_early_exit || impl == 10,
|
||||
"max_empty_result_buckets is only supported by "
|
||||
"IndexIVFFastScan range-search implem 10");
|
||||
|
||||
if (!multiple_threads && !cq.done()) {
|
||||
cq.quantize(quantizer, n, x, quantizer_params);
|
||||
invlists->prefetch_lists(cq.ids, static_cast<int>(n * cq.nprobe));
|
||||
@@ -740,12 +790,13 @@ void IndexIVFFastScan::range_search_dispatch_implem(
|
||||
search_implem_10(
|
||||
n,
|
||||
x,
|
||||
/*k=*/0, // range search has no k
|
||||
*handler,
|
||||
cq,
|
||||
&ndis,
|
||||
&nlist_visited,
|
||||
context,
|
||||
nullptr,
|
||||
params,
|
||||
*scanner);
|
||||
} else {
|
||||
FAISS_THROW_FMT("Range search implem %d not implemented", impl);
|
||||
@@ -784,12 +835,13 @@ void IndexIVFFastScan::range_search_dispatch_implem(
|
||||
search_implem_10(
|
||||
i1 - i0,
|
||||
x + i0 * d,
|
||||
/*k=*/0,
|
||||
*handler,
|
||||
cq_i,
|
||||
&ndis,
|
||||
&nlist_visited,
|
||||
context,
|
||||
nullptr,
|
||||
params,
|
||||
*scanner);
|
||||
}
|
||||
}
|
||||
@@ -962,12 +1014,13 @@ void IndexIVFFastScan::search_implem_2(
|
||||
void IndexIVFFastScan::search_implem_10(
|
||||
idx_t n,
|
||||
const float* x,
|
||||
idx_t k,
|
||||
SIMDResultHandlerToFloat& handler,
|
||||
const CoarseQuantized& cq,
|
||||
size_t* ndis_out,
|
||||
size_t* nlist_out,
|
||||
const FastScanDistancePostProcessing& context,
|
||||
const IVFSearchParameters* /* params */,
|
||||
const IVFSearchParameters* params,
|
||||
FastScanCodeScanner& scanner) const {
|
||||
size_t dim12 = ksub * M2;
|
||||
AlignedTable<uint8_t> dis_tables;
|
||||
@@ -984,6 +1037,27 @@ void IndexIVFFastScan::search_implem_10(
|
||||
handler.begin(skip & 16 ? nullptr : normalizers.get());
|
||||
size_t cur_nprobe = cq.nprobe;
|
||||
|
||||
// Per-query early-stop options from SearchParametersIVF.
|
||||
const size_t param_max_codes = params ? params->max_codes : 0;
|
||||
const size_t param_max_lists_num = params ? params->max_lists_num : 0;
|
||||
const bool ensure_topk_full = params ? params->ensure_topk_full : false;
|
||||
const size_t cur_max_codes = (param_max_codes == 0)
|
||||
? std::numeric_limits<size_t>::max()
|
||||
: param_max_codes;
|
||||
const size_t cur_max_lists_num =
|
||||
(param_max_lists_num == 0) ? cur_nprobe : param_max_lists_num;
|
||||
// Effective budgets are the values tested in the probe loop below.
|
||||
// ensure_topk_full raises small budgets to reduce empty result slots.
|
||||
const size_t effective_max_codes = ensure_topk_full
|
||||
? std::max(cur_max_codes, (size_t)k)
|
||||
: cur_max_codes;
|
||||
const size_t effective_max_lists_num = ensure_topk_full
|
||||
? std::max(cur_max_lists_num, (size_t)k)
|
||||
: cur_max_lists_num;
|
||||
const bool is_range_search = k == 0;
|
||||
const size_t max_empty_result_buckets =
|
||||
(is_range_search && params) ? params->max_empty_result_buckets : 0;
|
||||
|
||||
// Allocate probe_map once and reuse it
|
||||
std::vector<int> probe_map;
|
||||
probe_map.reserve(1);
|
||||
@@ -995,7 +1069,24 @@ void IndexIVFFastScan::search_implem_10(
|
||||
if (single_LUT) {
|
||||
LUT = dis_tables.get() + i * dim12;
|
||||
}
|
||||
// Per-query counters. For k-NN, the handler count excludes rows
|
||||
// filtered by IDSelector.
|
||||
const size_t scan0 = handler.count_scanned_rows();
|
||||
size_t nscan_q = 0;
|
||||
size_t nlists_visited_q = 0;
|
||||
size_t nempty_result_buckets = 0;
|
||||
for (size_t j = 0; j < cur_nprobe; j++) {
|
||||
if (!is_range_search) {
|
||||
nscan_q = handler.count_scanned_rows() - scan0;
|
||||
}
|
||||
// Early-stop check: apply k-NN max_codes/max_lists_num before
|
||||
// starting the next list. nscan_q excludes IDSelector-filtered
|
||||
// rows.
|
||||
if (nscan_q >= effective_max_codes ||
|
||||
nlists_visited_q >= effective_max_lists_num) {
|
||||
break;
|
||||
}
|
||||
const size_t prev_in_range_num = handler.in_range_num;
|
||||
size_t ij = i * cur_nprobe + j;
|
||||
if (!single_LUT) {
|
||||
LUT = dis_tables.get() + ij * dim12;
|
||||
@@ -1006,10 +1097,22 @@ void IndexIVFFastScan::search_implem_10(
|
||||
|
||||
idx_t list_no = cq.ids[ij];
|
||||
if (list_no < 0) {
|
||||
// Early-stop check: invalid probes count as empty range
|
||||
// buckets.
|
||||
if (max_empty_result_buckets > 0 &&
|
||||
++nempty_result_buckets >= max_empty_result_buckets) {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
size_t ls = invlists->list_size(list_no);
|
||||
if (ls == 0) {
|
||||
// Early-stop check: empty inverted lists count as empty range
|
||||
// buckets.
|
||||
if (max_empty_result_buckets > 0 &&
|
||||
++nempty_result_buckets >= max_empty_result_buckets) {
|
||||
break;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -1028,7 +1131,7 @@ void IndexIVFFastScan::search_implem_10(
|
||||
1,
|
||||
roundup(ls, bbs),
|
||||
bbs,
|
||||
static_cast<int>(M2),
|
||||
M2,
|
||||
codes.get(),
|
||||
LUT,
|
||||
context.pq2x4_scale,
|
||||
@@ -1036,6 +1139,23 @@ void IndexIVFFastScan::search_implem_10(
|
||||
|
||||
ndis += ls;
|
||||
nlist_visited++;
|
||||
if (is_range_search) {
|
||||
nscan_q += ls;
|
||||
}
|
||||
nlists_visited_q++;
|
||||
|
||||
if (max_empty_result_buckets > 0) {
|
||||
// Early-stop check: apply the range-search empty-bucket
|
||||
// budget after each visited list; any hit resets the counter.
|
||||
if (handler.in_range_num == prev_in_range_num) {
|
||||
nempty_result_buckets++;
|
||||
if (nempty_result_buckets >= max_empty_result_buckets) {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
nempty_result_buckets = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -292,10 +292,11 @@ struct IndexIVFFastScan : IndexIVF {
|
||||
const IVFSearchParameters* params = nullptr) const;
|
||||
|
||||
// implem 10 and 12 are not multithreaded internally, so
|
||||
// export search stats
|
||||
// export search stats. k is required for ensure_topk_full support.
|
||||
void search_implem_10(
|
||||
idx_t n,
|
||||
const float* x,
|
||||
idx_t k,
|
||||
SIMDResultHandlerToFloat& handler,
|
||||
const CoarseQuantized& cq,
|
||||
size_t* ndis_out,
|
||||
|
||||
@@ -140,6 +140,11 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
|
||||
local_stats);
|
||||
});
|
||||
|
||||
// num_active is the count of codes for which exact distance
|
||||
// was computed in this batch (post-filter, post-pruning).
|
||||
handler.stats.scan_cnt += num_active;
|
||||
|
||||
// Add batch survivors to heap.
|
||||
for (size_t i = 0; i < num_active; i++) {
|
||||
uint32_t idx = active_indices_[i];
|
||||
size_t global_idx = batch_start + idx;
|
||||
@@ -148,8 +153,10 @@ struct IVFFlatScannerPanorama : InvertedListScanner {
|
||||
if (C::cmp(handler.threshold, dis)) {
|
||||
int64_t id = store_pairs ? lo_build(list_no, global_idx)
|
||||
: ids[global_idx];
|
||||
handler.add_result(dis, id);
|
||||
nup++;
|
||||
if (handler.add_result(dis, id)) {
|
||||
handler.stats.nheap_updates++;
|
||||
nup++;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -799,10 +799,15 @@ struct WrappedSearchResult {
|
||||
}
|
||||
|
||||
inline void add(idx_t j, float dis) {
|
||||
// Reached only for codes that passed skip_entry — i.e. distance
|
||||
// was actually computed for this code (post-filter).
|
||||
res.stats.scan_cnt++;
|
||||
if (C::cmp(res.threshold, dis)) {
|
||||
idx_t id = ids ? ids[j] : lo_build(this->list_no, j);
|
||||
res.add_result(dis, id);
|
||||
nup++;
|
||||
if (res.add_result(dis, id)) {
|
||||
res.stats.nheap_updates++;
|
||||
nup++;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
|
||||
#include <faiss/impl/AuxIndexStructures.h>
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/ResultHandler.h>
|
||||
#include <faiss/impl/simdlib/simdlib_dispatch.h>
|
||||
#include <faiss/utils/Heap.h>
|
||||
#include <faiss/utils/distances.h>
|
||||
@@ -394,64 +395,86 @@ struct IVFPQFastScanScanner : InvertedListScanner {
|
||||
size_t ntotal,
|
||||
const uint8_t* codes,
|
||||
const idx_t* ids,
|
||||
float* distances,
|
||||
idx_t* labels,
|
||||
size_t k) const override {
|
||||
// initialize the current iteration heap to the worst possible value of
|
||||
// the prior loop
|
||||
std::vector<float> curr_dists(k, distances[0]);
|
||||
std::vector<idx_t> curr_labels(k, labels[0]);
|
||||
ResultHandler& handler) const override {
|
||||
auto scan_with_heap = [&](auto* heap_handler) -> size_t {
|
||||
const size_t k = heap_handler->k;
|
||||
if (k == 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
auto scanner = index.make_knn_scanner(
|
||||
!keep_max, nq, k, curr_dists.data(), curr_labels.data(), sel);
|
||||
// initialize the current iteration heap to the worst possible value
|
||||
// of the caller-owned result handler.
|
||||
std::vector<float> curr_dists(k, handler.threshold);
|
||||
std::vector<idx_t> curr_labels(k, -1);
|
||||
|
||||
SIMDResultHandlerToFloat* rh = scanner->handler();
|
||||
|
||||
// This does not quite match search_implem_10, but it is fine because
|
||||
// the scanner operates on a single query at a time, and this value is
|
||||
// used as the query index. For a single query, the value is always 0.
|
||||
int qmap1[1] = {0};
|
||||
|
||||
rh->q_map = qmap1;
|
||||
rh->begin(&normalizers[0]);
|
||||
|
||||
rh->dbias = biases.get();
|
||||
rh->ntotal = ntotal;
|
||||
rh->id_map = ids;
|
||||
|
||||
scanner->accumulate_loop(
|
||||
1,
|
||||
roundup(ntotal, index.bbs),
|
||||
index.bbs,
|
||||
static_cast<int>(index.M2),
|
||||
codes,
|
||||
dis_tables.get(),
|
||||
0,
|
||||
index.get_block_stride());
|
||||
|
||||
// The handler is for the results of this iteration.
|
||||
// Then we need a second heap to combine across iterations.
|
||||
rh->end();
|
||||
|
||||
if (keep_max) {
|
||||
minheap_addn(
|
||||
auto scanner = index.make_knn_scanner(
|
||||
!keep_max,
|
||||
nq,
|
||||
k,
|
||||
distances,
|
||||
labels,
|
||||
curr_dists.data(),
|
||||
curr_labels.data(),
|
||||
k);
|
||||
sel);
|
||||
|
||||
SIMDResultHandlerToFloat* rh = scanner->handler();
|
||||
|
||||
// This does not quite match search_implem_10, but it is fine
|
||||
// because the scanner operates on a single query at a time, and
|
||||
// this value is used as the query index. For a single query, the
|
||||
// value is always 0.
|
||||
int qmap1[1] = {0};
|
||||
|
||||
rh->q_map = qmap1;
|
||||
rh->begin(&normalizers[0]);
|
||||
|
||||
rh->dbias = biases.get();
|
||||
rh->ntotal = ntotal;
|
||||
rh->id_map = ids;
|
||||
|
||||
scanner->accumulate_loop(
|
||||
1,
|
||||
roundup(ntotal, index.bbs),
|
||||
index.bbs,
|
||||
static_cast<int>(index.M2),
|
||||
codes,
|
||||
dis_tables.get(),
|
||||
0,
|
||||
index.get_block_stride());
|
||||
|
||||
const size_t scan_cnt = rh->count_scanned_rows();
|
||||
rh->end();
|
||||
|
||||
handler.stats.scan_cnt += scan_cnt;
|
||||
size_t nup = 0;
|
||||
for (size_t j = 0; j < k; j++) {
|
||||
if (curr_labels[j] < 0) {
|
||||
continue;
|
||||
}
|
||||
if (handler.add_result(curr_dists[j], curr_labels[j])) {
|
||||
handler.stats.nheap_updates++;
|
||||
nup++;
|
||||
}
|
||||
}
|
||||
return nup;
|
||||
};
|
||||
|
||||
if (!keep_max) {
|
||||
using C = CMax<float, idx_t>;
|
||||
if (auto* heap_handler =
|
||||
dynamic_cast<HeapResultHandler<C, false>*>(&handler)) {
|
||||
return scan_with_heap(heap_handler);
|
||||
}
|
||||
} else {
|
||||
maxheap_addn(
|
||||
k,
|
||||
distances,
|
||||
labels,
|
||||
curr_dists.data(),
|
||||
curr_labels.data(),
|
||||
k);
|
||||
using C = CMin<float, idx_t>;
|
||||
if (auto* heap_handler =
|
||||
dynamic_cast<HeapResultHandler<C, false>*>(&handler)) {
|
||||
return scan_with_heap(heap_handler);
|
||||
}
|
||||
}
|
||||
|
||||
return rh->num_updates();
|
||||
FAISS_THROW_MSG(
|
||||
"IVFPQFastScanScanner::scan_codes requires HeapResultHandler; "
|
||||
"custom ResultHandler scan is not supported by this optimized "
|
||||
"scanner");
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -255,10 +255,14 @@ struct RaBitInvertedListScanner : InvertedListScanner {
|
||||
handler.threshold,
|
||||
keep_max);
|
||||
if (should_refine) {
|
||||
// Refining computes the full distance — counts as a
|
||||
// post-filter "distance computed" for stats purposes.
|
||||
handler.stats.scan_cnt++;
|
||||
float dis = distance_to_code(codes);
|
||||
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
|
||||
|
||||
if (handler.add_result(dis, id)) {
|
||||
handler.stats.nheap_updates++;
|
||||
nup++;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/RaBitQUtils.h>
|
||||
#include <faiss/impl/RaBitQuantizerMultiBit.h>
|
||||
#include <faiss/impl/ResultHandler.h>
|
||||
#include <faiss/impl/fast_scan/FastScanDistancePostProcessing.h>
|
||||
#include <faiss/impl/fast_scan/fast_scan.h>
|
||||
#include <faiss/invlists/BlockInvertedLists.h>
|
||||
@@ -874,64 +875,84 @@ struct IVFRaBitQFastScanScanner : InvertedListScanner {
|
||||
size_t ntotal,
|
||||
const uint8_t* codes,
|
||||
const idx_t* ids,
|
||||
float* distances,
|
||||
idx_t* labels,
|
||||
size_t k) const override {
|
||||
std::vector<float> curr_dists(k, distances[0]);
|
||||
std::vector<idx_t> curr_labels(k, labels[0]);
|
||||
ResultHandler& result_handler) const override {
|
||||
auto scan_with_heap = [&](auto* heap_handler) -> size_t {
|
||||
const size_t k = heap_handler->k;
|
||||
if (k == 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
auto scanner = index.make_knn_scanner(
|
||||
!keep_max,
|
||||
nq,
|
||||
k,
|
||||
curr_dists.data(),
|
||||
curr_labels.data(),
|
||||
sel,
|
||||
0,
|
||||
context);
|
||||
auto* handler = scanner->handler();
|
||||
std::vector<float> curr_dists(k, result_handler.threshold);
|
||||
std::vector<idx_t> curr_labels(k, -1);
|
||||
|
||||
int qmap1[1] = {0};
|
||||
handler->q_map = qmap1;
|
||||
handler->begin(&normalizers[0]);
|
||||
handler->dbias = biases.get();
|
||||
handler->ntotal = ntotal;
|
||||
handler->id_map = ids;
|
||||
|
||||
handler->set_list_context(list_no, probe_map);
|
||||
if (!handler->list_codes_ptr) {
|
||||
handler->list_codes_ptr = codes;
|
||||
}
|
||||
|
||||
scanner->accumulate_loop(
|
||||
1,
|
||||
roundup(ntotal, index.bbs),
|
||||
index.bbs,
|
||||
static_cast<int>(index.M2),
|
||||
codes,
|
||||
dis_tables.get(),
|
||||
0,
|
||||
index.get_block_stride());
|
||||
|
||||
handler->end();
|
||||
if (keep_max) {
|
||||
minheap_addn(
|
||||
auto scanner = index.make_knn_scanner(
|
||||
!keep_max,
|
||||
nq,
|
||||
k,
|
||||
distances,
|
||||
labels,
|
||||
curr_dists.data(),
|
||||
curr_labels.data(),
|
||||
k);
|
||||
sel,
|
||||
0,
|
||||
context);
|
||||
auto* handler = scanner->handler();
|
||||
|
||||
int qmap1[1] = {0};
|
||||
handler->q_map = qmap1;
|
||||
handler->begin(&normalizers[0]);
|
||||
handler->dbias = biases.get();
|
||||
handler->ntotal = ntotal;
|
||||
handler->id_map = ids;
|
||||
|
||||
handler->set_list_context(list_no, probe_map);
|
||||
if (!handler->list_codes_ptr) {
|
||||
handler->list_codes_ptr = codes;
|
||||
}
|
||||
|
||||
scanner->accumulate_loop(
|
||||
1,
|
||||
roundup(ntotal, index.bbs),
|
||||
index.bbs,
|
||||
static_cast<int>(index.M2),
|
||||
codes,
|
||||
dis_tables.get(),
|
||||
0,
|
||||
index.get_block_stride());
|
||||
|
||||
const size_t scan_cnt = handler->count_scanned_rows();
|
||||
handler->end();
|
||||
|
||||
result_handler.stats.scan_cnt += scan_cnt;
|
||||
size_t nup = 0;
|
||||
for (size_t j = 0; j < k; j++) {
|
||||
if (curr_labels[j] < 0) {
|
||||
continue;
|
||||
}
|
||||
if (result_handler.add_result(curr_dists[j], curr_labels[j])) {
|
||||
result_handler.stats.nheap_updates++;
|
||||
nup++;
|
||||
}
|
||||
}
|
||||
return nup;
|
||||
};
|
||||
|
||||
if (!keep_max) {
|
||||
using C = CMax<float, idx_t>;
|
||||
if (auto* heap_handler = dynamic_cast<HeapResultHandler<C, false>*>(
|
||||
&result_handler)) {
|
||||
return scan_with_heap(heap_handler);
|
||||
}
|
||||
} else {
|
||||
maxheap_addn(
|
||||
k,
|
||||
distances,
|
||||
labels,
|
||||
curr_dists.data(),
|
||||
curr_labels.data(),
|
||||
k);
|
||||
using C = CMin<float, idx_t>;
|
||||
if (auto* heap_handler = dynamic_cast<HeapResultHandler<C, false>*>(
|
||||
&result_handler)) {
|
||||
return scan_with_heap(heap_handler);
|
||||
}
|
||||
}
|
||||
return handler->num_updates();
|
||||
|
||||
FAISS_THROW_MSG(
|
||||
"IVFRaBitQFastScanScanner::scan_codes requires "
|
||||
"HeapResultHandler; custom ResultHandler scan is not supported "
|
||||
"by this optimized scanner");
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -18,6 +18,8 @@
|
||||
#include <mutex>
|
||||
#include <vector>
|
||||
|
||||
#include <faiss/impl/InvertedListScannerStats.h>
|
||||
|
||||
#include <faiss/MetricType.h>
|
||||
#include <faiss/impl/platform_macros.h>
|
||||
|
||||
@@ -92,6 +94,7 @@ struct RangeQueryResult {
|
||||
idx_t qno; //< id of the query
|
||||
size_t nres; //< nb of results for this query
|
||||
RangeSearchPartialResult* pres;
|
||||
InvertedListScannerStats stats;
|
||||
|
||||
/// called by search function to report a new result
|
||||
void add(float dis, idx_t id);
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
/*
|
||||
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#ifndef FAISS_INVERTED_LIST_SCANNER_STATS_H
|
||||
#define FAISS_INVERTED_LIST_SCANNER_STATS_H
|
||||
|
||||
#include <cstddef>
|
||||
|
||||
namespace faiss {
|
||||
|
||||
/** Per-list statistics returned by inverted-list scanners. */
|
||||
struct InvertedListScannerStats {
|
||||
/// Number of distances computed after IDSelector filtering.
|
||||
size_t scan_cnt = 0;
|
||||
|
||||
/// Number of heap updates.
|
||||
size_t nheap_updates = 0;
|
||||
};
|
||||
|
||||
} // namespace faiss
|
||||
|
||||
#endif
|
||||
@@ -15,6 +15,7 @@
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/FaissException.h>
|
||||
#include <faiss/impl/IDSelector.h>
|
||||
#include <faiss/impl/InvertedListScannerStats.h>
|
||||
#include <faiss/utils/Heap.h>
|
||||
#include <faiss/utils/partitioning.h>
|
||||
#include <algorithm>
|
||||
@@ -29,6 +30,9 @@ struct ResultHandlerUnordered {
|
||||
// if not better than threshold, then not necessary to call add_result
|
||||
T threshold{};
|
||||
|
||||
// per-list scan statistics populated by inverted-list scanners
|
||||
InvertedListScannerStats stats;
|
||||
|
||||
// return whether threshold was updated
|
||||
virtual bool add_result(T dis, TI idx) = 0;
|
||||
|
||||
|
||||
@@ -26,6 +26,7 @@
|
||||
#include <faiss/VectorTransform.h>
|
||||
#include <faiss/impl/AuxIndexStructures.h>
|
||||
#include <faiss/impl/FaissAssert.h>
|
||||
#include <faiss/impl/ResultHandler.h>
|
||||
#include <faiss/impl/binary_hamming/dispatch.h>
|
||||
#include <faiss/utils/hamming.h>
|
||||
|
||||
@@ -103,17 +104,20 @@ struct IVFScanner : InvertedListScanner {
|
||||
size_t list_size,
|
||||
const uint8_t* codes,
|
||||
const idx_t* ids,
|
||||
float* simi,
|
||||
idx_t* idxi,
|
||||
size_t k) const override {
|
||||
ResultHandler& handler) const override {
|
||||
size_t nup = 0;
|
||||
// get_InvertedListScanner asserts no IDSelector, so every code
|
||||
// in the list has its distance computed.
|
||||
for (size_t j = 0; j < list_size; j++) {
|
||||
float dis = hc.hamming(codes);
|
||||
handler.stats.scan_cnt++;
|
||||
|
||||
if (dis < simi[0]) {
|
||||
if (dis < handler.threshold) {
|
||||
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
|
||||
maxheap_replace_top(k, simi, idxi, dis, id);
|
||||
nup++;
|
||||
if (handler.add_result(dis, id)) {
|
||||
handler.stats.nheap_updates++;
|
||||
nup++;
|
||||
}
|
||||
}
|
||||
codes += code_size;
|
||||
}
|
||||
@@ -128,9 +132,11 @@ struct IVFScanner : InvertedListScanner {
|
||||
RangeQueryResult& result) const override {
|
||||
for (size_t j = 0; j < list_size; j++) {
|
||||
float dis = hc.hamming(codes);
|
||||
result.stats.scan_cnt++;
|
||||
if (dis < radius) {
|
||||
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
|
||||
result.add(dis, id);
|
||||
result.stats.nheap_updates++;
|
||||
}
|
||||
codes += code_size;
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
#include <cstdio>
|
||||
|
||||
#include <faiss/IndexIVF.h>
|
||||
#include <faiss/impl/InvertedListScannerStats.h>
|
||||
#include <faiss/impl/ResultHandler.h>
|
||||
|
||||
/* This is the inner loop of the inverted list scanners. The default version
|
||||
@@ -45,12 +46,16 @@ size_t run_scan_codes1(
|
||||
}
|
||||
}
|
||||
|
||||
// post-IDSelector: distance is about to be computed for this code.
|
||||
handler.stats.scan_cnt++;
|
||||
float dis = scanner.distance_to_code(codes); // will be inlined if final
|
||||
if (C::cmp(threshold, dis)) {
|
||||
int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
|
||||
handler.add_result(dis, id);
|
||||
threshold = handler.threshold;
|
||||
nup++;
|
||||
if (handler.add_result(dis, id)) {
|
||||
handler.stats.nheap_updates++;
|
||||
nup++;
|
||||
threshold = handler.threshold;
|
||||
}
|
||||
}
|
||||
codes += code_size;
|
||||
}
|
||||
|
||||
@@ -2268,6 +2268,11 @@ def range_search_with_parameters(
|
||||
class IVFSearchParameters(SearchParameters):
|
||||
nprobe: int
|
||||
max_codes: int
|
||||
max_lists_num: int
|
||||
ensure_topk_full: bool
|
||||
max_empty_result_buckets: int
|
||||
quantizer_params: SearchParameters | None
|
||||
inverted_list_context: Any
|
||||
sel: IDSelector | None
|
||||
def __init__(self) -> None: ...
|
||||
|
||||
|
||||
@@ -93,6 +93,7 @@ typedef uint64_t size_t;
|
||||
#include <faiss/IndexLSH.h>
|
||||
#include <faiss/IndexPQ.h>
|
||||
#include <faiss/IndexAdditiveQuantizer.h>
|
||||
#include <faiss/impl/InvertedListScannerStats.h>
|
||||
#include <faiss/IndexIVF.h>
|
||||
#include <faiss/IndexIVFPQ.h>
|
||||
#include <faiss/Index2Layer.h>
|
||||
@@ -547,6 +548,7 @@ void gpu_sync_all_devices()
|
||||
%ignore BlockInvertedListsIOHook;
|
||||
%include <faiss/invlists/BlockInvertedLists.h>
|
||||
%include <faiss/invlists/DirectMap.h>
|
||||
%include <faiss/impl/InvertedListScannerStats.h>
|
||||
%include <faiss/IndexIVF.h>
|
||||
// NOTE(hoss): SWIG (wrongly) believes the overloaded const version shadows the
|
||||
// non-const one.
|
||||
|
||||
@@ -30,6 +30,7 @@ set(FAISS_TEST_SRC
|
||||
test_fastscan_perf.cpp
|
||||
test_pqfs_unaligned.cpp
|
||||
test_common_ivf_empty_index.cpp
|
||||
test_ivf_early_termination.cpp
|
||||
test_callback.cpp
|
||||
test_utils.cpp
|
||||
test_hamming.cpp
|
||||
|
||||
@@ -1019,3 +1019,162 @@ class TestRangeSearchImplem10(TestRangeSearchImplem12):
|
||||
@for_all_simd_levels
|
||||
class TestRangeSearchImplem110(TestRangeSearchImplem12):
|
||||
IMPLEM = 110
|
||||
|
||||
|
||||
class TestFastScanEarlyTerminationKnobs(unittest.TestCase):
|
||||
"""
|
||||
FastScan early-stop options use per-query implementations.
|
||||
Incompatible explicit implementations throw on opt-in.
|
||||
"""
|
||||
|
||||
def _build_index(self):
|
||||
ds = datasets.SyntheticDataset(32, 2000, 500, 30)
|
||||
index = faiss.index_factory(32, "IVF32,PQ4x4fs", faiss.METRIC_L2)
|
||||
index.train(ds.get_train())
|
||||
index.add(ds.get_database())
|
||||
return ds, index
|
||||
|
||||
def test_fastscan_defaults_roundtrip(self):
|
||||
ds, index = self._build_index()
|
||||
k = 10
|
||||
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 4
|
||||
# defaults should not raise
|
||||
D_p, I_p = index.search(ds.get_queries(), k, params=params)
|
||||
# and should match a no-params search with the same nprobe.
|
||||
index.nprobe = 4
|
||||
D_ref, I_ref = index.search(ds.get_queries(), k)
|
||||
np.testing.assert_array_equal(I_p, I_ref)
|
||||
|
||||
def test_fastscan_max_lists_num_caps_probes(self):
|
||||
ds, index = self._build_index()
|
||||
# implem auto-selection should pick 10 when max_lists_num is set
|
||||
index.implem = 0
|
||||
params_full = faiss.SearchParametersIVF()
|
||||
params_full.nprobe = 16
|
||||
D_full, I_full = index.search(ds.get_queries(), 10, params=params_full)
|
||||
|
||||
params_cap = faiss.SearchParametersIVF()
|
||||
params_cap.nprobe = 16
|
||||
params_cap.max_lists_num = 2
|
||||
D_cap, I_cap = index.search(ds.get_queries(), 10, params=params_cap)
|
||||
|
||||
# Tight cap should produce no more valid hits than full nprobe and
|
||||
# at least one query should differ (unless data is unusually small).
|
||||
miss_full = np.sum(I_full == -1)
|
||||
miss_cap = np.sum(I_cap == -1)
|
||||
self.assertLessEqual(miss_full, miss_cap)
|
||||
|
||||
def test_fastscan_ensure_topk_full_fills_heap(self):
|
||||
ds, index = self._build_index()
|
||||
index.implem = 0
|
||||
# Use one query so the missing-result count is easy to compare.
|
||||
q = ds.get_queries()[:1]
|
||||
k = 10
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 16
|
||||
params.max_codes = k // 2 # tight: forces truncation without
|
||||
# ensure_topk_full
|
||||
params.ensure_topk_full = False
|
||||
_, I_tight = index.search(q, k, params=params)
|
||||
params.ensure_topk_full = True
|
||||
_, I_full = index.search(q, k, params=params)
|
||||
self.assertGreaterEqual(
|
||||
int(np.sum(I_tight == -1)), int(np.sum(I_full == -1))
|
||||
)
|
||||
|
||||
def test_fastscan_ensure_topk_full_with_restrictive_selector(self):
|
||||
ds, index = self._build_index()
|
||||
index.implem = 0
|
||||
q = ds.get_queries()[:1]
|
||||
k = 10
|
||||
|
||||
# Keep one third of the ids: restrictive enough to catch raw
|
||||
# pre-filter max_codes accounting, while still leaving enough
|
||||
# approximate FastScan candidates to fill the heap.
|
||||
keep = np.arange(0, ds.nb, 3).astype("int64")
|
||||
sel = faiss.IDSelectorBatch(keep)
|
||||
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 32
|
||||
params.max_codes = 2
|
||||
params.ensure_topk_full = True
|
||||
params.sel = sel
|
||||
|
||||
_, I = index.search(q, k, params=params)
|
||||
self.assertEqual(int(np.sum(I == -1)), 0)
|
||||
|
||||
def test_fastscan_ensure_topk_full_multi_query_works(self):
|
||||
# Each query has its own per-query early-termination counters in
|
||||
# search_implem_10, so multi-query batches are supported.
|
||||
ds, index = self._build_index()
|
||||
index.implem = 0
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 4
|
||||
params.ensure_topk_full = True
|
||||
# No throw expected:
|
||||
index.search(ds.get_queries()[:2], 10, params=params)
|
||||
|
||||
def test_fastscan_explicit_implem_12_rejects_knobs(self):
|
||||
ds, index = self._build_index()
|
||||
index.implem = 12
|
||||
for knob, val in [
|
||||
("max_lists_num", 2),
|
||||
("ensure_topk_full", True),
|
||||
("max_codes", 5),
|
||||
]:
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 4
|
||||
setattr(params, knob, val)
|
||||
with self.assertRaises(RuntimeError, msg=f"knob {knob}"):
|
||||
index.search(ds.get_queries()[:1], 10, params=params)
|
||||
|
||||
def test_fastscan_range_honors_max_empty_result_buckets(self):
|
||||
ds, index = self._build_index()
|
||||
q = ds.get_queries()[:1]
|
||||
empty_sel = faiss.IDSelectorBatch(np.array([], dtype="int64"))
|
||||
stats = faiss.cvar.indexIVF_stats
|
||||
|
||||
params_full = faiss.SearchParametersIVF()
|
||||
params_full.nprobe = 4
|
||||
params_full.sel = empty_sel
|
||||
|
||||
stats.reset()
|
||||
lims_full, _, _ = index.range_search(q, 1.0, params=params_full)
|
||||
full_nlist = stats.nlist
|
||||
|
||||
params_early = faiss.SearchParametersIVF()
|
||||
params_early.nprobe = 4
|
||||
params_early.max_empty_result_buckets = 2
|
||||
params_early.sel = empty_sel
|
||||
|
||||
stats.reset()
|
||||
lims_early, _, _ = index.range_search(q, 1.0, params=params_early)
|
||||
early_nlist = stats.nlist
|
||||
|
||||
np.testing.assert_array_equal(lims_full, lims_early)
|
||||
self.assertEqual(lims_early[-1], 0)
|
||||
self.assertLess(early_nlist, full_nlist)
|
||||
|
||||
def test_fastscan_range_explicit_implem_12_rejects_max_empty(self):
|
||||
ds, index = self._build_index()
|
||||
index.implem = 12
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 4
|
||||
params.max_empty_result_buckets = 2
|
||||
with self.assertRaises(RuntimeError):
|
||||
index.range_search(ds.get_queries()[:1], 1.0, params=params)
|
||||
|
||||
def test_fastscan_fields_exposed(self):
|
||||
p = faiss.SearchParametersIVF()
|
||||
self.assertEqual(p.max_lists_num, 0)
|
||||
self.assertIs(p.ensure_topk_full, False)
|
||||
self.assertEqual(p.max_empty_result_buckets, 0)
|
||||
p.max_lists_num = 7
|
||||
p.ensure_topk_full = True
|
||||
p.max_empty_result_buckets = 3
|
||||
self.assertEqual(p.max_lists_num, 7)
|
||||
self.assertIs(p.ensure_topk_full, True)
|
||||
self.assertEqual(p.max_empty_result_buckets, 3)
|
||||
self.assertIs(faiss.SearchParametersIVF, faiss.IVFSearchParameters)
|
||||
|
||||
@@ -0,0 +1,755 @@
|
||||
/*
|
||||
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
// Unit tests for the three early-termination knobs on SearchParametersIVF:
|
||||
// - ensure_topk_full (knn, IndexIVF::search_preassigned)
|
||||
// - max_empty_result_buckets (range, IndexIVF::range_search_preassigned)
|
||||
// - max_lists_num (fastscan knn)
|
||||
//
|
||||
// The tests establish the default-behavior contract: when the new fields
|
||||
// are left at their defaults, search results must be byte-for-byte identical
|
||||
// to a search with no params at all.
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <random>
|
||||
#include <set>
|
||||
#include <vector>
|
||||
|
||||
#include <gtest/gtest.h>
|
||||
|
||||
#include <faiss/IndexFlat.h>
|
||||
#include <faiss/IndexIVFFlat.h>
|
||||
#include <faiss/IndexIVFPQFastScan.h>
|
||||
#include <faiss/impl/AuxIndexStructures.h>
|
||||
#include <faiss/impl/FaissException.h>
|
||||
#include <faiss/impl/IDSelector.h>
|
||||
#include <faiss/index_factory.h>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int d = 16;
|
||||
constexpr size_t nb = 500;
|
||||
constexpr size_t nq = 20;
|
||||
constexpr size_t nlist = 16;
|
||||
constexpr int k = 10;
|
||||
|
||||
std::vector<float> make_data(size_t n, uint32_t seed) {
|
||||
std::mt19937 rng(seed);
|
||||
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
|
||||
std::vector<float> out(n * d);
|
||||
for (auto& v : out) {
|
||||
v = dist(rng);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
std::unique_ptr<faiss::Index> build_ivf_flat() {
|
||||
auto xb = make_data(nb, 0xabcdef);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,Flat", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
return index;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Default-constructed SearchParametersIVF (all three new fields at 0/false)
|
||||
// must reproduce the result of a search without any SearchParameters.
|
||||
TEST(IVFEarlyTermination, DefaultsPreserveBaselineKnn) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->nprobe = 4;
|
||||
|
||||
auto xq = make_data(nq, 0x123456);
|
||||
std::vector<float> D_ref(k * nq), D_new(k * nq);
|
||||
std::vector<faiss::idx_t> I_ref(k * nq), I_new(k * nq);
|
||||
|
||||
index->search(nq, xq.data(), k, D_ref.data(), I_ref.data());
|
||||
|
||||
faiss::SearchParametersIVF params; // all defaults
|
||||
params.nprobe = 4;
|
||||
index->search(nq, xq.data(), k, D_new.data(), I_new.data(), ¶ms);
|
||||
|
||||
EXPECT_EQ(I_ref, I_new);
|
||||
EXPECT_EQ(D_ref, D_new);
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, DefaultsPreserveBaselineRange) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->nprobe = 4;
|
||||
|
||||
auto xq = make_data(nq, 0x789abc);
|
||||
const float radius = 2.0f;
|
||||
|
||||
faiss::RangeSearchResult ref(nq);
|
||||
faiss::RangeSearchResult tst(nq);
|
||||
|
||||
index->range_search(nq, xq.data(), radius, &ref);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
index->range_search(nq, xq.data(), radius, &tst, ¶ms);
|
||||
|
||||
ASSERT_EQ(ref.nq, tst.nq);
|
||||
for (size_t i = 0; i <= ref.nq; i++) {
|
||||
EXPECT_EQ(ref.lims[i], tst.lims[i]) << "at lims[" << i << "]";
|
||||
}
|
||||
// The underlying result order is implementation-defined across OMP
|
||||
// threads, so compare as sorted (id, distance) pairs per query.
|
||||
for (size_t q = 0; q < ref.nq; q++) {
|
||||
std::vector<std::pair<faiss::idx_t, float>> ref_pairs;
|
||||
std::vector<std::pair<faiss::idx_t, float>> tst_pairs;
|
||||
for (size_t j = ref.lims[q]; j < ref.lims[q + 1]; j++) {
|
||||
ref_pairs.emplace_back(ref.labels[j], ref.distances[j]);
|
||||
}
|
||||
for (size_t j = tst.lims[q]; j < tst.lims[q + 1]; j++) {
|
||||
tst_pairs.emplace_back(tst.labels[j], tst.distances[j]);
|
||||
}
|
||||
std::sort(ref_pairs.begin(), ref_pairs.end());
|
||||
std::sort(tst_pairs.begin(), tst_pairs.end());
|
||||
EXPECT_EQ(ref_pairs, tst_pairs);
|
||||
}
|
||||
}
|
||||
|
||||
// With a tight max_codes, ensure_topk_full=true must produce at least as
|
||||
// many valid (non -1) labels as ensure_topk_full=false.
|
||||
TEST(IVFEarlyTermination, EnsureTopkFullReducesMisses) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
|
||||
auto xq = make_data(nq, 0xdeadbe);
|
||||
std::vector<float> D_tight(k * nq), D_full(k * nq);
|
||||
std::vector<faiss::idx_t> I_tight(k * nq), I_full(k * nq);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = nlist;
|
||||
// Tight budget: strictly smaller than k to guarantee truncation.
|
||||
params.max_codes = k / 2;
|
||||
params.ensure_topk_full = false;
|
||||
index->search(nq, xq.data(), k, D_tight.data(), I_tight.data(), ¶ms);
|
||||
|
||||
params.ensure_topk_full = true;
|
||||
index->search(nq, xq.data(), k, D_full.data(), I_full.data(), ¶ms);
|
||||
|
||||
size_t miss_tight =
|
||||
std::count(I_tight.begin(), I_tight.end(), (faiss::idx_t)-1);
|
||||
size_t miss_full =
|
||||
std::count(I_full.begin(), I_full.end(), (faiss::idx_t)-1);
|
||||
|
||||
EXPECT_LE(miss_full, miss_tight);
|
||||
EXPECT_GT(miss_tight, size_t(0))
|
||||
<< "test precondition: tight budget should cause truncation";
|
||||
EXPECT_EQ(miss_full, size_t(0))
|
||||
<< "ensure_topk_full must fill the heap when enough data is "
|
||||
"available in the probed lists";
|
||||
}
|
||||
|
||||
// The field is checked against the existing max_codes throw guard for
|
||||
// parallel modes other than 0/3.
|
||||
TEST(IVFEarlyTermination, EnsureTopkFullRejectsUnsupportedParallelMode) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->parallel_mode = 1;
|
||||
|
||||
auto xq = make_data(1, 0xaaaa);
|
||||
std::vector<float> D(k);
|
||||
std::vector<faiss::idx_t> I(k);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.ensure_topk_full = true;
|
||||
EXPECT_THROW(
|
||||
index->search(1, xq.data(), k, D.data(), I.data(), ¶ms),
|
||||
faiss::FaissException);
|
||||
}
|
||||
|
||||
// Range-search early-exit: with max_empty_result_buckets set, the result
|
||||
// set is a subset of the full-scan result set.
|
||||
TEST(IVFEarlyTermination, RangeEarlyExitProducesSubset) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
|
||||
auto xq = make_data(nq, 0xfacefeed);
|
||||
// Pick a radius tight enough that some probed buckets legitimately
|
||||
// yield no hits.
|
||||
const float radius = 0.5f;
|
||||
|
||||
faiss::RangeSearchResult full(nq);
|
||||
faiss::RangeSearchResult early(nq);
|
||||
|
||||
faiss::SearchParametersIVF p_full;
|
||||
p_full.nprobe = nlist;
|
||||
p_full.max_empty_result_buckets = 0; // disabled
|
||||
index->range_search(nq, xq.data(), radius, &full, &p_full);
|
||||
|
||||
faiss::SearchParametersIVF p_early;
|
||||
p_early.nprobe = nlist;
|
||||
p_early.max_empty_result_buckets = 1; // aggressive
|
||||
index->range_search(nq, xq.data(), radius, &early, &p_early);
|
||||
|
||||
ASSERT_EQ(full.nq, early.nq);
|
||||
for (size_t q = 0; q < full.nq; q++) {
|
||||
std::set<faiss::idx_t> full_ids(
|
||||
full.labels + full.lims[q], full.labels + full.lims[q + 1]);
|
||||
std::set<faiss::idx_t> early_ids(
|
||||
early.labels + early.lims[q], early.labels + early.lims[q + 1]);
|
||||
// early_ids ⊆ full_ids
|
||||
for (auto id : early_ids) {
|
||||
EXPECT_TRUE(full_ids.count(id))
|
||||
<< "early-exit returned id " << id
|
||||
<< " not present in full-scan result for query " << q;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Verify that the empty-bucket early exit actually stops the probe loop.
|
||||
TEST(IVFEarlyTermination, RangeEarlyExitActuallyTruncates) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
|
||||
auto xq = make_data(nq, 0xfacefeed);
|
||||
// Radius chosen so that the dataset (uniform[0,1]^16, nb=500,
|
||||
// nlist=16) yields per-query results spread across multiple
|
||||
// probed lists with interspersed empty buckets — observable
|
||||
// shrinkage when max_empty_result_buckets fires.
|
||||
const float radius = 1.0f;
|
||||
|
||||
faiss::RangeSearchResult full(nq);
|
||||
faiss::RangeSearchResult early(nq);
|
||||
|
||||
faiss::SearchParametersIVF p_full;
|
||||
p_full.nprobe = nlist;
|
||||
p_full.max_empty_result_buckets = 0; // disabled
|
||||
index->range_search(nq, xq.data(), radius, &full, &p_full);
|
||||
|
||||
faiss::SearchParametersIVF p_early;
|
||||
p_early.nprobe = nlist;
|
||||
p_early.max_empty_result_buckets = 1; // exit on first empty bucket
|
||||
index->range_search(nq, xq.data(), radius, &early, &p_early);
|
||||
|
||||
ASSERT_EQ(full.nq, early.nq);
|
||||
|
||||
// Per-query result count must never grow under early-exit.
|
||||
for (size_t q = 0; q < full.nq; q++) {
|
||||
EXPECT_LE(
|
||||
early.lims[q + 1] - early.lims[q],
|
||||
full.lims[q + 1] - full.lims[q])
|
||||
<< "early-exit grew result count for query " << q;
|
||||
}
|
||||
|
||||
// At least one query must strictly shrink — proves the loop broke.
|
||||
bool any_shrunk = false;
|
||||
for (size_t q = 0; q < full.nq; q++) {
|
||||
if ((early.lims[q + 1] - early.lims[q]) <
|
||||
(full.lims[q + 1] - full.lims[q])) {
|
||||
any_shrunk = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
EXPECT_TRUE(any_shrunk)
|
||||
<< "no query shrank: max_empty_result_buckets has no effect";
|
||||
}
|
||||
|
||||
// max_empty_result_buckets counter must reset on probes that add hits.
|
||||
// Hand-verified: with ndup=2 we should NEVER terminate as long as a
|
||||
// non-empty bucket appears between any two empty buckets, even if every
|
||||
// other probe is empty. Test by forcing alternating-empty patterns.
|
||||
TEST(IVFEarlyTermination, RangeEarlyExitCounterResetsOnHit) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
|
||||
auto xq = make_data(nq, 0xc0ffee);
|
||||
const float radius = 1.0f;
|
||||
|
||||
// With max_empty_result_buckets=1, exit on first empty bucket.
|
||||
// With =3, only exit on three consecutive empties — strictly more
|
||||
// results expected (or at least equal) for the latter.
|
||||
faiss::RangeSearchResult tight(nq);
|
||||
faiss::RangeSearchResult relaxed(nq);
|
||||
|
||||
faiss::SearchParametersIVF p_tight;
|
||||
p_tight.nprobe = nlist;
|
||||
p_tight.max_empty_result_buckets = 1;
|
||||
index->range_search(nq, xq.data(), radius, &tight, &p_tight);
|
||||
|
||||
faiss::SearchParametersIVF p_relaxed;
|
||||
p_relaxed.nprobe = nlist;
|
||||
p_relaxed.max_empty_result_buckets = 3;
|
||||
index->range_search(nq, xq.data(), radius, &relaxed, &p_relaxed);
|
||||
|
||||
ASSERT_EQ(tight.nq, relaxed.nq);
|
||||
for (size_t q = 0; q < tight.nq; q++) {
|
||||
EXPECT_LE(
|
||||
tight.lims[q + 1] - tight.lims[q],
|
||||
relaxed.lims[q + 1] - relaxed.lims[q])
|
||||
<< "relaxed early-exit returned fewer results than tight "
|
||||
"for query "
|
||||
<< q;
|
||||
}
|
||||
}
|
||||
|
||||
// max_empty_result_buckets is only supported in parallel_mode == 0.
|
||||
TEST(IVFEarlyTermination, RangeEarlyExitRejectsUnsupportedParallelMode) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->parallel_mode = 1;
|
||||
|
||||
auto xq = make_data(1, 0xbbbb);
|
||||
faiss::RangeSearchResult res(1);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.max_empty_result_buckets = 2;
|
||||
EXPECT_THROW(
|
||||
index->range_search(1, xq.data(), 1.0f, &res, ¶ms),
|
||||
faiss::FaissException);
|
||||
}
|
||||
|
||||
// IndexIVFStats::ndis is documented as "nb of distances computed".
|
||||
// Verify that the post-filter scan_count flowing through
|
||||
// InvertedListScannerStats actually populates ndis according to that
|
||||
// docstring: under no IDSelector, ndis equals the sum of probed list
|
||||
// sizes; under a 50%-keep IDSelectorBitmap, ndis is roughly halved
|
||||
// because only codes that survive the selector have distances
|
||||
// computed.
|
||||
TEST(IVFEarlyTermination, NdisStatsMatchesDocstringPostFilter) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->nprobe = nlist;
|
||||
|
||||
auto xq = make_data(nq, 0xb0bafe11);
|
||||
|
||||
// Baseline: no selector. ndis equals the total raw codes the
|
||||
// scanner saw.
|
||||
auto& stats = faiss::indexIVF_stats;
|
||||
stats.reset();
|
||||
std::vector<float> D(k * nq);
|
||||
std::vector<faiss::idx_t> I(k * nq);
|
||||
index->search(nq, xq.data(), k, D.data(), I.data());
|
||||
const size_t ndis_no_sel = stats.ndis;
|
||||
EXPECT_GT(ndis_no_sel, size_t(0))
|
||||
<< "test precondition: at least some codes should be visited";
|
||||
|
||||
// 50%-keep bitmap selector. ndis should drop because rejected
|
||||
// codes don't get a distance computed.
|
||||
std::vector<uint8_t> packed((nb + 7) / 8, 0);
|
||||
for (size_t i = 0; i < nb; i += 2) {
|
||||
packed[i / 8] |= 1u << (i % 8);
|
||||
}
|
||||
faiss::IDSelectorBitmap sel(nb, packed.data());
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = nlist;
|
||||
params.sel = &sel;
|
||||
|
||||
stats.reset();
|
||||
index->search(nq, xq.data(), k, D.data(), I.data(), ¶ms);
|
||||
const size_t ndis_with_sel = stats.ndis;
|
||||
|
||||
// Strict: with a 50%-keep selector, ndis must be strictly less
|
||||
// than the no-selector ndis (each call processes the same lists
|
||||
// with the same nprobe).
|
||||
EXPECT_LT(ndis_with_sel, ndis_no_sel)
|
||||
<< "ndis did not decrease under restrictive IDSelector — "
|
||||
"check post-filter scan_cnt plumbing";
|
||||
// Loose-bound sanity: ~50% retention should put ndis somewhere
|
||||
// around half. Allow generous slop because per-list size variance
|
||||
// means the empirical ratio can drift.
|
||||
EXPECT_LT(ndis_with_sel * 4, ndis_no_sel * 3)
|
||||
<< "ndis_with_sel=" << ndis_with_sel
|
||||
<< " not noticeably below ~75% of ndis_no_sel=" << ndis_no_sel;
|
||||
}
|
||||
|
||||
// Same docstring contract for IndexIVF::range_search_preassigned: ndis
|
||||
// must be post-filter for the range path too.
|
||||
TEST(IVFEarlyTermination, NdisStatsMatchesDocstringPostFilterRange) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->nprobe = nlist;
|
||||
|
||||
auto xq = make_data(nq, 0xc01dca11);
|
||||
const float radius = 1.0f;
|
||||
|
||||
auto& stats = faiss::indexIVF_stats;
|
||||
stats.reset();
|
||||
faiss::RangeSearchResult res_no_sel(nq);
|
||||
index->range_search(nq, xq.data(), radius, &res_no_sel);
|
||||
const size_t ndis_no_sel = stats.ndis;
|
||||
EXPECT_GT(ndis_no_sel, size_t(0));
|
||||
|
||||
std::vector<uint8_t> packed((nb + 7) / 8, 0);
|
||||
for (size_t i = 0; i < nb; i += 2) {
|
||||
packed[i / 8] |= 1u << (i % 8);
|
||||
}
|
||||
faiss::IDSelectorBitmap sel(nb, packed.data());
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = nlist;
|
||||
params.sel = &sel;
|
||||
|
||||
stats.reset();
|
||||
faiss::RangeSearchResult res_with_sel(nq);
|
||||
index->range_search(nq, xq.data(), radius, &res_with_sel, ¶ms);
|
||||
const size_t ndis_with_sel = stats.ndis;
|
||||
|
||||
EXPECT_LT(ndis_with_sel, ndis_no_sel);
|
||||
EXPECT_LT(ndis_with_sel * 4, ndis_no_sel * 3);
|
||||
}
|
||||
|
||||
// With a restrictive bitmap selector, ensure_topk_full keeps probing until
|
||||
// enough selector-passing candidates have been scanned.
|
||||
TEST(IVFEarlyTermination, EnsureTopkFullFillsHeapUnderRestrictiveSelector) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
ivf->nprobe = nlist;
|
||||
|
||||
// ~10%-keep bitmap selector (every 10th id).
|
||||
std::vector<uint8_t> packed((nb + 7) / 8, 0);
|
||||
for (size_t i = 0; i < nb; i += 10) {
|
||||
packed[i / 8] |= 1u << (i % 8);
|
||||
}
|
||||
faiss::IDSelectorBitmap sel(nb, packed.data());
|
||||
|
||||
auto xq = make_data(nq, 0xdeadbe);
|
||||
std::vector<float> D(k * nq);
|
||||
std::vector<faiss::idx_t> I(k * nq);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = nlist;
|
||||
params.max_codes = k / 2; // pre-filter budget < k
|
||||
params.ensure_topk_full = true; // post-filter soft-cap kicks in
|
||||
params.sel = &sel;
|
||||
|
||||
index->search(nq, xq.data(), k, D.data(), I.data(), ¶ms);
|
||||
|
||||
// Valid hits available per query (~50 across all nlist buckets,
|
||||
// far exceeding k=10) — the strengthened break should fill the
|
||||
// heap completely.
|
||||
EXPECT_EQ(std::count(I.begin(), I.end(), (faiss::idx_t)-1), size_t(0))
|
||||
<< "ensure_topk_full failed to fill the heap under a "
|
||||
"restrictive IDSelector";
|
||||
}
|
||||
|
||||
// Without a selector, post-filter and raw scan counts are the same.
|
||||
TEST(IVFEarlyTermination, EnsureTopkFullNoSelectorByteIdentical) {
|
||||
auto index = build_ivf_flat();
|
||||
auto* ivf = dynamic_cast<faiss::IndexIVFFlat*>(index.get());
|
||||
ASSERT_NE(ivf, nullptr);
|
||||
|
||||
auto xq = make_data(nq, 0xa5a5a5);
|
||||
|
||||
// Reference: ensure_topk_full=false.
|
||||
faiss::SearchParametersIVF p_ref;
|
||||
p_ref.nprobe = nlist;
|
||||
p_ref.max_codes = k / 2;
|
||||
p_ref.ensure_topk_full = false;
|
||||
std::vector<float> D_ref(k * nq);
|
||||
std::vector<faiss::idx_t> I_ref(k * nq);
|
||||
index->search(nq, xq.data(), k, D_ref.data(), I_ref.data(), &p_ref);
|
||||
|
||||
// Soft budget: ensure_topk_full=true.
|
||||
faiss::SearchParametersIVF p_full;
|
||||
p_full.nprobe = nlist;
|
||||
p_full.max_codes = k / 2;
|
||||
p_full.ensure_topk_full = true;
|
||||
std::vector<float> D_full(k * nq);
|
||||
std::vector<faiss::idx_t> I_full(k * nq);
|
||||
index->search(nq, xq.data(), k, D_full.data(), I_full.data(), &p_full);
|
||||
|
||||
// Softening the budget must not increase the number of missing results.
|
||||
size_t miss_ref = std::count(I_ref.begin(), I_ref.end(), (faiss::idx_t)-1);
|
||||
size_t miss_full =
|
||||
std::count(I_full.begin(), I_full.end(), (faiss::idx_t)-1);
|
||||
EXPECT_LE(miss_full, miss_ref);
|
||||
}
|
||||
|
||||
// Current FastScan handlers do not report heap update counts.
|
||||
TEST(IVFEarlyTermination, FastscanNheapUpdatesIsZeroToday) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
auto xq = make_data(nq, 0x5678);
|
||||
std::vector<float> D(k * nq);
|
||||
std::vector<faiss::idx_t> I(k * nq);
|
||||
|
||||
auto& stats = faiss::indexIVF_stats;
|
||||
stats.reset();
|
||||
index->search(nq, xq.data(), k, D.data(), I.data());
|
||||
|
||||
// FastScan plumbs nq, ndis, and nlist into indexIVF_stats.
|
||||
EXPECT_GT(stats.ndis, size_t(0));
|
||||
EXPECT_EQ(stats.nheap_updates, size_t(0))
|
||||
<< "FastScan started reporting nheap_updates — update this "
|
||||
"test and verify the value is correct.";
|
||||
}
|
||||
|
||||
// FastScan early-stop options use the per-query implementations.
|
||||
TEST(IVFEarlyTermination, FastscanDefaultsWork) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
auto xq = make_data(nq, 0x5678);
|
||||
std::vector<float> D(k * nq);
|
||||
std::vector<faiss::idx_t> I(k * nq);
|
||||
|
||||
faiss::SearchParametersIVF params; // defaults
|
||||
params.nprobe = 4;
|
||||
EXPECT_NO_THROW(
|
||||
index->search(nq, xq.data(), k, D.data(), I.data(), ¶ms));
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanMaxListsNumHonored) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
std::vector<float> D_full(k), D_cap(k);
|
||||
std::vector<faiss::idx_t> I_full(k), I_cap(k);
|
||||
|
||||
faiss::SearchParametersIVF p_full;
|
||||
p_full.nprobe = nlist;
|
||||
EXPECT_NO_THROW(index->search(
|
||||
1, xq.data(), k, D_full.data(), I_full.data(), &p_full));
|
||||
|
||||
faiss::SearchParametersIVF p_cap;
|
||||
p_cap.nprobe = nlist;
|
||||
p_cap.max_lists_num = 1;
|
||||
EXPECT_NO_THROW(
|
||||
index->search(1, xq.data(), k, D_cap.data(), I_cap.data(), &p_cap));
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanEnsureTopkFullSingleQueryWorks) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
std::vector<float> D(k);
|
||||
std::vector<faiss::idx_t> I(k);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = nlist;
|
||||
params.max_codes = 2;
|
||||
params.ensure_topk_full = true;
|
||||
EXPECT_NO_THROW(
|
||||
index->search(1, xq.data(), k, D.data(), I.data(), ¶ms));
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanEnsureTopkFullRestrictiveSelectorWorks) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
std::vector<uint8_t> packed((nb + 7) / 8, 0);
|
||||
for (size_t i = 0; i < nb; i += 10) {
|
||||
packed[i / 8] |= 1u << (i % 8);
|
||||
}
|
||||
faiss::IDSelectorBitmap sel(nb, packed.data());
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
std::vector<float> D(k);
|
||||
std::vector<faiss::idx_t> I(k);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = nlist;
|
||||
params.max_codes = 2;
|
||||
params.ensure_topk_full = true;
|
||||
params.sel = &sel;
|
||||
index->search(1, xq.data(), k, D.data(), I.data(), ¶ms);
|
||||
|
||||
EXPECT_EQ(std::count(I.begin(), I.end(), (faiss::idx_t)-1), size_t(0))
|
||||
<< "FastScan max_codes must be counted after IDSelector filtering";
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanEnsureTopkFullMultiQueryWorks) {
|
||||
// Each query independently resets nscan_q; multi-query batches are
|
||||
// therefore supported.
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
auto xq = make_data(2, 0x9999);
|
||||
std::vector<float> D(k * 2);
|
||||
std::vector<faiss::idx_t> I(k * 2);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.ensure_topk_full = true;
|
||||
EXPECT_NO_THROW(
|
||||
index->search(2, xq.data(), k, D.data(), I.data(), ¶ms));
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanExplicitImplem12RejectsKnobs) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
auto* fs = dynamic_cast<faiss::IndexIVFFastScan*>(index.get());
|
||||
ASSERT_NE(fs, nullptr);
|
||||
fs->implem = 12;
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
std::vector<float> D(k);
|
||||
std::vector<faiss::idx_t> I(k);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.max_lists_num = 2;
|
||||
EXPECT_THROW(
|
||||
index->search(1, xq.data(), k, D.data(), I.data(), ¶ms),
|
||||
faiss::FaissException);
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanRangeDefaultsPreserveBaseline) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
auto* fs = dynamic_cast<faiss::IndexIVFFastScan*>(index.get());
|
||||
ASSERT_NE(fs, nullptr);
|
||||
fs->nprobe = 4;
|
||||
|
||||
auto xq = make_data(nq, 0x9999);
|
||||
const float radius = 1.0f;
|
||||
faiss::RangeSearchResult ref(nq);
|
||||
faiss::RangeSearchResult tst(nq);
|
||||
|
||||
index->range_search(nq, xq.data(), radius, &ref);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.max_empty_result_buckets = 0;
|
||||
index->range_search(nq, xq.data(), radius, &tst, ¶ms);
|
||||
|
||||
ASSERT_EQ(ref.nq, tst.nq);
|
||||
for (size_t i = 0; i <= ref.nq; i++) {
|
||||
EXPECT_EQ(ref.lims[i], tst.lims[i]) << "at lims[" << i << "]";
|
||||
}
|
||||
for (size_t q = 0; q < ref.nq; q++) {
|
||||
std::vector<std::pair<faiss::idx_t, float>> ref_pairs;
|
||||
std::vector<std::pair<faiss::idx_t, float>> tst_pairs;
|
||||
for (size_t j = ref.lims[q]; j < ref.lims[q + 1]; j++) {
|
||||
ref_pairs.emplace_back(ref.labels[j], ref.distances[j]);
|
||||
}
|
||||
for (size_t j = tst.lims[q]; j < tst.lims[q + 1]; j++) {
|
||||
tst_pairs.emplace_back(tst.labels[j], tst.distances[j]);
|
||||
}
|
||||
std::sort(ref_pairs.begin(), ref_pairs.end());
|
||||
std::sort(tst_pairs.begin(), tst_pairs.end());
|
||||
EXPECT_EQ(ref_pairs, tst_pairs);
|
||||
}
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanRangeHonorsMaxEmptyResultBuckets) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
const float radius = 1.0f;
|
||||
faiss::RangeSearchResult full(1);
|
||||
faiss::RangeSearchResult early(1);
|
||||
std::vector<faiss::idx_t> empty_subset;
|
||||
faiss::IDSelectorBatch empty_sel(0, empty_subset.data());
|
||||
|
||||
auto& stats = faiss::indexIVF_stats;
|
||||
faiss::SearchParametersIVF p_full;
|
||||
p_full.nprobe = nlist;
|
||||
p_full.max_empty_result_buckets = 0;
|
||||
p_full.sel = &empty_sel;
|
||||
stats.reset();
|
||||
index->range_search(1, xq.data(), radius, &full, &p_full);
|
||||
const size_t full_nlist = stats.nlist;
|
||||
|
||||
faiss::SearchParametersIVF p_early;
|
||||
p_early.nprobe = nlist;
|
||||
p_early.max_empty_result_buckets = 1;
|
||||
p_early.sel = &empty_sel;
|
||||
stats.reset();
|
||||
index->range_search(1, xq.data(), radius, &early, &p_early);
|
||||
const size_t early_nlist = stats.nlist;
|
||||
|
||||
EXPECT_GT(full_nlist, size_t(1))
|
||||
<< "test precondition: full scan should visit several lists";
|
||||
EXPECT_LT(early_nlist, full_nlist)
|
||||
<< "max_empty_result_buckets did not stop fastscan range early";
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanRangeRejectsMaxEmptyWithExplicitImpl12) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
auto* fs = dynamic_cast<faiss::IndexIVFFastScan*>(index.get());
|
||||
ASSERT_NE(fs, nullptr);
|
||||
fs->implem = 12;
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
faiss::RangeSearchResult res(1);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.max_empty_result_buckets = 2;
|
||||
EXPECT_THROW(
|
||||
index->range_search(1, xq.data(), 1.0f, &res, ¶ms),
|
||||
faiss::FaissException);
|
||||
}
|
||||
|
||||
TEST(IVFEarlyTermination, FastscanRangeRejectsMaxEmptyUnsupportedParallelMode) {
|
||||
auto xb = make_data(nb, 0xcafebabe);
|
||||
std::unique_ptr<faiss::Index> index(
|
||||
faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2));
|
||||
index->train(nb, xb.data());
|
||||
index->add(nb, xb.data());
|
||||
auto* fs = dynamic_cast<faiss::IndexIVFFastScan*>(index.get());
|
||||
ASSERT_NE(fs, nullptr);
|
||||
fs->parallel_mode = 1;
|
||||
|
||||
auto xq = make_data(1, 0x9999);
|
||||
faiss::RangeSearchResult res(1);
|
||||
|
||||
faiss::SearchParametersIVF params;
|
||||
params.nprobe = 4;
|
||||
params.max_empty_result_buckets = 2;
|
||||
EXPECT_THROW(
|
||||
index->range_search(1, xq.data(), 1.0f, &res, ¶ms),
|
||||
faiss::FaissException);
|
||||
}
|
||||
@@ -118,6 +118,73 @@ int test_params_override(const char* index_key, MetricType metric) {
|
||||
* Test subsets
|
||||
*************************************************************/
|
||||
|
||||
/*************************************************************
|
||||
* Test ensure_topk_full behavior
|
||||
*************************************************************/
|
||||
|
||||
// With a tight max_codes budget the default search (ensure_topk_full=false)
|
||||
// should truncate results, yielding -1 placeholders. With ensure_topk_full=true
|
||||
// the same budget becomes soft: the probe loop keeps running until k raw
|
||||
// candidates have been scanned, so placeholders disappear (data permitting).
|
||||
int test_ensure_topk_full(const char* index_key, MetricType metric) {
|
||||
std::vector<float> xb = make_data(nb);
|
||||
auto index = make_index(index_key, metric, xb);
|
||||
std::vector<float> xq = make_data(nq);
|
||||
|
||||
constexpr int k = 10;
|
||||
// Budget: smaller than k codes per query. Guaranteed truncation when
|
||||
// ensure_topk_full is false.
|
||||
const size_t tight_budget = k / 2;
|
||||
|
||||
IVFSearchParameters p;
|
||||
p.nprobe = 32;
|
||||
p.max_codes = tight_budget;
|
||||
p.ensure_topk_full = false;
|
||||
auto res_tight = search_index_with_params(index.get(), xq.data(), &p);
|
||||
|
||||
p.ensure_topk_full = true;
|
||||
auto res_full = search_index_with_params(index.get(), xq.data(), &p);
|
||||
|
||||
size_t miss_tight =
|
||||
std::count(res_tight.begin(), res_tight.end(), (idx_t)-1);
|
||||
size_t miss_full = std::count(res_full.begin(), res_full.end(), (idx_t)-1);
|
||||
|
||||
// ensure_topk_full=true must produce at least as many valid slots as
|
||||
// ensure_topk_full=false; in practice it must strictly reduce -1 count
|
||||
// given our tight budget.
|
||||
if (miss_full > miss_tight) {
|
||||
return 1;
|
||||
}
|
||||
if (miss_tight == 0) {
|
||||
// Test assumption violated: budget should have caused truncation.
|
||||
return 2;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Regression test: when all three new fields are left at defaults the
|
||||
// search behaves identically to a search with no params at all. This is the
|
||||
// backward-compatibility guarantee for pre-existing baseline callers.
|
||||
int test_defaults_preserve_baseline(const char* index_key, MetricType metric) {
|
||||
std::vector<float> xb = make_data(nb);
|
||||
auto index = make_index(index_key, metric, xb);
|
||||
std::vector<float> xq = make_data(nq);
|
||||
|
||||
ParameterSpace ps;
|
||||
ps.set_index_parameter(index.get(), "nprobe", 4);
|
||||
auto res_no_params = search_index(index.get(), xq.data());
|
||||
|
||||
IVFSearchParameters p; // default-constructed: all new fields at 0/false
|
||||
p.nprobe = 4;
|
||||
auto res_default_params =
|
||||
search_index_with_params(index.get(), xq.data(), &p);
|
||||
|
||||
if (res_no_params != res_default_params) {
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int test_selector(const char* index_key) {
|
||||
std::vector<float> xb = make_data(nb); // database vectors
|
||||
std::vector<float> xq = make_data(nq);
|
||||
@@ -192,6 +259,33 @@ TEST(TPO, IVFFlatPP) {
|
||||
EXPECT_EQ(err2, 0);
|
||||
}
|
||||
|
||||
TEST(TPOEnsureTopK, IVFFlat) {
|
||||
EXPECT_EQ(test_ensure_topk_full("IVF32,Flat", METRIC_L2), 0);
|
||||
EXPECT_EQ(test_ensure_topk_full("IVF32,Flat", METRIC_INNER_PRODUCT), 0);
|
||||
}
|
||||
|
||||
TEST(TPOEnsureTopK, IVFPQ) {
|
||||
EXPECT_EQ(test_ensure_topk_full("IVF32,PQ8np", METRIC_L2), 0);
|
||||
EXPECT_EQ(test_ensure_topk_full("IVF32,PQ8np", METRIC_INNER_PRODUCT), 0);
|
||||
}
|
||||
|
||||
TEST(TPOEnsureTopK, IVFSQ) {
|
||||
EXPECT_EQ(test_ensure_topk_full("IVF32,SQ8", METRIC_L2), 0);
|
||||
EXPECT_EQ(test_ensure_topk_full("IVF32,SQ8", METRIC_INNER_PRODUCT), 0);
|
||||
}
|
||||
|
||||
TEST(TPODefaults, IVFFlat) {
|
||||
EXPECT_EQ(test_defaults_preserve_baseline("IVF32,Flat", METRIC_L2), 0);
|
||||
}
|
||||
|
||||
TEST(TPODefaults, IVFPQ) {
|
||||
EXPECT_EQ(test_defaults_preserve_baseline("IVF32,PQ8np", METRIC_L2), 0);
|
||||
}
|
||||
|
||||
TEST(TPODefaults, IVFSQ) {
|
||||
EXPECT_EQ(test_defaults_preserve_baseline("IVF32,SQ8", METRIC_L2), 0);
|
||||
}
|
||||
|
||||
TEST(TSEL, IVFFlat) {
|
||||
int err = test_selector("PCA16,IVF32,Flat");
|
||||
EXPECT_EQ(err, 0);
|
||||
|
||||
@@ -581,3 +581,197 @@ class TestPrecomputed(unittest.TestCase):
|
||||
|
||||
def test_knn_and_range_FS(self):
|
||||
self.do_test_knn_and_range("IVF32,PQ8x4fs", range=False)
|
||||
|
||||
|
||||
class TestIVFEarlyTermination(unittest.TestCase):
|
||||
"""
|
||||
Coverage for the new SearchParametersIVF early-stop fields:
|
||||
- ensure_topk_full
|
||||
- max_empty_result_buckets
|
||||
- max_lists_num
|
||||
"""
|
||||
|
||||
def _build(self, factory, mt=faiss.METRIC_L2, d=32):
|
||||
ds = datasets.SyntheticDataset(d, 2000, 500, 30)
|
||||
index = faiss.index_factory(d, factory, mt)
|
||||
index.train(ds.get_train())
|
||||
index.add(ds.get_database())
|
||||
return ds, index
|
||||
|
||||
def test_ensure_topk_full_fills_heap(self):
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
k = 10
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 32
|
||||
params.max_codes = k // 2 # tight budget: forces truncation
|
||||
params.ensure_topk_full = False
|
||||
_, I_tight = index.search(ds.get_queries(), k, params=params)
|
||||
|
||||
params.ensure_topk_full = True
|
||||
_, I_full = index.search(ds.get_queries(), k, params=params)
|
||||
|
||||
# ensure_topk_full=true must produce at least as many valid slots.
|
||||
miss_tight = np.sum(I_tight == -1)
|
||||
miss_full = np.sum(I_full == -1)
|
||||
self.assertLessEqual(miss_full, miss_tight)
|
||||
# Test precondition: tight budget produced some misses.
|
||||
self.assertGreater(miss_tight, 0)
|
||||
|
||||
def test_range_default_matches_baseline(self):
|
||||
"""Default max_empty_result_buckets=0 must preserve pre-existing
|
||||
range_search results byte-for-byte."""
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
# Set nprobe via the index attribute for the no-params baseline.
|
||||
index.nprobe = 8
|
||||
radius = 2.0
|
||||
Lref, Dref, Iref = index.range_search(ds.get_queries(), radius)
|
||||
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 8
|
||||
# all new fields are at defaults
|
||||
Lnew, Dnew, Inew = index.range_search(
|
||||
ds.get_queries(), radius, params=params
|
||||
)
|
||||
|
||||
np.testing.assert_array_equal(Lref, Lnew)
|
||||
# Per-query (id, distance) pairs must match as sorted sets since
|
||||
# within-query order is not guaranteed across OMP threads.
|
||||
for q in range(ds.nq):
|
||||
a = sorted(zip(
|
||||
Iref[Lref[q]: Lref[q + 1]].tolist(),
|
||||
Dref[Lref[q]: Lref[q + 1]].tolist(),
|
||||
))
|
||||
b = sorted(zip(
|
||||
Inew[Lnew[q]: Lnew[q + 1]].tolist(),
|
||||
Dnew[Lnew[q]: Lnew[q + 1]].tolist(),
|
||||
))
|
||||
self.assertEqual(a, b)
|
||||
|
||||
def test_range_early_exit_is_subset(self):
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
# Radius chosen empirically against this synthetic dataset
|
||||
# (d=32, nb=500, IVF32) so that several probed buckets per
|
||||
# query yield zero in-radius hits — required to exercise
|
||||
# the empty-bucket exit path. Too-tight or too-wide radii
|
||||
# leave the test vacuous.
|
||||
radius = 15.0
|
||||
p_full = faiss.SearchParametersIVF()
|
||||
p_full.nprobe = 32
|
||||
p_full.max_empty_result_buckets = 0 # disabled
|
||||
Lf, Df, If = index.range_search(
|
||||
ds.get_queries(), radius, params=p_full
|
||||
)
|
||||
|
||||
p_early = faiss.SearchParametersIVF()
|
||||
p_early.nprobe = 32
|
||||
p_early.max_empty_result_buckets = 1 # aggressive early-exit
|
||||
Le, De, Ie = index.range_search(
|
||||
ds.get_queries(), radius, params=p_early
|
||||
)
|
||||
|
||||
# Early-exit must never expand the result set.
|
||||
self.assertTrue(np.all(np.diff(Le) <= np.diff(Lf)))
|
||||
# And each per-query id-set must be a subset of the full result.
|
||||
for q in range(ds.nq):
|
||||
full = set(If[Lf[q]: Lf[q + 1]].tolist())
|
||||
early = set(Ie[Le[q]: Le[q + 1]].tolist())
|
||||
self.assertTrue(early.issubset(full))
|
||||
# At least one query must strictly shrink.
|
||||
any_shrunk = bool(np.any(np.diff(Le) < np.diff(Lf)))
|
||||
self.assertTrue(
|
||||
any_shrunk,
|
||||
"no query shrank: max_empty_result_buckets has no effect",
|
||||
)
|
||||
|
||||
def test_range_early_exit_counter_resets_on_hit(self):
|
||||
"""Tightening max_empty_result_buckets (1 vs 3) must never
|
||||
produce *more* results. With 3 the loop tolerates two
|
||||
intervening empty buckets between non-empty ones; with 1 it
|
||||
terminates on the first empty bucket. Per-query result
|
||||
counts must therefore satisfy tight ≤ relaxed."""
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
radius = 15.0
|
||||
p_tight = faiss.SearchParametersIVF()
|
||||
p_tight.nprobe = 32
|
||||
p_tight.max_empty_result_buckets = 1
|
||||
Lt, _, _ = index.range_search(
|
||||
ds.get_queries(), radius, params=p_tight
|
||||
)
|
||||
|
||||
p_relaxed = faiss.SearchParametersIVF()
|
||||
p_relaxed.nprobe = 32
|
||||
p_relaxed.max_empty_result_buckets = 3
|
||||
Lr, _, _ = index.range_search(
|
||||
ds.get_queries(), radius, params=p_relaxed
|
||||
)
|
||||
|
||||
self.assertTrue(np.all(np.diff(Lt) <= np.diff(Lr)))
|
||||
|
||||
def test_range_rejects_non_default_parallel_mode(self):
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
ivf = faiss.downcast_index(index)
|
||||
ivf.parallel_mode = 1
|
||||
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 4
|
||||
params.max_empty_result_buckets = 2
|
||||
with self.assertRaises(RuntimeError):
|
||||
index.range_search(ds.get_queries()[:1], 1.0, params=params)
|
||||
|
||||
def test_ndis_stats_post_filter(self):
|
||||
"""IndexIVFStats::ndis is documented as 'nb of distances
|
||||
computed' — a post-filter quantity. With a 50%-keep selector,
|
||||
ndis must drop because rejected codes don't have a distance
|
||||
computed."""
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
index.nprobe = 32
|
||||
stats = faiss.cvar.indexIVF_stats
|
||||
|
||||
stats.reset()
|
||||
index.search(ds.get_queries(), 10)
|
||||
ndis_no_sel = stats.ndis
|
||||
self.assertGreater(ndis_no_sel, 0)
|
||||
|
||||
# 50%-keep batch selector.
|
||||
keep = np.arange(0, ds.nb, 2).astype("int64")
|
||||
sel = faiss.IDSelectorBatch(keep)
|
||||
params = faiss.SearchParametersIVF(nprobe=32, sel=sel)
|
||||
stats.reset()
|
||||
index.search(ds.get_queries(), 10, params=params)
|
||||
ndis_with_sel = stats.ndis
|
||||
|
||||
self.assertLess(ndis_with_sel, ndis_no_sel)
|
||||
# Loose-bound sanity: ~50% retention should drop ndis at least
|
||||
# noticeably below the no-selector value.
|
||||
self.assertLess(ndis_with_sel * 4, ndis_no_sel * 3)
|
||||
|
||||
def test_ensure_topk_full_with_restrictive_selector(self):
|
||||
"""Under a 10%-keep selector and
|
||||
max_codes < k, the heap must still be filled (no -1 entries)
|
||||
because the post-filter break condition keeps probing until k
|
||||
valid candidates have been observed."""
|
||||
ds, index = self._build("IVF32,Flat")
|
||||
# 10%-keep selector — 50 valid ids in 500-vector dataset.
|
||||
keep = np.arange(0, ds.nb, 10).astype("int64")
|
||||
sel = faiss.IDSelectorBatch(keep)
|
||||
|
||||
params = faiss.SearchParametersIVF()
|
||||
params.nprobe = 32
|
||||
params.max_codes = 5 # tight pre-filter budget
|
||||
params.ensure_topk_full = True # post-filter soft cap
|
||||
params.sel = sel
|
||||
_, I = index.search(ds.get_queries(), 10, params=params)
|
||||
|
||||
self.assertEqual(int(np.sum(I == -1)), 0)
|
||||
|
||||
def test_fields_exposed_and_defaulted(self):
|
||||
p = faiss.SearchParametersIVF()
|
||||
self.assertEqual(p.max_lists_num, 0)
|
||||
self.assertIs(p.ensure_topk_full, False)
|
||||
self.assertEqual(p.max_empty_result_buckets, 0)
|
||||
p.max_lists_num = 5
|
||||
p.ensure_topk_full = True
|
||||
p.max_empty_result_buckets = 2
|
||||
self.assertEqual(p.max_lists_num, 5)
|
||||
self.assertIs(p.ensure_topk_full, True)
|
||||
self.assertEqual(p.max_empty_result_buckets, 2)
|
||||
|
||||
Reference in New Issue
Block a user