IDSelectorWithContext: scan-context hook for locality-aware selectors (#5490)

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

@@ -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;

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<const IDSelectorWithContext*>(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;
             }

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<float> xb = make_data(nb); // database vectors
     std::vector<float> 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<float> xb = make_data(nb);
+    std::vector<float> xq = make_data(nq);
+    ParameterSpace ps;
+
+    std::vector<idx_t> 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
  *************************************************************/