IDSelectorWithContext: extend context hook to IVFPQ/PQR, multibit RaBitQ, and Panorama scanners (facebookresearch#5491)

Summary:
Pull Request resolved: facebookresearch#5491

D113139231 added `faiss::IDSelectorWithContext` and wired exactly one scan loop (`run_scan_codes1`, used by IVFFlat / IVFSQ). This diff extends the same "hand the selector its scan context" hook to the remaining linear IVF per-candidate selector sites, so `IDSelectorWithContext` is honored uniformly instead of being silently dropped depending on index type or a runtime parameter (`nb_bits`).

A new shared helper `IDSelectorContextDispatch` (in `IDSelector.h`) centralizes the once-per-list `dynamic_cast` + per-candidate dispatch. `run_scan_codes1` is refactored onto it, and the new sites reuse it:

- `IVFPQ` / `IVFPQR`: `WrappedSearchResult::skip_entry` — one edit covers the table / pointer / on-the-fly / polysemous variants, and IVFPQR reuses the IVFPQ scanner.
- multibit `RaBitQ` (`nb_bits >= 2`): the two hand-written `scan_codes_multibit` loops (scalar + SIMD). This closes a real footgun where the hook was honored at 1 bit (via `run_scan_codes1`) but silently dropped at `nb_bits >= 2` within the same index type.
- `IVFFlatPanorama`: `progressive_filter_batch` — the scan-order position is the full-list `global_idx`, so a lookahead crosses batch boundaries correctly.

No behavior change: `is_member_with_context` returns a verdict identical to `is_member`; the context is a prefetch-only side channel. Selectors that do not implement the sub-interface, and `store_pairs` scans (synthetic ids), fall back to `is_member`. The base `IDSelector` is untouched, so every existing selector is unaffected.

Binary size: one small inline struct plus a cached pointer; no new template dimensions.

FastScan (`IVFPQFastScan` / `IVFRaBitQFastScan`) is intentionally deferred to a benchmark-gated follow-up: its result-lane scan is sparse and it is LUT/SIMD-bound, so the prefetch benefit there is speculative (correctness holds either way). Out of scope: `IVFSpectralHash` (asserts no selector), binary IVF, and HNSW (separate selector paths).

Stacks on D113139231.

Reviewed By: mnorris11

Differential Revision: D113913501

fbshipit-source-id: 011f52015f3e87b5b7cf91291edbe020bedadbeb

Author: Lengning Liu

faiss/impl/IDSelector.h

@@ -43,6 +43,30 @@ struct IDSelectorWithContext : IDSelector {
             const = 0;
 };
 
+/** Routes each per-candidate membership test to is_member_with_context() when
+ * the selector implements IDSelectorWithContext, else to plain is_member().
+ * Construct one per inverted-list scan: the dynamic_cast is the only RTTI cost
+ * and the per-candidate cost is a single predicted branch (cf. the
+ * IDSelectorRange dynamic_cast in IndexIVF.cpp). The scan context is only
+ * meaningful when the scan exposes a real id array (i.e. !store_pairs); when
+ * store_pairs the context path is disabled and every test falls back to
+ * is_member. */
+struct IDSelectorContextDispatch {
+    const IDSelector* sel;
+    const IDSelectorWithContext* ctx_sel;
+
+    IDSelectorContextDispatch(const IDSelector* sel, bool store_pairs)
+            : sel(sel),
+              ctx_sel((sel != nullptr && !store_pairs)
+                              ? dynamic_cast<const IDSelectorWithContext*>(sel)
+                              : nullptr) {}
+
+    bool is_member(idx_t id, const IDScanContext& ctx) const {
+        return ctx_sel ? ctx_sel->is_member_with_context(id, ctx)
+                       : sel->is_member(id);
+    }
+};
+
 /** ids between [imin, imax) */
 struct IDSelectorRange : IDSelector {
     idx_t imin, imax;

faiss/impl/Panorama.h

@@ -323,12 +323,20 @@ struct Panorama {
         size_t batch_offset = batch_no * batch_size * code_size;
         const uint8_t* storage_base = codes_base + batch_offset;
 
+        // Honor IDSelectorWithContext: the scan-order position within the whole
+        // list is global_idx, so a lookahead crosses batch boundaries correctly
+        // (ids is the full list, length list_size).
+        IDSelectorContextDispatch sel_dispatch(
+                sel, /*store_pairs=*/ids == nullptr);
+
         // Initialize active set with ID-filtered vectors.
         size_t num_active = 0;
         for (size_t i = 0; i < curr_batch_size; i++) {
             size_t global_idx = batch_start + i;
             idx_t id = (ids == nullptr) ? global_idx : ids[global_idx];
-            bool include = !use_sel || sel->is_member(id);
+            bool include = !use_sel ||
+                    sel_dispatch.is_member(
+                            id, IDScanContext{ids, list_size, global_idx});
 
             active_indices[num_active] = i;
             float cum_sum = batch_cum_sums[i];

faiss/impl/RaBitQuantizer.cpp

@@ -390,11 +390,17 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer {
         const size_t ex_bits = nb_bits - 1;
         FAISS_ASSERT(ex_bits > 0);
 
+        // Honor IDSelectorWithContext on the multibit path too, so a RaBitQ
+        // index does not silently lose the context hook once nb_bits >= 2 (the
+        // 1-bit path already routes through run_scan_codes1).
+        const IDSelectorContextDispatch sel_dispatch(sel, store_pairs);
+
         size_t nup = 0;
         for (size_t j = 0; j < list_size; j++) {
             if (sel != nullptr) {
                 idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                if (!sel->is_member(id)) {
+                if (!sel_dispatch.is_member(
+                            id, IDScanContext{ids, list_size, j})) {
                     codes += code_size;
                     continue;
                 }
@@ -601,11 +607,17 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer {
         const size_t ex_bits = nb_bits - 1;
         FAISS_ASSERT(ex_bits > 0);
 
+        // Honor IDSelectorWithContext on the multibit path too, so a RaBitQ
+        // index does not silently lose the context hook once nb_bits >= 2 (the
+        // 1-bit path already routes through run_scan_codes1).
+        const IDSelectorContextDispatch sel_dispatch(sel, store_pairs);
+
         size_t nup = 0;
         for (size_t j = 0; j < list_size; j++) {
             if (sel != nullptr) {
                 idx_t id = store_pairs ? lo_build(list_no, j) : ids[j];
-                if (!sel->is_member(id)) {
+                if (!sel_dispatch.is_member(
+                            id, IDScanContext{ids, list_size, j})) {
                     codes += code_size;
                     continue;
                 }

faiss/impl/expanded_scanners.h

@@ -36,24 +36,15 @@ size_t run_scan_codes1(
     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;
+    // context (ids, list_size, j) so it can exploit scan-order locality; the
+    // dispatch caches the once-per-list RTTI detection.
+    const IDSelectorContextDispatch sel_dispatch(sel, store_pairs);
     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
-            const bool member = ctx_sel
-                    ? ctx_sel->is_member_with_context(
-                              id, IDScanContext{ids, list_size, j})
-                    : sel->is_member(id);
-            if (!member) {
+            if (!sel_dispatch.is_member(id, IDScanContext{ids, list_size, j})) {
                 codes += code_size;
                 continue;
             }

faiss/impl/pq_code_distance/IVFPQScanner_impl.h

@@ -29,19 +29,31 @@ struct WrappedSearchResult {
     ResultHandler& res;
     size_t nup = 0;
     idx_t list_no;
-
+    size_t list_size;
     const idx_t* ids;
     const IDSelector* sel;
+    IDSelectorContextDispatch dispatch;
 
     WrappedSearchResult(
             idx_t list_no_in,
+            size_t list_size_in,
             const idx_t* ids_in,
             const IDSelector* sel_in,
             ResultHandler& res_in)
-            : res(res_in), list_no(list_no_in), ids(ids_in), sel(sel_in) {}
+            : res(res_in),
+              list_no(list_no_in),
+              list_size(list_size_in),
+              ids(ids_in),
+              sel(sel_in),
+              // A selector implies real ids, so ids==nullptr iff store_pairs;
+              // that disables the context path exactly when ids[] is synthetic.
+              dispatch(sel_in, /*store_pairs=*/ids_in == nullptr) {}
 
     inline bool skip_entry(idx_t j) {
-        return use_sel && !sel->is_member(ids[j]);
+        return use_sel &&
+                !dispatch.is_member(
+                        ids[j],
+                        IDScanContext{ids, list_size, static_cast<size_t>(j)});
     }
 
     inline void add(idx_t j, float dis) {
@@ -474,6 +486,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>,
             ResultHandler& handler) const override {
         WrappedSearchResult<C, use_sel> res(
                 this->key,
+                ncode,
                 this->store_pairs ? nullptr : ids,
                 this->sel,
                 handler);

tests/test_params_override.cpp

@@ -17,7 +17,9 @@
 #include <faiss/AutoTune.h>
 #include <faiss/IVFlib.h>
 #include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexFlat.h>
 #include <faiss/IndexIVF.h>
+#include <faiss/IndexIVFRaBitQ.h>
 #include <faiss/clone_index.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/IDSelector.h>
@@ -186,18 +188,21 @@ int test_defaults_preserve_baseline(const char* index_key, MetricType metric) {
 }
 
 // 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.
+// exercising the context parameters. Used to prove that a scanner (a) returns
+// results identical to the plain is_member path and (b) actually routes through
+// is_member_with_context — n_context_calls counts how often the context path
+// fired, so a scanner that silently dropped the hook is detectable.
 struct ContextSelector : IDSelectorWithContext {
     const IDSelector& inner;
+    mutable size_t n_context_calls = 0;
     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;
+        ++n_context_calls;
         return inner.is_member(id);
     }
 };
@@ -242,24 +247,22 @@ 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);
+// Runs the same membership set (every id where i % 10 == 2) through a plain
+// IDSelectorBatch and through a ContextSelector wrapping it, on an
+// already-trained index. Returns:
+//   0 if the context path fired AND labels are byte-identical to the plain
+//   path, 1 if the results diverged (a correctness regression), 2 if the
+//   scanner never routed through is_member_with_context (the hook was
+//     silently dropped for this index type).
+int check_ctx_matches_plain(Index& index) {
     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.
@@ -268,24 +271,50 @@ int test_selector_with_context(const char* index_key) {
     plain_params.nprobe = 3;
     plain_params.sel = &batch;
     auto plain_result =
-            search_index_with_params(index.get(), xq.data(), &plain_params);
+            search_index_with_params(&index, 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);
+    auto ctx_result = search_index_with_params(&index, xq.data(), &ctx_params);
 
     if (plain_result != ctx_result) {
         return 1;
     }
-
+    if (ctx_sel.n_context_calls == 0) {
+        return 2;
+    }
     return 0;
 }
 
+// Same membership set as test_selector, but driven through an
+// IDSelectorWithContext to prove the context scan path is functionally
+// identical to — and actually exercised by — the given index type.
+int test_selector_with_context(const char* index_key) {
+    std::vector<float> xb = make_data(nb);
+    auto index = make_index(index_key, METRIC_L2, xb);
+    // nprobe is set explicitly in the IVFSearchParameters inside
+    // check_ctx_matches_plain, so no need to set it on the index here.
+    return check_ctx_matches_plain(*index);
+}
+
+// Multibit RaBitQ (nb_bits >= 2) scans through the two hand-written loops in
+// RaBitQuantizer.cpp rather than run_scan_codes1, so it needs the context hook
+// wired separately. nb_bits=3 => ex_bits=2 => scan_codes_multibit; the default
+// qb keeps a real quantized-query distance computer (no 1-bit fallback).
+int test_selector_with_context_rabitq_multibit() {
+    std::vector<float> xb = make_data(nb);
+    IndexFlatL2 quantizer(d);
+    IndexIVFRaBitQ index(
+            &quantizer, d, 32, METRIC_L2, /*own_invlists=*/true, /*nb_bits=*/3);
+    index.train(nb, xb.data());
+    index.add(nb, xb.data());
+    return check_ctx_matches_plain(index);
+}
+
 } // namespace
 
 /*************************************************************
@@ -362,10 +391,11 @@ 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.
+// An IDSelectorWithContext must be functionally identical to a plain
+// IDSelector AND actually exercised (results must not silently fall back to
+// is_member) on every applicable IVF scanner. IVFFlat/IVFSQ route through
+// run_scan_codes1; IVFPQ through WrappedSearchResult::skip_entry; multibit
+// RaBitQ through its two hand-written scan loops.
 TEST(TSELCtx, IVFFlat) {
     EXPECT_EQ(test_selector_with_context("PCA16,IVF32,Flat"), 0);
 }
@@ -378,6 +408,14 @@ TEST(TSELCtx, IVFFSQ) {
     EXPECT_EQ(test_selector_with_context("PCA16,IVF32,SQ8"), 0);
 }
 
+TEST(TSELCtx, RaBitQMultibit) {
+    EXPECT_EQ(test_selector_with_context_rabitq_multibit(), 0);
+}
+
+TEST(TSELCtx, Panorama) {
+    EXPECT_EQ(test_selector_with_context("IVF32,FlatPanorama"), 0);
+}
+
 /*************************************************************
  * Same for binary indexes
  *************************************************************/