Skip to content

Commit

Permalink
Optimize DFS while marking connected components (#14022)
Browse files Browse the repository at this point in the history
  • Loading branch information
viswanathk authored Jan 6, 2025
1 parent d92d045 commit 5460da8
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions lucene/core/src/java/org/apache/lucene/util/hnsw/HnswUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand All @@ -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);
Expand Down

0 comments on commit 5460da8

Please sign in to comment.