HNSW Index: Unify heuristic policies

src/index/hnsw/graph_ops.h

@@ -75,6 +75,30 @@ template <typename Graph> class GraphOperations {
   using Node = typename Graph::Node;
   using NodeData = typename Graph::NodeData;
   using LevelId = typename Graph::LevelId;
+  using LayerOps = LayerOperations<Graph, AlwaysVisiblePolicy>;
+  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,15 +110,25 @@ template <typename Graph> class GraphOperations {
                   uint32_t ef_search, std::vector<Node> &nearest_nodes);
 
 private:
-  using LayerOps = LayerOperations<Graph, AlwaysVisiblePolicy>;
-
   // Algorithm 1, lines 5-7: ef is fixed at 1 for the greedy descent.
   static constexpr uint32_t GREEDY_DESCENT_EF = 1;
 
   // Replaces each candidate with its counterpart at the next lower level.
   // level is the level every candidate currently lives at.
   bool advance_to_next_level(std::vector<Node> &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<Node> &out,
+      std::vector<Node> *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

src/index/hnsw/graph_ops_impl.h

@@ -47,10 +47,6 @@ bool GraphOperations<Graph>::insert(const NodeData &new_node_data) {
   using LockLevels = typename Graph::LockLevels;
   using DescendPolicy = typename Graph::LockLevels::DescendPolicy;
 
-  using LayerOps = LayerOperations<Graph, AlwaysVisiblePolicy>;
-  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<Graph>::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<Graph>::remove(const Node &target_node,
   using DescendPolicy = typename Graph::LockLevels::DescendPolicy;
   using UnlinkOrphans = typename Graph::UnlinkOrphans;
 
-  using LayerOps = LayerOperations<Graph, AlwaysVisiblePolicy>;
-  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<Graph>::remove(const Node &target_node,
           return true;
         }
         std::vector<Node> 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<Graph>::advance_to_next_level(
   return false;
 }
 
+template <typename Graph>
+bool GraphOperations<Graph>::consume_heuristic(
+    LayerOps &layer, uint32_t M, std::vector<Node> &out,
+    std::vector<Node> *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 <typename Graph>
 bool GraphOperations<Graph>::shrink_neighbours(
     LayerOps &layer, const Node &linked_node, LevelId level,
     const std::vector<Node> &overflowed) {
-  using ExtendCandidates = typename LayerOps::ExtendCandidates;
-  using KeepPrunedConnections = typename LayerOps::KeepPrunedConnections;
-
   for (const Node &neighbour : overflowed) {
     std::vector<Node> connections;
     if (m_graph.neighbours(neighbour, level, connections)) {
@@ -517,8 +507,7 @@ bool GraphOperations<Graph>::shrink_neighbours(
       return true;
     }
     std::vector<Node> 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)) {

src/index/hnsw/layer_ops_impl.h

@@ -320,6 +320,12 @@ bool LayerOperations<Graph, Policy>::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<Node>(&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) {

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<int>{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<MockGraph>;
+
+  // 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<int>{1, 0}));
+  } else {
+    assert((g.adjacency[0][2] == std::vector<int>{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<int>{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<int>{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() {