You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by pa...@apache.org on 2025/09/18 06:03:18 UTC
(lucene) branch main updated: Refactoring HNSWGraphBuilder's API and adding more comments about concurrency (#15184)
This is an automated email from the ASF dual-hosted git repository.
patrickz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new c243698f573 Refactoring HNSWGraphBuilder's API and adding more comments about concurrency (#15184)
c243698f573 is described below
commit c243698f573943cba9cec907cca732c781cec98f
Author: Patrick Zhai <zh...@users.noreply.github.com>
AuthorDate: Thu Sep 18 01:03:11 2025 -0500
Refactoring HNSWGraphBuilder's API and adding more comments about concurrency (#15184)
---
lucene/CHANGES.txt | 4 +-
.../codecs/lucene99/Lucene99HnswVectorsWriter.java | 6 +-
.../org/apache/lucene/util/hnsw/HnswBuilder.java | 10 ++-
.../util/hnsw/HnswConcurrentMergeBuilder.java | 19 ++----
.../apache/lucene/util/hnsw/HnswGraphBuilder.java | 79 ++++++++++++----------
.../util/hnsw/InitializedHnswGraphBuilder.java | 8 ---
.../lucene/util/hnsw/MergingHnswGraphBuilder.java | 5 +-
7 files changed, 65 insertions(+), 66 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index be887da77c7..43b64c2977a 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -134,7 +134,9 @@ New Features
Improvements
---------------------
-# GITHUB#15148: Add support uint8 distance and allow 8 bit scalar quantization (Trevor McCulloch)
+* GITHUB#15148: Add support uint8 distance and allow 8 bit scalar quantization (Trevor McCulloch)
+
+* GITHUB#15184: Refactoring internal HNSWGraphBuilder's APIs and avoid creating new scorer for each call (Patrick Zhai)
Optimizations
---------------------
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
index 49973956192..32e745f6c4f 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsWriter.java
@@ -59,7 +59,6 @@ import org.apache.lucene.util.hnsw.IncrementalHnswGraphMerger;
import org.apache.lucene.util.hnsw.NeighborArray;
import org.apache.lucene.util.hnsw.OnHeapHnswGraph;
import org.apache.lucene.util.hnsw.RandomVectorScorerSupplier;
-import org.apache.lucene.util.hnsw.UpdateableRandomVectorScorer;
import org.apache.lucene.util.packed.DirectMonotonicWriter;
/**
@@ -586,7 +585,6 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
private int lastDocID = -1;
private int node = 0;
private final FlatFieldVectorsWriter<T> flatFieldVectorsWriter;
- private UpdateableRandomVectorScorer scorer;
@SuppressWarnings("unchecked")
static FieldWriter<?> create(
@@ -642,7 +640,6 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
(List<float[]>) flatFieldVectorsWriter.getVectors(),
fieldInfo.getVectorDimension()));
};
- this.scorer = scorerSupplier.scorer();
hnswGraphBuilder =
HnswGraphBuilder.create(scorerSupplier, M, beamWidth, HnswGraphBuilder.randSeed);
hnswGraphBuilder.setInfoStream(infoStream);
@@ -658,8 +655,7 @@ public final class Lucene99HnswVectorsWriter extends KnnVectorsWriter {
+ "\" appears more than once in this document (only one value is allowed per field)");
}
flatFieldVectorsWriter.addValue(docID, vectorValue);
- scorer.setScoringOrdinal(node);
- hnswGraphBuilder.addGraphNode(node, scorer);
+ hnswGraphBuilder.addGraphNode(node);
node++;
lastDocID = docID;
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java
index aa27525b7f1..38109c9c95e 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswBuilder.java
@@ -18,6 +18,7 @@
package org.apache.lucene.util.hnsw;
import java.io.IOException;
+import org.apache.lucene.internal.hppc.IntHashSet;
import org.apache.lucene.util.InfoStream;
/**
@@ -34,16 +35,21 @@ public interface HnswBuilder {
*/
OnHeapHnswGraph build(int maxOrd) throws IOException;
- /** Inserts a doc with vector value to the graph */
+ /** Inserts a doc with a vector value to the graph */
void addGraphNode(int node) throws IOException;
+ /**
+ * Inserts a doc with a vector value to the graph, searching on level 0 with provided entry points
+ */
+ void addGraphNode(int node, IntHashSet eps) throws IOException;
+
/** Set info-stream to output debugging information */
void setInfoStream(InfoStream infoStream);
OnHeapHnswGraph getGraph();
/**
- * Once this method is called no further updates to the graph are accepted (addGraphNode will
+ * Once this method is called, no further updates to the graph are accepted (addGraphNode will
* throw IllegalStateException). Final modifications to the graph (eg patching up disconnected
* components, re-ordering node ids for better delta compression) may be triggered, so callers
* should expect this call to take some time.
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java
index 415e6cd2314..9fbf69558e5 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswConcurrentMergeBuilder.java
@@ -25,6 +25,7 @@ import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
+import org.apache.lucene.internal.hppc.IntHashSet;
import org.apache.lucene.search.TaskExecutor;
import org.apache.lucene.util.BitSet;
import org.apache.lucene.util.FixedBitSet;
@@ -98,6 +99,11 @@ public class HnswConcurrentMergeBuilder implements HnswBuilder {
throw new UnsupportedOperationException("This builder is for merge only");
}
+ @Override
+ public void addGraphNode(int node, IntHashSet eps) throws IOException {
+ throw new UnsupportedOperationException("This builder is for merge only");
+ }
+
@Override
public void setInfoStream(InfoStream infoStream) {
this.infoStream = infoStream;
@@ -142,7 +148,6 @@ public class HnswConcurrentMergeBuilder implements HnswBuilder {
private final BitSet initializedNodes;
private int batchSize = DEFAULT_BATCH_SIZE;
- private final UpdateableRandomVectorScorer scorer;
private ConcurrentMergeWorker(
RandomVectorScorerSupplier scorerSupplier,
@@ -163,7 +168,6 @@ public class HnswConcurrentMergeBuilder implements HnswBuilder {
new NeighborQueue(beamWidth, true), hnswLock, new FixedBitSet(hnsw.maxNodeId() + 1)));
this.workProgress = workProgress;
this.initializedNodes = initializedNodes;
- this.scorer = scorerSupplier.scorer();
}
/**
@@ -192,21 +196,12 @@ public class HnswConcurrentMergeBuilder implements HnswBuilder {
}
}
- @Override
- public void addGraphNode(int node, UpdateableRandomVectorScorer scorer) throws IOException {
- if (initializedNodes != null && initializedNodes.get(node)) {
- return;
- }
- super.addGraphNode(node, scorer);
- }
-
@Override
public void addGraphNode(int node) throws IOException {
if (initializedNodes != null && initializedNodes.get(node)) {
return;
}
- scorer.setScoringOrdinal(node);
- addGraphNode(node, scorer);
+ super.addGraphNode(node);
}
}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
index cd437069ff2..cd54443ab76 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphBuilder.java
@@ -39,6 +39,10 @@ import org.apache.lucene.util.hnsw.HnswUtil.Component;
/**
* Builder for HNSW graph. See {@link HnswGraph} for a gloss on the algorithm and the meaning of the
* hyper-parameters.
+ *
+ * <p>Thread-safety: This class is NOT thread safe, it cannot be shared across threads, however, it
+ * IS safe for multiple HnswGraphBuilder to build the same graph, if the graph's size is known in
+ * the beginning (like when doing merge)
*/
public class HnswGraphBuilder implements HnswBuilder {
@@ -64,7 +68,7 @@ public class HnswGraphBuilder implements HnswBuilder {
private final double ml;
private final SplittableRandom random;
- protected final RandomVectorScorerSupplier scorerSupplier;
+ private final UpdateableRandomVectorScorer scorer;
private final HnswGraphSearcher graphSearcher;
private final GraphBuilderKnnCollector entryCandidates; // for upper levels of graph search
private final GraphBuilderKnnCollector
@@ -144,8 +148,8 @@ public class HnswGraphBuilder implements HnswBuilder {
throw new IllegalArgumentException("beamWidth must be positive");
}
this.M = hnsw.maxConn();
- this.scorerSupplier =
- Objects.requireNonNull(scorerSupplier, "scorer supplier must not be null");
+ this.scorer =
+ Objects.requireNonNull(scorerSupplier, "scorer supplier must not be null").scorer();
// normalization factor for level generation; currently not configurable
this.ml = M == 1 ? 1 : 1 / Math.log(1.0 * M);
this.random = new SplittableRandom(seed);
@@ -196,10 +200,8 @@ public class HnswGraphBuilder implements HnswBuilder {
if (infoStream.isEnabled(HNSW_COMPONENT)) {
infoStream.message(HNSW_COMPONENT, "addVectors [" + minOrd + " " + maxOrd + ")");
}
- UpdateableRandomVectorScorer scorer = scorerSupplier.scorer();
for (int node = minOrd; node < maxOrd; node++) {
- scorer.setScoringOrdinal(node);
- addGraphNode(node, scorer);
+ addGraphNode(node);
if ((node % 10000 == 0) && infoStream.isEnabled(HNSW_COMPONENT)) {
t = printGraphBuildStatus(node, start, t);
}
@@ -210,10 +212,24 @@ public class HnswGraphBuilder implements HnswBuilder {
addVectors(0, maxOrd);
}
- public void addGraphNode(int node, UpdateableRandomVectorScorer scorer) throws IOException {
- addGraphNodeInternal(node, scorer, null);
- }
-
+ /**
+ * Note: this implementation is thread safe when the graph size is fixed (e.g. when merging) The
+ * process of adding a node is roughly: 1. Add the node to all levels from top to the bottom, but
+ * do not connect it to any other node, nor try to promote itself to an entry node before the
+ * connection is done. (Unless the graph is empty and this is the first node, in that case we set
+ * the entry node and return) 2. Do the search from top to bottom, remember all the possible
+ * neighbours on each level the node is on. 3. Add the neighbor to the node from bottom to top
+ * level. When adding the neighbour, we always add all the outgoing links first before adding an
+ * incoming link such that when a search visits this node, it can always find a way out 4. If the
+ * node has a level that is less or equal to the graph's max level, then we're done here. If the
+ * node has a level larger than the graph's max level, then we need to promote the node as the
+ * entry node. If, while we add the node to the graph, the entry node has changed (which means the
+ * graph level has changed as well), we need to reinsert the node to the newly introduced levels
+ * (repeating step 2,3 for new levels) and again try to promote the node to entry node.
+ *
+ * @param eps0 If specified, we will use it as the entry points of search on level 0, is useful
+ * when you have some prior knowledge, e.g. in {@link MergingHnswGraphBuilder}
+ */
private void addGraphNodeInternal(int node, UpdateableRandomVectorScorer scorer, IntHashSet eps0)
throws IOException {
if (frozen) {
@@ -224,7 +240,8 @@ public class HnswGraphBuilder implements HnswBuilder {
for (int level = nodeLevel; level >= 0; level--) {
hnsw.addNode(level, node);
}
- // then promote itself as entry node if entry node is not set
+ // then promote itself as entry node if entry node is not set (this is the first ever node of
+ // the graph)
if (hnsw.trySetNewEntryNode(node, nodeLevel)) {
return;
}
@@ -235,8 +252,12 @@ public class HnswGraphBuilder implements HnswBuilder {
int curMaxLevel;
do {
curMaxLevel = hnsw.numLevels() - 1;
- // NOTE: the entry node and max level may not be paired, but because we get the level first
+ // NOTE: the entry node and max level are not retrieved synchronously, which could lead to a
+ // situation where
+ // the entry node's level is different from the graph's max level, but because we get the
+ // level first,
// we ensure that the entry node we get later will always exist on the curMaxLevel
+ // e.g., curMaxLevel <= entryNode.level
int[] eps = new int[] {hnsw.entryNode()};
// we first do the search from top to bottom
@@ -271,15 +292,21 @@ public class HnswGraphBuilder implements HnswBuilder {
}
lowestUnsetLevel += scratchPerLevel.length;
assert lowestUnsetLevel == Math.min(nodeLevel, curMaxLevel) + 1;
- if (lowestUnsetLevel > nodeLevel) {
+ if (lowestUnsetLevel == nodeLevel + 1) {
+ // we have already set all the levels we need for this node
return;
}
assert lowestUnsetLevel == curMaxLevel + 1 && nodeLevel > curMaxLevel;
+ // The node's level is higher than the graph's max level, so we need to
+ // try to promote this node as the graph's entry node
if (hnsw.tryPromoteNewEntryNode(node, nodeLevel, curMaxLevel)) {
return;
}
+ // If we're not able to promote, it means the graph must have already changed
+ // and has a new max level and some other entry node
if (hnsw.numLevels() == curMaxLevel + 1) {
- // This should never happen if all the calculations are correct
+ // This is an impossible situation, if happens, then something above is
+ // not hold
throw new IllegalStateException(
"We're not able to promote node "
+ node
@@ -294,31 +321,12 @@ public class HnswGraphBuilder implements HnswBuilder {
@Override
public void addGraphNode(int node) throws IOException {
- /*
- * Note: this implementation is thread safe when graph size is fixed (e.g. when merging)
- * The process of adding a node is roughly:
- * 1. Add the node to all level from top to the bottom, but do not connect it to any other node,
- * nor try to promote itself to an entry node before the connection is done. (Unless the graph is empty
- * and this is the first node, in that case we set the entry node and return)
- * 2. Do the search from top to bottom, remember all the possible neighbours on each level the node
- * is on.
- * 3. Add the neighbor to the node from bottom to top level, when adding the neighbour,
- * we always add all the outgoing links first before adding incoming link such that
- * when a search visits this node, it can always find a way out
- * 4. If the node has level that is less or equal to graph level, then we're done here.
- * If the node has level larger than graph level, then we need to promote the node
- * as the entry node. If, while we add the node to the graph, the entry node has changed
- * (which means the graph level has changed as well), we need to reinsert the node
- * to the newly introduced levels (repeating step 2,3 for new levels) and again try to
- * promote the node to entry node.
- */
- UpdateableRandomVectorScorer scorer = scorerSupplier.scorer();
scorer.setScoringOrdinal(node);
addGraphNodeInternal(node, scorer, null);
}
- public void addGraphNodeWithEps(int node, IntHashSet eps0) throws IOException {
- UpdateableRandomVectorScorer scorer = scorerSupplier.scorer();
+ @Override
+ public void addGraphNode(int node, IntHashSet eps0) throws IOException {
scorer.setScoringOrdinal(node);
addGraphNodeInternal(node, scorer, eps0);
}
@@ -486,7 +494,6 @@ public class HnswGraphBuilder implements HnswBuilder {
// while linking
GraphBuilderKnnCollector beam = new GraphBuilderKnnCollector(2);
int[] eps = new int[1];
- UpdateableRandomVectorScorer scorer = scorerSupplier.scorer();
for (Component c : components) {
if (c != c0) {
if (c.start() == NO_MORE_DOCS) {
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/InitializedHnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/InitializedHnswGraphBuilder.java
index c4c790c7025..7dff036ddde 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/InitializedHnswGraphBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/InitializedHnswGraphBuilder.java
@@ -98,14 +98,6 @@ public final class InitializedHnswGraphBuilder extends HnswGraphBuilder {
this.initializedNodes = initializedNodes;
}
- @Override
- public void addGraphNode(int node, UpdateableRandomVectorScorer scorer) throws IOException {
- if (initializedNodes.get(node)) {
- return;
- }
- super.addGraphNode(node, scorer);
- }
-
@Override
public void addGraphNode(int node) throws IOException {
if (initializedNodes.get(node)) {
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java
index e7c3908d891..08366927b24 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/MergingHnswGraphBuilder.java
@@ -84,7 +84,8 @@ public final class MergingHnswGraphBuilder extends HnswGraphBuilder {
* @param ordMaps the ordinal maps for the graphs
* @param totalNumberOfVectors the total number of vectors in the new graph, this should include
* all vectors expected to be added to the graph in the future
- * @param initializedNodes the nodes will be initialized through the merging
+ * @param initializedNodes the nodes will be initialized through the merging, if null, all nodes
+ * should be already initialized after {@link #updateGraph(HnswGraph, int[])} being called
* @return a new HnswGraphBuilder that is initialized with the provided HnswGraph
* @throws IOException when reading the graph fails
*/
@@ -172,7 +173,7 @@ public final class MergingHnswGraphBuilder extends HnswGraphBuilder {
}
}
}
- addGraphNodeWithEps(ordMapS[u], eps);
+ addGraphNode(ordMapS[u], eps);
}
}
}