From 1d0539decb0200ae6cd66454e60bbb188cf1477d Mon Sep 17 00:00:00 2001
From: Viswanath Kuchibhotla <vishwanath_243@yahoo.com>
Date: Wed, 27 Nov 2024 22:35:50 +0530
Subject: [PATCH] Optimize DFS while marking connected components

---
 .../src/java/org/apache/lucene/util/hnsw/HnswUtil.java   | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

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
index a4b1d0c7c53..10aee15c205 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswUtil.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswUtil.java
@@ -163,6 +163,10 @@ public class HnswUtil {
       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);
+    }
+    FixedBitSet nodesInStack = new FixedBitSet(hnswGraph.size());
     Deque<Integer> stack = new ArrayDeque<>();
     stack.push(entryPoint);
     int count = 0;
@@ -178,7 +182,10 @@ public class HnswUtil {
       int friendCount = 0;
       while ((friendOrd = hnswGraph.nextNeighbor()) != NO_MORE_DOCS) {
         ++friendCount;
-        stack.push(friendOrd);
+        if (connectedNodes.get(friendOrd) == false && nodesInStack.get(friendOrd) == false) {
+          stack.push(friendOrd);
+          nodesInStack.set(friendOrd);
+        }
       }
       if (friendCount < maxConn && notFullyConnected != null) {
         notFullyConnected.set(node);