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 <faiss/utils/Heap.h>`.
- **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
This commit is contained in:
Michael Norris
2026-05-04 10:25:15 -07:00
committed by meta-codesync[bot]
parent 1e69d5ab20
commit 40a8cc0034
5 changed files with 77 additions and 5 deletions
+1 -2
View File
@@ -421,7 +421,6 @@ 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);
@@ -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 &&
-1
View File
@@ -16,7 +16,6 @@
#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>
#include <faiss/utils/distances_dispatch.h>
#include <faiss/utils/extra_distances.h>
+1
View File
@@ -596,6 +596,7 @@ struct RangeResultHandler : ResultHandlerT<C> {
bool add_result(T dis, TI idx) final {
if (C::cmp(threshold, dis)) {
qr->add(dis, idx);
return true;
}
return false;
}
+2 -2
View File
@@ -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
+73
View File
@@ -28,6 +28,7 @@
#include <faiss/impl/AuxIndexStructures.h>
#include <faiss/impl/FaissException.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/index_factory.h>
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<float, faiss::idx_t>;
faiss::RangeSearchResult rsr(1);
faiss::RangeSearchPartialResult pres(&rsr);
auto& qr = pres.new_result(0);
faiss::RangeResultHandler<C> 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<faiss::idx_t> ids(rsr.labels, rsr.labels + 2);
EXPECT_EQ(ids, (std::set<faiss::idx_t>{42, 45}));
std::set<float> dists(rsr.distances, rsr.distances + 2);
EXPECT_EQ(dists, (std::set<float>{1.0f, 3.0f}));
}
// Same test with CMin (inner product semantics): adds when threshold < dis.
TEST(IVFEarlyTermination, RangeResultHandlerAddResultCMin) {
using C = faiss::CMin<float, faiss::idx_t>;
faiss::RangeSearchResult rsr(1);
faiss::RangeSearchPartialResult pres(&rsr);
auto& qr = pres.new_result(0);
faiss::RangeResultHandler<C> 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<faiss::idx_t> ids(rsr.labels, rsr.labels + 2);
EXPECT_EQ(ids, (std::set<faiss::idx_t>{100, 103}));
std::set<float> dists(rsr.distances, rsr.distances + 2);
EXPECT_EQ(dists, (std::set<float>{7.0f, 10.0f}));
}
// FastScan early-stop options use the per-query implementations.
TEST(IVFEarlyTermination, FastscanDefaultsWork) {
auto xb = make_data(nb, 0xcafebabe);