From 799bf3aaeca910113fd2f1b0287568f07eb002fa Mon Sep 17 00:00:00 2001 From: Alexandr Guzhva Date: Fri, 1 May 2026 12:23:05 -0700 Subject: [PATCH] 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 --- c_api/CMakeLists.txt | 4 + c_api/IndexIVF_c.cpp | 9 + c_api/IndexIVF_c.h | 6 + c_api/example_search_params_ivf_c.c | 60 ++ faiss/IndexIVF.cpp | 85 +- faiss/IndexIVF.h | 21 +- faiss/IndexIVFFastScan.cpp | 134 +++- faiss/IndexIVFFastScan.h | 3 +- faiss/IndexIVFFlatPanorama.cpp | 11 +- faiss/IndexIVFPQ.cpp | 9 +- faiss/IndexIVFPQFastScan.cpp | 123 +-- faiss/IndexIVFRaBitQ.cpp | 4 + faiss/IndexIVFRaBitQFastScan.cpp | 123 +-- faiss/impl/AuxIndexStructures.h | 3 + faiss/impl/InvertedListScannerStats.h | 26 + faiss/impl/ResultHandler.h | 4 + .../IndexIVFSpectralHash_impl.h | 18 +- faiss/impl/expanded_scanners.h | 11 +- faiss/python/__init__.pyi | 5 + faiss/python/swigfaiss.swig | 2 + tests/CMakeLists.txt | 1 + tests/test_fast_scan_ivf.py | 159 ++++ tests/test_ivf_early_termination.cpp | 755 ++++++++++++++++++ tests/test_params_override.cpp | 94 +++ tests/test_search_params.py | 194 +++++ 25 files changed, 1730 insertions(+), 134 deletions(-) create mode 100644 c_api/example_search_params_ivf_c.c create mode 100644 faiss/impl/InvertedListScannerStats.h create mode 100644 tests/test_ivf_early_termination.cpp diff --git a/c_api/CMakeLists.txt b/c_api/CMakeLists.txt index 67a945d52..5f2de5140 100644 --- a/c_api/CMakeLists.txt +++ b/c_api/CMakeLists.txt @@ -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() diff --git a/c_api/IndexIVF_c.cpp b/c_api/IndexIVF_c.cpp index 0e42e0f76..0cab7ea01 100644 --- a/c_api/IndexIVF_c.cpp +++ b/c_api/IndexIVF_c.cpp @@ -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) diff --git a/c_api/IndexIVF_c.h b/c_api/IndexIVF_c.h index da2f2e357..288daa1a6 100644 --- a/c_api/IndexIVF_c.h +++ b/c_api/IndexIVF_c.h @@ -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) * diff --git a/c_api/example_search_params_ivf_c.c b/c_api/example_search_params_ivf_c.c new file mode 100644 index 000000000..148e2df7c --- /dev/null +++ b/c_api/example_search_params_ivf_c.c @@ -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 +#include +#include + +#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; +} diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index 1cdde3ac6..52aae8a3f 100644 --- a/faiss/IndexIVF.cpp +++ b/faiss/IndexIVF.cpp @@ -420,6 +420,9 @@ void IndexIVF::search_preassigned( const idx_t unlimited_list_size = std::numeric_limits::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(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 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 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 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 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; RangeResultHandler 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; RangeResultHandler 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 diff --git a/faiss/IndexIVF.h b/faiss/IndexIVF.h index d3a4bb891..a58ecf765 100644 --- a/faiss/IndexIVF.h +++ b/faiss/IndexIVF.h @@ -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, diff --git a/faiss/IndexIVFFastScan.cpp b/faiss/IndexIVFFastScan.cpp index 26fd54caf..a4ef6edc3 100644 --- a/faiss/IndexIVFFastScan.cpp +++ b/faiss/IndexIVFFastScan.cpp @@ -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(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(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 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::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 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(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; + } + } } } diff --git a/faiss/IndexIVFFastScan.h b/faiss/IndexIVFFastScan.h index fa7662852..7f052f7d6 100644 --- a/faiss/IndexIVFFastScan.h +++ b/faiss/IndexIVFFastScan.h @@ -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, diff --git a/faiss/IndexIVFFlatPanorama.cpp b/faiss/IndexIVFFlatPanorama.cpp index ec57162f4..313a057ee 100644 --- a/faiss/IndexIVFFlatPanorama.cpp +++ b/faiss/IndexIVFFlatPanorama.cpp @@ -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++; + } } } } diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp index 922ac4df2..157e161da 100644 --- a/faiss/IndexIVFPQ.cpp +++ b/faiss/IndexIVFPQ.cpp @@ -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++; + } } } }; diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp index bcfdf41c1..84a9b15b3 100644 --- a/faiss/IndexIVFPQFastScan.cpp +++ b/faiss/IndexIVFPQFastScan.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -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 curr_dists(k, distances[0]); - std::vector 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 curr_dists(k, handler.threshold); + std::vector 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(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(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; + if (auto* heap_handler = + dynamic_cast*>(&handler)) { + return scan_with_heap(heap_handler); + } } else { - maxheap_addn( - k, - distances, - labels, - curr_dists.data(), - curr_labels.data(), - k); + using C = CMin; + if (auto* heap_handler = + dynamic_cast*>(&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"); } }; diff --git a/faiss/IndexIVFRaBitQ.cpp b/faiss/IndexIVFRaBitQ.cpp index d651f66c6..43cb17bdb 100644 --- a/faiss/IndexIVFRaBitQ.cpp +++ b/faiss/IndexIVFRaBitQ.cpp @@ -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++; } } diff --git a/faiss/IndexIVFRaBitQFastScan.cpp b/faiss/IndexIVFRaBitQFastScan.cpp index ae1574c83..02bbdb484 100644 --- a/faiss/IndexIVFRaBitQFastScan.cpp +++ b/faiss/IndexIVFRaBitQFastScan.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -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 curr_dists(k, distances[0]); - std::vector 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 curr_dists(k, result_handler.threshold); + std::vector 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(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(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; + if (auto* heap_handler = dynamic_cast*>( + &result_handler)) { + return scan_with_heap(heap_handler); + } } else { - maxheap_addn( - k, - distances, - labels, - curr_dists.data(), - curr_labels.data(), - k); + using C = CMin; + if (auto* heap_handler = dynamic_cast*>( + &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"); } }; diff --git a/faiss/impl/AuxIndexStructures.h b/faiss/impl/AuxIndexStructures.h index 7533d206a..58a80997a 100644 --- a/faiss/impl/AuxIndexStructures.h +++ b/faiss/impl/AuxIndexStructures.h @@ -18,6 +18,8 @@ #include #include +#include + #include #include @@ -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); diff --git a/faiss/impl/InvertedListScannerStats.h b/faiss/impl/InvertedListScannerStats.h new file mode 100644 index 000000000..57be7b7b9 --- /dev/null +++ b/faiss/impl/InvertedListScannerStats.h @@ -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 + +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 diff --git a/faiss/impl/ResultHandler.h b/faiss/impl/ResultHandler.h index 44b692ef8..7387cf7f8 100644 --- a/faiss/impl/ResultHandler.h +++ b/faiss/impl/ResultHandler.h @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -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; diff --git a/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h b/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h index e84728c0b..828752aa2 100644 --- a/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h +++ b/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h @@ -26,6 +26,7 @@ #include #include #include +#include #include #include @@ -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; } diff --git a/faiss/impl/expanded_scanners.h b/faiss/impl/expanded_scanners.h index 6e1253e37..8eab6ab2a 100644 --- a/faiss/impl/expanded_scanners.h +++ b/faiss/impl/expanded_scanners.h @@ -10,6 +10,7 @@ #include #include +#include #include /* 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; } diff --git a/faiss/python/__init__.pyi b/faiss/python/__init__.pyi index a6dc8044a..99d8d4f5d 100644 --- a/faiss/python/__init__.pyi +++ b/faiss/python/__init__.pyi @@ -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: ... diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig index 42d6b25a2..b3577284f 100644 --- a/faiss/python/swigfaiss.swig +++ b/faiss/python/swigfaiss.swig @@ -93,6 +93,7 @@ typedef uint64_t size_t; #include #include #include +#include #include #include #include @@ -547,6 +548,7 @@ void gpu_sync_all_devices() %ignore BlockInvertedListsIOHook; %include %include +%include %include // NOTE(hoss): SWIG (wrongly) believes the overloaded const version shadows the // non-const one. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 3a27b981e..24f9df0e4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -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 diff --git a/tests/test_fast_scan_ivf.py b/tests/test_fast_scan_ivf.py index 7e2d74737..f55e61bf9 100644 --- a/tests/test_fast_scan_ivf.py +++ b/tests/test_fast_scan_ivf.py @@ -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) diff --git a/tests/test_ivf_early_termination.cpp b/tests/test_ivf_early_termination.cpp new file mode 100644 index 000000000..762fe52be --- /dev/null +++ b/tests/test_ivf_early_termination.cpp @@ -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 +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +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 make_data(size_t n, uint32_t seed) { + std::mt19937 rng(seed); + std::uniform_real_distribution dist(0.0f, 1.0f); + std::vector out(n * d); + for (auto& v : out) { + v = dist(rng); + } + return out; +} + +std::unique_ptr build_ivf_flat() { + auto xb = make_data(nb, 0xabcdef); + std::unique_ptr 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(index.get()); + ASSERT_NE(ivf, nullptr); + ivf->nprobe = 4; + + auto xq = make_data(nq, 0x123456); + std::vector D_ref(k * nq), D_new(k * nq); + std::vector 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(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> ref_pairs; + std::vector> 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(index.get()); + ASSERT_NE(ivf, nullptr); + + auto xq = make_data(nq, 0xdeadbe); + std::vector D_tight(k * nq), D_full(k * nq); + std::vector 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(index.get()); + ASSERT_NE(ivf, nullptr); + ivf->parallel_mode = 1; + + auto xq = make_data(1, 0xaaaa); + std::vector D(k); + std::vector 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(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 full_ids( + full.labels + full.lims[q], full.labels + full.lims[q + 1]); + std::set 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(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(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(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(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 D(k * nq); + std::vector 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 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(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 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(index.get()); + ASSERT_NE(ivf, nullptr); + ivf->nprobe = nlist; + + // ~10%-keep bitmap selector (every 10th id). + std::vector 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 D(k * nq); + std::vector 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(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 D_ref(k * nq); + std::vector 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 D_full(k * nq); + std::vector 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 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 D(k * nq); + std::vector 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 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 D(k * nq); + std::vector 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 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 D_full(k), D_cap(k); + std::vector 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 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 D(k); + std::vector 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 index( + faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2)); + index->train(nb, xb.data()); + index->add(nb, xb.data()); + + std::vector 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 D(k); + std::vector 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 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 D(k * 2); + std::vector 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 index( + faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2)); + index->train(nb, xb.data()); + index->add(nb, xb.data()); + auto* fs = dynamic_cast(index.get()); + ASSERT_NE(fs, nullptr); + fs->implem = 12; + + auto xq = make_data(1, 0x9999); + std::vector D(k); + std::vector 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 index( + faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2)); + index->train(nb, xb.data()); + index->add(nb, xb.data()); + auto* fs = dynamic_cast(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> ref_pairs; + std::vector> 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 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 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 index( + faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2)); + index->train(nb, xb.data()); + index->add(nb, xb.data()); + auto* fs = dynamic_cast(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 index( + faiss::index_factory(d, "IVF16,PQ4x4fs", faiss::METRIC_L2)); + index->train(nb, xb.data()); + index->add(nb, xb.data()); + auto* fs = dynamic_cast(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); +} diff --git a/tests/test_params_override.cpp b/tests/test_params_override.cpp index 182cb7ccd..15d5fef09 100644 --- a/tests/test_params_override.cpp +++ b/tests/test_params_override.cpp @@ -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 xb = make_data(nb); + auto index = make_index(index_key, metric, xb); + std::vector 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 xb = make_data(nb); + auto index = make_index(index_key, metric, xb); + std::vector 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 xb = make_data(nb); // database vectors std::vector 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); diff --git a/tests/test_search_params.py b/tests/test_search_params.py index 246ccaa71..00e89f925 100644 --- a/tests/test_search_params.py +++ b/tests/test_search_params.py @@ -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)