Fix IDSelector handling in unbounded HNSW search (#5508)

Summary:
- apply `SearchParametersHNSW.sel` to the unbounded HNSW search path
- keep filtered nodes available for graph traversal while excluding them from the result set
- add regression coverage for L2 and inner-product search, restrictive `efSearch`, and an empty selector

## Root cause

The bounded HNSW path applies the selector through its result handler. The unbounded path did not receive the selector and used the same heap for traversal pruning and final results, so IDs rejected by `SearchParameters.sel` could be returned.

## Implementation

When a selector is present, the unbounded path maintains a separate result heap containing only accepted IDs. The existing candidate heap remains unchanged and can still contain rejected nodes, preserving traversal through them. When no selector is present, the existing unfiltered result behavior is preserved.

## Validation

- issue reproducer: unbounded search now returns ID 1 at distance 100.0 instead of rejected ID 0
- `tests/test_search_params.py`: 52 passed
- HNSW C++ tests: 14 passed
- full CPU C++ test suite: 278 passed, 5 skipped
- clang-format and `git diff --check`: passed

Fixes https://github.com/facebookresearch/faiss/issues/5504

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5508

Reviewed By: alibeklfc

Differential Revision: D120364656

Pulled By: mnorris11

fbshipit-source-id: 074234b80cf7b2e06a85d57489a1c214dfd3195b
This commit is contained in:
YingzuoLiu
2026-09-16 17:55:51 -07:00
committed by meta-codesync[bot]
parent 1dbb3d25b0
commit 3778f5892d
2 changed files with 105 additions and 10 deletions
+46 -10
View File
@@ -1568,6 +1568,8 @@ void reservePriorityQueue(
/// Templated body of `search_from_candidate_unbounded`. The choice of
/// max-heap vs min-heap for both `top_candidates` and `candidates` is
/// derived from C via `TopCandidatesQueue` / `CandidatesQueue`.
/// `sel` filters the results only. The traversal still visits the rejected
/// nodes.
template <typename VTType, class C>
TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
const HNSW& hnsw,
@@ -1575,16 +1577,48 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
DistanceComputer& qdis,
int ef,
VTType& vt,
HNSWStats& stats) {
HNSWStats& stats,
const IDSelector* sel) {
const auto max_size = static_cast<size_t>(ef);
int ndis = 0;
// `top_candidates` bounds the traversal and keeps the nodes that `sel`
// rejects, because a rejected node is still a valid hop towards an
// accepted one.
TopCandidatesQueue<C> top_candidates;
reservePriorityQueue(top_candidates, ef);
// `selected_candidates` keeps only the accepted nodes, and
// the function returns it instead when `sel` is set.
TopCandidatesQueue<C> selected_candidates;
if (sel) {
reservePriorityQueue(selected_candidates, ef);
}
CandidatesQueue<C> candidates;
reservePriorityQueue(candidates, ef);
// Adds a node to a queue of the best `max_size` nodes. The size test
// comes first because `top()` needs a queue that is not empty.
auto push_bounded = [max_size](
TopCandidatesQueue<C>& queue,
const float dis,
const size_t idx) {
if (queue.size() >= max_size && !C::cmp(queue.top().first, dis)) {
return false;
}
queue.emplace(dis, idx);
if (queue.size() > max_size) {
queue.pop();
}
return true;
};
top_candidates.push(node);
candidates.push(node);
if (sel && sel->is_member(node.second)) {
selected_candidates.push(node);
}
vt.set(node.second);
@@ -1619,14 +1653,11 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
size_t saved_j[4];
auto add_to_heap = [&](const size_t idx, const float dis) {
if (C::cmp(top_candidates.top().first, dis) ||
top_candidates.size() < static_cast<size_t>(ef)) {
if (sel && sel->is_member(idx)) {
push_bounded(selected_candidates, dis, idx);
}
if (push_bounded(top_candidates, dis, idx)) {
candidates.emplace(dis, idx);
top_candidates.emplace(dis, idx);
if (top_candidates.size() > static_cast<size_t>(ef)) {
top_candidates.pop();
}
}
};
@@ -1674,6 +1705,9 @@ TopCandidatesQueue<C> search_from_candidate_unbounded_fixVT(
}
stats.ndis += ndis;
if (sel) {
return selected_candidates;
}
return top_candidates;
}
@@ -1693,7 +1727,7 @@ std::priority_queue<HNSW::Node> hnsw_detail::search_from_candidate_unbounded(
using C = HNSW::C_distance;
auto call = [&]<typename VTType>(VTType& vt_concrete) {
return search_from_candidate_unbounded_fixVT<VTType, C>(
hnsw, node, qdis, ef, vt_concrete, stats);
hnsw, node, qdis, ef, vt_concrete, stats, nullptr);
};
if (VisitedTableVector* vtv = dynamic_cast<VisitedTableVector*>(vt)) {
return call(*vtv);
@@ -1792,6 +1826,7 @@ HNSWStats search_impl(
hnsw.search_method == HNSW::SM_DEFAULT ||
hnsw.search_method == HNSW::SM_PANORAMA,
"invalid HNSW search method");
const IDSelector* sel = params ? params->sel : nullptr;
auto call = [&]<typename VTType>(VTType& vt_concrete) {
return search_from_candidate_unbounded_fixVT<VTType, C>(
hnsw,
@@ -1799,7 +1834,8 @@ HNSWStats search_impl(
qdis,
ef,
vt_concrete,
stats);
stats,
sel);
};
TopCandidatesQueue<C> top_candidates;
if (VisitedTableVector* vtv = dynamic_cast<VisitedTableVector*>(&vt)) {
+59
View File
@@ -302,6 +302,65 @@ class TestSelector(unittest.TestCase):
def test_HSNW(self):
self.do_test_id_selector_weak("HNSW")
def make_line_index(self, metric=faiss.METRIC_L2):
"""build an HNSW index over 4 points on a line"""
xb = np.array(
[[0.0, 0.0], [10.0, 0.0], [20.0, 0.0], [30.0, 0.0]],
dtype=np.float32,
)
index = faiss.IndexHNSWFlat(2, 4, metric)
index.hnsw.efConstruction = 40
index.add(xb)
return xb, index
def unbounded_params(self, ef_search, sel):
return faiss.SearchParametersHNSW(
efSearch=ef_search, bounded_queue=False, sel=sel
)
def test_HNSW_unbounded_queue_id_selector(self):
cases = (
(faiss.METRIC_L2, [[0.0, 0.0]], 100.0),
(faiss.METRIC_INNER_PRODUCT, [[1.0, 0.0]], 10.0),
)
for metric, query, expected_distance in cases:
_, index = self.make_line_index(metric)
xq = np.array(query, dtype=np.float32)
for ef_search in (1, 40):
with self.subTest(metric=metric, ef_search=ef_search):
params = self.unbounded_params(
ef_search, faiss.IDSelectorRange(1, 2)
)
D, I = index.search(xq, 1, params=params)
np.testing.assert_array_equal(I, [[1]])
np.testing.assert_array_equal(D, [[expected_distance]])
def test_HNSW_unbounded_queue_selector_rejects_all(self):
_, index = self.make_line_index()
params = self.unbounded_params(40, faiss.IDSelectorRange(4, 4))
xq = np.zeros((1, 2), dtype=np.float32)
_, I = index.search(xq, 1, params=params)
np.testing.assert_array_equal(I, [[-1]])
def test_HNSW_unbounded_queue_rejected_starting_node(self):
xb, index = self.make_line_index()
starting_id = index.hnsw.entry_point
allowed_id = (starting_id + 1) % len(xb)
params = self.unbounded_params(
1, faiss.IDSelectorRange(allowed_id, allowed_id + 1)
)
# The query coincides with the entry point (distance zero), so the
# upper-level search keeps it as the level-0 starting node.
xq = xb[starting_id : starting_id + 1]
expected_distance = np.sum((xq[0] - xb[allowed_id]) ** 2)
D, I = index.search(xq, 1, params=params)
np.testing.assert_array_equal(I, [[allowed_id]])
np.testing.assert_array_equal(D, [[expected_distance]])
def test_idmap(self):
ds = datasets.SyntheticDataset(32, 100, 100, 20)
rs = np.random.RandomState(123)