diff --git a/src/index/hnsw/graph_ops.h b/src/index/hnsw/graph_ops.h index b9821e6..cb27058 100644 --- a/src/index/hnsw/graph_ops.h +++ b/src/index/hnsw/graph_ops.h @@ -75,6 +75,30 @@ template class GraphOperations { using Node = typename Graph::Node; using NodeData = typename Graph::NodeData; using LevelId = typename Graph::LevelId; + using LayerOps = LayerOperations; + using ExtendCandidates = typename LayerOps::ExtendCandidates; + using KeepPrunedConnections = typename LayerOps::KeepPrunedConnections; + + // Algorithm 4 heuristic parameters used by all GraphOperations calls to + // consume_heuristic(): insert()'s own-neighbour selection, remove()'s + // repair-neighbour selection, and shrink_neighbours()'s Mmax reselection. + // + // Per Section 3 of the HNSW paper, extendCandidates, when enabled, extends + // the candidate set with neighbours of the candidates. It defaults to + // false, and we keep that default. + // + // keepPrunedConnections retains pruned candidates to maintain a fixed + // number of connections per element. We disable it, allowing the heuristic + // to select the most diverse neighbours without retaining pruned candidates. + // + // Keep these centralized as named constants rather than literals at each + // call site, so there is a single place to change the policy if needed. + // TODO(villagesql-indexing): Consider making SHOULD_EXTEND_CANDIDATES and + // SHOULD_KEEP_PRUNED_CONNECTIONS configurable. + static constexpr ExtendCandidates SHOULD_EXTEND_CANDIDATES = + ExtendCandidates::No; + static constexpr KeepPrunedConnections SHOULD_KEEP_PRUNED_CONNECTIONS = + KeepPrunedConnections::No; GraphOperations(Graph &graph) : m_graph(graph) {}; @@ -86,8 +110,6 @@ template class GraphOperations { uint32_t ef_search, std::vector &nearest_nodes); private: - using LayerOps = LayerOperations; - // Algorithm 1, lines 5-7: ef is fixed at 1 for the greedy descent. static constexpr uint32_t GREEDY_DESCENT_EF = 1; @@ -95,6 +117,18 @@ template class GraphOperations { // level is the level every candidate currently lives at. bool advance_to_next_level(std::vector &candidates, LevelId level); + // Wraps LayerOps::consume_heuristic(), defaulting extend_candidates and + // keep_pruned_connections to SHOULD_EXTEND_CANDIDATES and + // SHOULD_KEEP_PRUNED_CONNECTIONS so call sites only need to supply what + // varies -- M, the output, and (optionally) candidate_pool -- while still + // allowing either policy to be overridden explicitly. + bool consume_heuristic( + LayerOps &layer, uint32_t M, std::vector &out, + std::vector *candidate_pool = nullptr, + ExtendCandidates extend_candidates = SHOULD_EXTEND_CANDIDATES, + KeepPrunedConnections keep_pruned_connections = + SHOULD_KEEP_PRUNED_CONNECTIONS); + // Reselects up to Mmax(level) neighbours for each node in 'overflowed', // after 'linked_node' was withheld from being linked back to it by // link_neighbours() (Algorithm 1, lines 14-15). level is the level shared diff --git a/src/index/hnsw/graph_ops_impl.h b/src/index/hnsw/graph_ops_impl.h index fa412d0..d2bf8d0 100644 --- a/src/index/hnsw/graph_ops_impl.h +++ b/src/index/hnsw/graph_ops_impl.h @@ -47,10 +47,6 @@ bool GraphOperations::insert(const NodeData &new_node_data) { using LockLevels = typename Graph::LockLevels; using DescendPolicy = typename Graph::LockLevels::DescendPolicy; - using LayerOps = LayerOperations; - using ExtendCandidates = typename LayerOps::ExtendCandidates; - using KeepPrunedConnections = typename LayerOps::KeepPrunedConnections; - // Lock graph in Shared Mode. LockGraph graph_lock(m_graph, LockMode::Shared); @@ -140,13 +136,7 @@ bool GraphOperations::insert(const NodeData &new_node_data) { if (layer.search(candidates, level, m_graph.ef_construction())) { return true; } - // Use the standard HNSW Algorithm 4 settings: - // - extend_candidates = false - // - keep_pruned_connections = true - // TODO(villagesql-indexing): Consider making these options configurable. - if (layer.consume_heuristic(m_graph.M(), ExtendCandidates::No, - KeepPrunedConnections::Yes, neighbours, - &candidates)) { + if (consume_heuristic(layer, m_graph.M(), neighbours, &candidates)) { return true; } } @@ -223,10 +213,6 @@ bool GraphOperations::remove(const Node &target_node, using DescendPolicy = typename Graph::LockLevels::DescendPolicy; using UnlinkOrphans = typename Graph::UnlinkOrphans; - using LayerOps = LayerOperations; - using ExtendCandidates = typename LayerOps::ExtendCandidates; - using KeepPrunedConnections = typename LayerOps::KeepPrunedConnections; - // Lock graph in Shared Mode. LockGraph graph_lock(m_graph, LockMode::Shared); @@ -323,9 +309,7 @@ bool GraphOperations::remove(const Node &target_node, return true; } std::vector new_neighbours; - if (layer.consume_heuristic(m_graph.M(), ExtendCandidates::No, - KeepPrunedConnections::Yes, - new_neighbours)) { + if (consume_heuristic(layer, m_graph.M(), new_neighbours)) { return true; } if (m_graph.replace_neighbours(orphan, level, new_neighbours)) { @@ -489,6 +473,15 @@ bool GraphOperations::advance_to_next_level( return false; } +template +bool GraphOperations::consume_heuristic( + LayerOps &layer, uint32_t M, std::vector &out, + std::vector *candidate_pool, ExtendCandidates extend_candidates, + KeepPrunedConnections keep_pruned_connections) { + return layer.consume_heuristic(M, extend_candidates, keep_pruned_connections, + out, candidate_pool); +} + // Algorithm 1, lines 14-15: shrink connections of neighbours whose degree // would otherwise exceed Mmax. link_neighbours() withholds the edge to // linked_node for exactly these nodes rather than adding it outright, so @@ -498,9 +491,6 @@ template bool GraphOperations::shrink_neighbours( LayerOps &layer, const Node &linked_node, LevelId level, const std::vector &overflowed) { - using ExtendCandidates = typename LayerOps::ExtendCandidates; - using KeepPrunedConnections = typename LayerOps::KeepPrunedConnections; - for (const Node &neighbour : overflowed) { std::vector connections; if (m_graph.neighbours(neighbour, level, connections)) { @@ -517,8 +507,7 @@ bool GraphOperations::shrink_neighbours( return true; } std::vector shrunk; - if (layer.consume_heuristic(m_graph.Mmax(level), ExtendCandidates::No, - KeepPrunedConnections::No, shrunk)) { + if (consume_heuristic(layer, m_graph.Mmax(level), shrunk)) { return true; } if (m_graph.replace_neighbours(neighbour, level, shrunk)) { diff --git a/src/index/hnsw/layer_ops_impl.h b/src/index/hnsw/layer_ops_impl.h index 7c2d959..b2b1f7e 100644 --- a/src/index/hnsw/layer_ops_impl.h +++ b/src/index/hnsw/layer_ops_impl.h @@ -320,6 +320,12 @@ bool LayerOperations::seed_impl( assert(m_results.empty()); assert(m_expand_buf.empty()); + // If the query is an existing graph node, mark it visited up front so it + // cannot reappear as its own candidate during neighbour expansion. + if (const Node *query_node = std::get_if(&m_query)) { + m_visited.insert(query_node->key()); + } + // Algorithm 2, lines 1-3: v = C = W = ep. for (const Node &node : entry_points) { if (m_visited.insert(node.key()).second) { diff --git a/unittest/graph_ops_test.cc b/unittest/graph_ops_test.cc index 48f5e8e..e81af4d 100644 --- a/unittest/graph_ops_test.cc +++ b/unittest/graph_ops_test.cc @@ -790,28 +790,54 @@ void test_insert_shrinks_overflowed_neighbours_past_mmax() { assert(!do_insert(g, 2, 3)); - // 2's own outgoing edges are unaffected: shrinking only touches the - // *neighbour's* list (graph_ops_impl.h's shrink_neighbours()), never the - // newly inserted node's own list. - assert((g.adjacency[0][2] == std::vector{1, 0})); - // 1 is closer to 0 (distance 1) than to 2 (distance 2), so with only one - // slot available the heuristic keeps 0 and drops the tentative edge to 2. - // Symmetric for 0, which is closer to 1 (distance 1) than to 2 (distance - // 3). Both neighbours end up asymmetrically linked: 2 points at them, but - // they don't point back -- exactly what link_neighbours() withholding the - // edge and shrink_neighbours() reselecting is meant to produce. + using GraphOps = svector::hnsw::GraphOperations; + + // 2's own neighbour selection: candidates are 1 (distance 2) and 0 + // (distance 3). 1 is nearest and always kept. 0 is dominated -- closer to + // 1 (distance 1) than to the query -- so whether it survives depends on + // GraphOps::SHOULD_KEEP_PRUNED_CONNECTIONS: Yes backfills it in anyway, No + // discards it outright. Branching on the same constant GraphOperations + // itself uses means this test keeps working whichever way that default is + // set, without needing hand-editing every time it's revisited. + // + // 2's own outgoing edges are unaffected either way: shrinking only touches + // the *neighbour's* list (graph_ops_impl.h's shrink_neighbours()), never + // the newly inserted node's own list. + if (GraphOps::SHOULD_KEEP_PRUNED_CONNECTIONS == + GraphOps::KeepPrunedConnections::Yes) { + assert((g.adjacency[0][2] == std::vector{1, 0})); + } else { + assert((g.adjacency[0][2] == std::vector{1})); + } + + // 1's reciprocal link back to 2 would push it over Mmax=1, so + // shrink_neighbours() reselects from 1's existing connection (0) plus the + // tentative one (2): 0 is nearer to 1 (distance 1) than 2 is (distance 2), + // so 0 wins regardless of keep_pruned_connections -- 1 ends up back on its + // original edge. assert((g.adjacency[0][1] == std::vector{0})); + + // 0 only gets a reciprocal-link/shrink attempt at all if 2 selected it as + // a neighbour in the first place, which is exactly the branch above. + // Either way 0 keeps its original edge to 1: with + // SHOULD_KEEP_PRUNED_CONNECTIONS == Yes, 0's own shrink (candidates {1, 2}, + // Mmax=1) keeps nearer 1 over 2; with No, 0 is never a candidate for + // anything and is simply untouched. assert((g.adjacency[0][0] == std::vector{1})); - // Both overflowed neighbours must have gone through the shrink-and-replace - // path, not a plain reciprocal link. + // Exactly as many neighbours as 2 ended up linking to (1, or 1 and 0) go + // through the shrink-and-replace path. int replace_count = 0; for (const std::string &entry : g.call_log) { if (entry.rfind("replace:", 0) == 0) { ++replace_count; } } - assert(replace_count == 2); + int expected_replace_count = GraphOps::SHOULD_KEEP_PRUNED_CONNECTIONS == + GraphOps::KeepPrunedConnections::Yes + ? 2 + : 1; + assert(replace_count == expected_replace_count); } void test_insert_replace_neighbours_failure_propagates() {