Optimize DFS while marking connected components (#14022) (#14105)

* Optimize DFS while marking connected components (#14022) * Add CHANGES.txt entry for HNSW DFS Optimization #14022
apache · Jan 7, 2025 · c795d74 · c795d74
1 parent 7146046
commit c795d74
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
@@ -141,6 +141,9 @@ Optimizations
 * GITHUB#14040: Specialized top-level DisjunctionMaxQuery evaluation when the
   tie break multiplier is 0. (Adrien Grand)
 
+* GITHUB#14022: Optimize DFS marking of connected components in HNSW by reducing stack depth, improving performance
+  and reducing allocations. (Viswanath Kuchibhotla)
+
 Bug Fixes
 ---------------------
 * GITHUB#13832: Fixed an issue where the DefaultPassageFormatter.format method did not format passages as intended

diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswUtil.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswUtil.java
@@ -29,6 +29,7 @@
 import org.apache.lucene.index.FilterLeafReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.internal.hppc.IntHashSet;
 import org.apache.lucene.util.FixedBitSet;
 
 /** Utilities for use in tests involving HNSW graphs */
@@ -105,7 +106,9 @@ static List<Component> components(
     } else {
       entryPoint = connectedNodes.nextSetBit(0);
     }
-    components.add(new Component(entryPoint, total));
+    if (total > 0) {
+      components.add(new Component(entryPoint, total));
+    }
     if (level == 0) {
       int nextClear = nextClearBit(connectedNodes, 0);
       while (nextClear != NO_MORE_DOCS) {
@@ -163,6 +166,10 @@ private static Component markRooted(
       throws IOException {
     // Start at entry point and search all nodes on this level
     // System.out.println("markRooted level=" + level + " entryPoint=" + entryPoint);
+    if (connectedNodes.get(entryPoint)) {
+      return new Component(entryPoint, 0);
+    }
+    IntHashSet nodesInStack = new IntHashSet();
     Deque<Integer> stack = new ArrayDeque<>();
     stack.push(entryPoint);
     int count = 0;
@@ -178,7 +185,10 @@ private static Component markRooted(
       int friendCount = 0;
       while ((friendOrd = hnswGraph.nextNeighbor()) != NO_MORE_DOCS) {
         ++friendCount;
-        stack.push(friendOrd);
+        if (connectedNodes.get(friendOrd) == false && nodesInStack.contains(friendOrd) == false) {
+          stack.push(friendOrd);
+          nodesInStack.add(friendOrd);
+        }
       }
       if (friendCount < maxConn && notFullyConnected != null) {
         notFullyConnected.set(node);