diff --git a/faiss/impl/HNSW.cpp b/faiss/impl/HNSW.cpp index 136818648..d3b5465b7 100644 --- a/faiss/impl/HNSW.cpp +++ b/faiss/impl/HNSW.cpp @@ -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 TopCandidatesQueue search_from_candidate_unbounded_fixVT( const HNSW& hnsw, @@ -1575,16 +1577,48 @@ TopCandidatesQueue search_from_candidate_unbounded_fixVT( DistanceComputer& qdis, int ef, VTType& vt, - HNSWStats& stats) { + HNSWStats& stats, + const IDSelector* sel) { + const auto max_size = static_cast(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 top_candidates; reservePriorityQueue(top_candidates, ef); + // `selected_candidates` keeps only the accepted nodes, and + // the function returns it instead when `sel` is set. + TopCandidatesQueue selected_candidates; + if (sel) { + reservePriorityQueue(selected_candidates, ef); + } + CandidatesQueue 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& 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 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(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(ef)) { - top_candidates.pop(); - } } }; @@ -1674,6 +1705,9 @@ TopCandidatesQueue search_from_candidate_unbounded_fixVT( } stats.ndis += ndis; + if (sel) { + return selected_candidates; + } return top_candidates; } @@ -1693,7 +1727,7 @@ std::priority_queue hnsw_detail::search_from_candidate_unbounded( using C = HNSW::C_distance; auto call = [&](VTType& vt_concrete) { return search_from_candidate_unbounded_fixVT( - hnsw, node, qdis, ef, vt_concrete, stats); + hnsw, node, qdis, ef, vt_concrete, stats, nullptr); }; if (VisitedTableVector* vtv = dynamic_cast(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 = [&](VTType& vt_concrete) { return search_from_candidate_unbounded_fixVT( hnsw, @@ -1799,7 +1834,8 @@ HNSWStats search_impl( qdis, ef, vt_concrete, - stats); + stats, + sel); }; TopCandidatesQueue top_candidates; if (VisitedTableVector* vtv = dynamic_cast(&vt)) { diff --git a/tests/test_search_params.py b/tests/test_search_params.py index d11a90775..d56a0a308 100644 --- a/tests/test_search_params.py +++ b/tests/test_search_params.py @@ -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)