From 40a8cc003460855a810b2063191eaac0382bb85e Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Mon, 4 May 2026 10:25:15 -0700 Subject: [PATCH] Fix lints associated with early stop facilities for IVF (#5172) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5172 Follow up to D103044128 Fix some lint issues and one stats-tracking issue (and add test for it) in IVF early-stop code: - **ResultHandler.h**: `RangeResultHandler::add_result` now returns `true` when a result is actually added, fixing `nheap_updates` stat tracking through the `scan_codes_range` / `expanded_scanners.h` path. Previously always returned `false`, making `nheap_updates` a perpetual no-op for range search via the default scanner. This is inconsistent with `iterate_codes_range` and `IndexIVFSpectralHash_impl.h`, which track correctly. - **IndexIVF.cpp**: Move `effective_max_codes` declaration to its only use site, eliminating a dead-store lint warning. - **IndexIVFPQFastScan.cpp**: Remove unused `#include `. - **test_fast_scan_ivf.py**: Reformat trailing comment for readability. - **test_ivf_early_termination.cpp**: Add unit test verifying `RangeResultHandler::add_result` return value. Register test in BUCK (was CMake-only). Reviewed By: mdouze Differential Revision: D103438123 fbshipit-source-id: b8878975e1f280ff1590363541a100d3ae51ba0e --- faiss/IndexIVF.cpp | 3 +- faiss/IndexIVFPQFastScan.cpp | 1 - faiss/impl/ResultHandler.h | 1 + tests/test_fast_scan_ivf.py | 4 +- tests/test_ivf_early_termination.cpp | 73 ++++++++++++++++++++++++++++ 5 files changed, 77 insertions(+), 5 deletions(-) diff --git a/faiss/IndexIVF.cpp b/faiss/IndexIVF.cpp index 52aae8a3f..f4081c00a 100644 --- a/faiss/IndexIVF.cpp +++ b/faiss/IndexIVF.cpp @@ -421,7 +421,6 @@ 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); @@ -466,7 +465,7 @@ void IndexIVF::search_preassigned( } // 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 = + idx_t 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 && diff --git a/faiss/IndexIVFPQFastScan.cpp b/faiss/IndexIVFPQFastScan.cpp index 84a9b15b3..d3f140300 100644 --- a/faiss/IndexIVFPQFastScan.cpp +++ b/faiss/IndexIVFPQFastScan.cpp @@ -16,7 +16,6 @@ #include #include #include -#include #include #include #include diff --git a/faiss/impl/ResultHandler.h b/faiss/impl/ResultHandler.h index 7387cf7f8..93e34d5c5 100644 --- a/faiss/impl/ResultHandler.h +++ b/faiss/impl/ResultHandler.h @@ -596,6 +596,7 @@ struct RangeResultHandler : ResultHandlerT { bool add_result(T dis, TI idx) final { if (C::cmp(threshold, dis)) { qr->add(dis, idx); + return true; } return false; } diff --git a/tests/test_fast_scan_ivf.py b/tests/test_fast_scan_ivf.py index f55e61bf9..836d088f2 100644 --- a/tests/test_fast_scan_ivf.py +++ b/tests/test_fast_scan_ivf.py @@ -1074,8 +1074,8 @@ class TestFastScanEarlyTerminationKnobs(unittest.TestCase): k = 10 params = faiss.SearchParametersIVF() params.nprobe = 16 - params.max_codes = k // 2 # tight: forces truncation without - # ensure_topk_full + # k // 2 max_codes forces truncation without ensure_topk_full + params.max_codes = k // 2 params.ensure_topk_full = False _, I_tight = index.search(q, k, params=params) params.ensure_topk_full = True diff --git a/tests/test_ivf_early_termination.cpp b/tests/test_ivf_early_termination.cpp index 762fe52be..51f50869c 100644 --- a/tests/test_ivf_early_termination.cpp +++ b/tests/test_ivf_early_termination.cpp @@ -28,6 +28,7 @@ #include #include #include +#include #include namespace { @@ -504,6 +505,78 @@ TEST(IVFEarlyTermination, FastscanNheapUpdatesIsZeroToday) { "test and verify the value is correct."; } +// RangeResultHandler::add_result must return true when a result is added +// (distance passes the threshold) so that callers in expanded_scanners.h +// correctly track nheap_updates. +TEST(IVFEarlyTermination, RangeResultHandlerAddResultReturnValue) { + using C = faiss::CMax; + + faiss::RangeSearchResult rsr(1); + faiss::RangeSearchPartialResult pres(&rsr); + auto& qr = pres.new_result(0); + + faiss::RangeResultHandler handler(&qr, /*threshold=*/5.0f); + + // CMax::cmp(5.0, 3.0) is true: distance within radius, should add + EXPECT_TRUE(handler.add_result(3.0f, 42)); + EXPECT_EQ(qr.nres, size_t(1)); + + // CMax::cmp(5.0, 5.0) is false: at boundary, should not add + EXPECT_FALSE(handler.add_result(5.0f, 43)); + EXPECT_EQ(qr.nres, size_t(1)); + + // CMax::cmp(5.0, 10.0) is false: distance exceeds radius + EXPECT_FALSE(handler.add_result(10.0f, 44)); + EXPECT_EQ(qr.nres, size_t(1)); + + // Second valid result + EXPECT_TRUE(handler.add_result(1.0f, 45)); + EXPECT_EQ(qr.nres, size_t(2)); + + // Finalize and verify stored (dis, id) pairs + pres.finalize(); + ASSERT_EQ(rsr.lims[0], size_t(0)); + ASSERT_EQ(rsr.lims[1], size_t(2)); + std::set ids(rsr.labels, rsr.labels + 2); + EXPECT_EQ(ids, (std::set{42, 45})); + std::set dists(rsr.distances, rsr.distances + 2); + EXPECT_EQ(dists, (std::set{1.0f, 3.0f})); +} + +// Same test with CMin (inner product semantics): adds when threshold < dis. +TEST(IVFEarlyTermination, RangeResultHandlerAddResultCMin) { + using C = faiss::CMin; + + faiss::RangeSearchResult rsr(1); + faiss::RangeSearchPartialResult pres(&rsr); + auto& qr = pres.new_result(0); + + faiss::RangeResultHandler handler(&qr, /*threshold=*/5.0f); + + // CMin::cmp(5.0, 10.0) is true: dis above threshold, should add + EXPECT_TRUE(handler.add_result(10.0f, 100)); + EXPECT_EQ(qr.nres, size_t(1)); + + // CMin::cmp(5.0, 5.0) is false: at boundary, should not add + EXPECT_FALSE(handler.add_result(5.0f, 101)); + EXPECT_EQ(qr.nres, size_t(1)); + + // CMin::cmp(5.0, 3.0) is false: dis below threshold + EXPECT_FALSE(handler.add_result(3.0f, 102)); + EXPECT_EQ(qr.nres, size_t(1)); + + // Second valid result + EXPECT_TRUE(handler.add_result(7.0f, 103)); + EXPECT_EQ(qr.nres, size_t(2)); + + pres.finalize(); + ASSERT_EQ(rsr.lims[1], size_t(2)); + std::set ids(rsr.labels, rsr.labels + 2); + EXPECT_EQ(ids, (std::set{100, 103})); + std::set dists(rsr.distances, rsr.distances + 2); + EXPECT_EQ(dists, (std::set{7.0f, 10.0f})); +} + // FastScan early-stop options use the per-query implementations. TEST(IVFEarlyTermination, FastscanDefaultsWork) { auto xb = make_data(nb, 0xcafebabe);