From edddbc998e858f16330f9b226dfd72f73d4d1d46 Mon Sep 17 00:00:00 2001 From: Lengning Liu Date: Wed, 5 Aug 2026 10:36:30 -0700 Subject: [PATCH] IDSelectorWithContext: scan-context hook for locality-aware selectors (#5490) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5490 Adds `faiss::IDSelectorWithContext`, an `IDSelector` sub-interface whose `is_member_with_context(id, ids, list_size, j)` also receives the surrounding scan context, so an implementation can exploit locality across a scan (e.g. prefetch data for an entry it will be asked about soon). The IVF inner scan loop `run_scan_codes1` `dynamic_cast`s the selector once per inverted list (RTTI, as already used for `IDSelectorRange` in IndexIVF.cpp) and, when it matches, calls `is_member_with_context` instead of `is_member` — a single predicted branch per candidate, no template dimension, and the base `IDSelector` is unchanged so every existing selector is unaffected. No behavior change on its own: a selector that does not implement the sub-interface takes the plain `is_member` path. First consumer is admarket's `FilterableIDSelector` (stacked on top). Reviewed By: mnorris11 Differential Revision: D113139231 fbshipit-source-id: 48a5ede7b502771b0b9a66593719502c5202c6da --- faiss/impl/IDSelector.h | 20 +++++++++ faiss/impl/expanded_scanners.h | 15 ++++++- tests/test_params_override.cpp | 77 ++++++++++++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) diff --git a/faiss/impl/IDSelector.h b/faiss/impl/IDSelector.h index 631991bd1..6923864c6 100644 --- a/faiss/impl/IDSelector.h +++ b/faiss/impl/IDSelector.h @@ -23,6 +23,26 @@ struct IDSelector { virtual ~IDSelector() {} }; +/** Scan context handed to IDSelectorWithContext::is_member_with_context: the + * contiguous id block being scanned and the position of the id under test. */ +struct IDScanContext { + /// the contiguous block of ids being scanned + const idx_t* ids; + /// number of entries in `ids` + size_t list_size; + /// index of the tested id within `ids` (i.e. ids[j] == id) + size_t j; +}; + +/** IDSelector that also receives the surrounding scan context on each + * membership test, letting an implementation exploit locality across a scan + * (for example, prefetching data for an entry it will be asked about soon). + * The policy is entirely up to the implementation. */ +struct IDSelectorWithContext : IDSelector { + virtual bool is_member_with_context(idx_t id, const IDScanContext& ctx) + const = 0; +}; + /** ids between [imin, imax) */ struct IDSelectorRange : IDSelector { idx_t imin, imax; diff --git a/faiss/impl/expanded_scanners.h b/faiss/impl/expanded_scanners.h index 8eab6ab2a..fbc9722c1 100644 --- a/faiss/impl/expanded_scanners.h +++ b/faiss/impl/expanded_scanners.h @@ -35,12 +35,25 @@ size_t run_scan_codes1( size_t list_no = scanner.list_no; size_t code_size = scanner.code_size; const IDSelector* sel = scanner.sel; + // If the selector implements IDSelectorWithContext, hand it the scan + // context (ids, list_size, j) so it can exploit scan-order locality. Detect + // it once per list via RTTI (cf. the IDSelectorRange dynamic_cast in + // IndexIVF.cpp) so the per-candidate cost is a single predicted branch. + // store_pairs never coexists with a selector, and ids[] is only valid when + // !store_pairs. + const IDSelectorWithContext* ctx_sel = (use_sel && !store_pairs) + ? dynamic_cast(sel) + : nullptr; float threshold = handler.threshold; for (size_t j = 0; j < list_size; j++) { if (use_sel) { int64_t id = store_pairs ? lo_build(list_no, j) : ids[j]; // skip code without computing distance - if (!sel->is_member(id)) { + const bool member = ctx_sel + ? ctx_sel->is_member_with_context( + id, IDScanContext{ids, list_size, j}) + : sel->is_member(id); + if (!member) { codes += code_size; continue; } diff --git a/tests/test_params_override.cpp b/tests/test_params_override.cpp index 15d5fef09..82dc0779a 100644 --- a/tests/test_params_override.cpp +++ b/tests/test_params_override.cpp @@ -185,6 +185,23 @@ int test_defaults_preserve_baseline(const char* index_key, MetricType metric) { return 0; } +// IDSelectorWithContext that delegates the verdict to an inner IDSelector while +// exercising the context parameters — used to prove the scan's +// is_member_with_context path returns results identical to the plain is_member +// path. +struct ContextSelector : IDSelectorWithContext { + const IDSelector& inner; + explicit ContextSelector(const IDSelector& inner) : inner(inner) {} + bool is_member(idx_t id) const override { + return inner.is_member(id); + } + bool is_member_with_context(idx_t id, const IDScanContext& ctx) + const override { + (void)ctx; + return inner.is_member(id); + } +}; + int test_selector(const char* index_key) { std::vector xb = make_data(nb); // database vectors std::vector xq = make_data(nq); @@ -225,6 +242,50 @@ int test_selector(const char* index_key) { return 0; } +// Same membership set as test_selector, but driven through an +// IDSelectorWithContext to prove the context scan path is functionally +// identical to the plain IDSelector path. +int test_selector_with_context(const char* index_key) { + std::vector xb = make_data(nb); + std::vector xq = make_data(nq); + ParameterSpace ps; + + std::vector kept; + for (size_t i = 0; i < nb; i++) { + if (i % 10 == 2) { + kept.push_back(i); + } + } + + auto index = make_index(index_key, METRIC_L2, xb); + ps.set_index_parameter(index.get(), "nprobe", 3); + + IDSelectorBatch batch(kept.size(), kept.data()); + + // Plain IDSelector path. + IVFSearchParameters plain_params; + plain_params.max_codes = 0; + plain_params.nprobe = 3; + plain_params.sel = &batch; + auto plain_result = + search_index_with_params(index.get(), xq.data(), &plain_params); + + // IDSelectorWithContext path over the same membership set. + ContextSelector ctx_sel(batch); + IVFSearchParameters ctx_params; + ctx_params.max_codes = 0; + ctx_params.nprobe = 3; + ctx_params.sel = &ctx_sel; + auto ctx_result = + search_index_with_params(index.get(), xq.data(), &ctx_params); + + if (plain_result != ctx_result) { + return 1; + } + + return 0; +} + } // namespace /************************************************************* @@ -301,6 +362,22 @@ TEST(TSEL, IVFFSQ) { EXPECT_EQ(err, 0); } +// IDSelectorWithContext must be functionally identical to a plain IDSelector. +// IVFFlat and IVFSQ route through run_scan_codes1 (which dispatches to +// is_member_with_context); IVFPQ uses its own scanner (plain is_member) — all +// must return identical results. +TEST(TSELCtx, IVFFlat) { + EXPECT_EQ(test_selector_with_context("PCA16,IVF32,Flat"), 0); +} + +TEST(TSELCtx, IVFFPQ) { + EXPECT_EQ(test_selector_with_context("PCA16,IVF32,PQ4x8np"), 0); +} + +TEST(TSELCtx, IVFFSQ) { + EXPECT_EQ(test_selector_with_context("PCA16,IVF32,SQ8"), 0); +} + /************************************************************* * Same for binary indexes *************************************************************/