You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ju...@apache.org on 2022/02/04 00:11:53 UTC
[lucene] branch branch_9x updated: LUCENE-10391: Reuse data structures across HnswGraph#searchLevel calls (#641)
This is an automated email from the ASF dual-hosted git repository.
julietibs pushed a commit to branch branch_9x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_9x by this push:
new 612feb9 LUCENE-10391: Reuse data structures across HnswGraph#searchLevel calls (#641)
612feb9 is described below
commit 612feb955dfefa24a2c71e7215e8d19c6d0c77f8
Author: Julie Tibshirani <ju...@apache.org>
AuthorDate: Thu Feb 3 16:00:09 2022 -0800
LUCENE-10391: Reuse data structures across HnswGraph#searchLevel calls (#641)
A couple of the data structures used in HNSW search are pretty large and
expensive to allocate. This commit creates a shared candidates queue and
visited set that are reused across calls to HnswGraph#searchLevel. Now the same
data structures are used for building the entire graph, which can cut down on
allocations during indexing. For graph building it also switches the visited
set to FixedBitSet for better performance.
---
.../codecs/lucene91/Lucene91HnswVectorsReader.java | 6 +-
.../org/apache/lucene/util/hnsw/HnswGraph.java | 118 --------------
.../apache/lucene/util/hnsw/HnswGraphBuilder.java | 15 +-
.../apache/lucene/util/hnsw/HnswGraphSearcher.java | 173 +++++++++++++++++++++
.../org/apache/lucene/util/hnsw/NeighborQueue.java | 5 +
.../org/apache/lucene/util/hnsw/TestHnswGraph.java | 8 +-
.../{TestNeighbors.java => TestNeighborQueue.java} | 13 +-
7 files changed, 208 insertions(+), 130 deletions(-)
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java
index 92135f4..f196160 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene91/Lucene91HnswVectorsReader.java
@@ -46,7 +46,7 @@ import org.apache.lucene.util.Bits;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.IOUtils;
import org.apache.lucene.util.RamUsageEstimator;
-import org.apache.lucene.util.hnsw.HnswGraph;
+import org.apache.lucene.util.hnsw.HnswGraphSearcher;
import org.apache.lucene.util.hnsw.NeighborQueue;
/**
@@ -229,15 +229,15 @@ public final class Lucene91HnswVectorsReader extends KnnVectorsReader {
k = Math.min(k, fieldEntry.size());
OffHeapVectorValues vectorValues = getOffHeapVectorValues(fieldEntry);
- // use a seed that is fixed for the index so we get reproducible results for the same query
NeighborQueue results =
- HnswGraph.search(
+ HnswGraphSearcher.search(
target,
k,
vectorValues,
fieldEntry.similarityFunction,
getGraphValues(fieldEntry),
getAcceptOrds(acceptDocs, fieldEntry));
+
int i = 0;
ScoreDoc[] scoreDocs = new ScoreDoc[Math.min(results.size(), k)];
while (results.size() > 0) {
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java
index 3a98b54..5869395 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraph.java
@@ -19,16 +19,11 @@ package org.apache.lucene.util.hnsw;
import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
-import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.apache.lucene.index.KnnGraphValues;
-import org.apache.lucene.index.RandomAccessVectorValues;
-import org.apache.lucene.index.VectorSimilarityFunction;
import org.apache.lucene.util.ArrayUtil;
-import org.apache.lucene.util.Bits;
-import org.apache.lucene.util.SparseFixedBitSet;
/**
* Hierarchical Navigable Small World graph. Provides efficient approximate nearest neighbor search
@@ -94,119 +89,6 @@ public final class HnswGraph extends KnnGraphValues {
}
/**
- * Searches HNSW graph for the nearest neighbors of a query vector.
- *
- * @param query search query vector
- * @param topK the number of nodes to be returned
- * @param vectors vector values
- * @param graphValues the graph values. May represent the entire graph, or a level in a
- * hierarchical graph.
- * @param acceptOrds {@link Bits} that represents the allowed document ordinals to match, or
- * {@code null} if they are all allowed to match.
- * @return a priority queue holding the closest neighbors found
- */
- public static NeighborQueue search(
- float[] query,
- int topK,
- RandomAccessVectorValues vectors,
- VectorSimilarityFunction similarityFunction,
- KnnGraphValues graphValues,
- Bits acceptOrds)
- throws IOException {
-
- NeighborQueue results;
- int[] eps = new int[] {graphValues.entryNode()};
- for (int level = graphValues.numLevels() - 1; level >= 1; level--) {
- results = searchLevel(query, 1, level, eps, vectors, similarityFunction, graphValues, null);
- eps[0] = results.pop();
- }
- results =
- searchLevel(query, topK, 0, eps, vectors, similarityFunction, graphValues, acceptOrds);
- return results;
- }
-
- /**
- * Searches for the nearest neighbors of a query vector in a given level
- *
- * @param query search query vector
- * @param topK the number of nearest to query results to return
- * @param level level to search
- * @param eps the entry points for search at this level expressed as level 0th ordinals
- * @param vectors vector values
- * @param similarityFunction similarity function
- * @param graphValues the graph values
- * @param acceptOrds {@link Bits} that represents the allowed document ordinals to match, or
- * {@code null} if they are all allowed to match.
- * @return a priority queue holding the closest neighbors found
- */
- static NeighborQueue searchLevel(
- float[] query,
- int topK,
- int level,
- final int[] eps,
- RandomAccessVectorValues vectors,
- VectorSimilarityFunction similarityFunction,
- KnnGraphValues graphValues,
- Bits acceptOrds)
- throws IOException {
-
- int size = graphValues.size();
- // MIN heap, holding the top results
- NeighborQueue results = new NeighborQueue(topK, similarityFunction.reversed);
- // MAX heap, from which to pull the candidate nodes
- NeighborQueue candidates = new NeighborQueue(topK, !similarityFunction.reversed);
- // set of ordinals that have been visited by search on this layer, used to avoid backtracking
- SparseFixedBitSet visited = new SparseFixedBitSet(size);
- for (int ep : eps) {
- if (visited.getAndSet(ep) == false) {
- float score = similarityFunction.compare(query, vectors.vectorValue(ep));
- candidates.add(ep, score);
- if (acceptOrds == null || acceptOrds.get(ep)) {
- results.add(ep, score);
- }
- }
- }
-
- // A bound that holds the minimum similarity to the query vector that a candidate vector must
- // have to be considered.
- BoundsChecker bound = BoundsChecker.create(similarityFunction.reversed);
- if (results.size() >= topK) {
- bound.set(results.topScore());
- }
- while (candidates.size() > 0) {
- // get the best candidate (closest or best scoring)
- float topCandidateScore = candidates.topScore();
- if (bound.check(topCandidateScore)) {
- break;
- }
- int topCandidateNode = candidates.pop();
- graphValues.seek(level, topCandidateNode);
- int friendOrd;
- while ((friendOrd = graphValues.nextNeighbor()) != NO_MORE_DOCS) {
- assert friendOrd < size : "friendOrd=" + friendOrd + "; size=" + size;
- if (visited.getAndSet(friendOrd)) {
- continue;
- }
-
- float score = similarityFunction.compare(query, vectors.vectorValue(friendOrd));
- if (bound.check(score) == false) {
- candidates.add(friendOrd, score);
- if (acceptOrds == null || acceptOrds.get(friendOrd)) {
- if (results.insertWithOverflow(friendOrd, score) && results.size() >= topK) {
- bound.set(results.topScore());
- }
- }
- }
- }
- }
- while (results.size() > topK) {
- results.pop();
- }
- results.setVisitedCount(visited.approximateCardinality());
- return results;
- }
-
- /**
* Returns the {@link NeighborQueue} connected to the given node.
*
* @param level level of the graph
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 041fd18..43efcb3 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
@@ -26,6 +26,7 @@ import java.util.SplittableRandom;
import org.apache.lucene.index.RandomAccessVectorValues;
import org.apache.lucene.index.RandomAccessVectorValuesProducer;
import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.InfoStream;
/**
@@ -51,6 +52,8 @@ public final class HnswGraphBuilder {
private final RandomAccessVectorValues vectorValues;
private final SplittableRandom random;
private final BoundsChecker bound;
+ private final HnswGraphSearcher graphSearcher;
+
final HnswGraph hnsw;
private InfoStream infoStream = InfoStream.getDefault();
@@ -93,6 +96,11 @@ public final class HnswGraphBuilder {
this.random = new SplittableRandom(seed);
int levelOfFirstNode = getRandomGraphLevel(ml, random);
this.hnsw = new HnswGraph(maxConn, levelOfFirstNode);
+ this.graphSearcher =
+ new HnswGraphSearcher(
+ similarityFunction,
+ new NeighborQueue(beamWidth, similarityFunction.reversed == false),
+ new FixedBitSet(vectorValues.size()));
bound = BoundsChecker.create(similarityFunction.reversed);
scratch = new NeighborArray(Math.max(beamWidth, maxConn + 1));
}
@@ -140,17 +148,16 @@ public final class HnswGraphBuilder {
for (int level = nodeLevel; level > curMaxLevel; level--) {
hnsw.addNode(level, node);
}
+
// for levels > nodeLevel search with topk = 1
for (int level = curMaxLevel; level > nodeLevel; level--) {
- candidates =
- HnswGraph.searchLevel(value, 1, level, eps, vectorValues, similarityFunction, hnsw, null);
+ candidates = graphSearcher.searchLevel(value, 1, level, eps, vectorValues, hnsw, null);
eps = new int[] {candidates.pop()};
}
// for levels <= nodeLevel search with topk = beamWidth, and add connections
for (int level = Math.min(nodeLevel, curMaxLevel); level >= 0; level--) {
candidates =
- HnswGraph.searchLevel(
- value, beamWidth, level, eps, vectorValues, similarityFunction, hnsw, null);
+ graphSearcher.searchLevel(value, beamWidth, level, eps, vectorValues, hnsw, null);
eps = candidates.nodes();
hnsw.addNode(level, node);
addDiverseNeighbors(level, node, candidates);
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java
new file mode 100644
index 0000000..71475c1
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java
@@ -0,0 +1,173 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.lucene.util.hnsw;
+
+import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
+
+import java.io.IOException;
+import org.apache.lucene.index.KnnGraphValues;
+import org.apache.lucene.index.RandomAccessVectorValues;
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.BitSet;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.SparseFixedBitSet;
+
+/**
+ * Searches an HNSW graph to find nearest neighbors to a query vector. For more background on the
+ * search algorithm, see {@link HnswGraph}.
+ */
+public final class HnswGraphSearcher {
+ private final VectorSimilarityFunction similarityFunction;
+ /**
+ * Scratch data structures that are used in each {@link #searchLevel} call. These can be expensive
+ * to allocate, so they're cleared and reused across calls.
+ */
+ private final NeighborQueue candidates;
+
+ private final BitSet visited;
+
+ /**
+ * Creates a new graph searcher.
+ *
+ * @param similarityFunction the similarity function to compare vectors
+ * @param candidates max heap that will track the candidate nodes to explore
+ * @param visited bit set that will track nodes that have already been visited
+ */
+ HnswGraphSearcher(
+ VectorSimilarityFunction similarityFunction, NeighborQueue candidates, BitSet visited) {
+ this.similarityFunction = similarityFunction;
+ this.candidates = candidates;
+ this.visited = visited;
+ }
+
+ /**
+ * Searches HNSW graph for the nearest neighbors of a query vector.
+ *
+ * @param query search query vector
+ * @param topK the number of nodes to be returned
+ * @param vectors the vector values
+ * @param similarityFunction the similarity function to compare vectors
+ * @param graphValues the graph values. May represent the entire graph, or a level in a
+ * hierarchical graph.
+ * @param acceptOrds {@link Bits} that represents the allowed document ordinals to match, or
+ * {@code null} if they are all allowed to match.
+ * @return a priority queue holding the closest neighbors found
+ */
+ public static NeighborQueue search(
+ float[] query,
+ int topK,
+ RandomAccessVectorValues vectors,
+ VectorSimilarityFunction similarityFunction,
+ KnnGraphValues graphValues,
+ Bits acceptOrds)
+ throws IOException {
+ HnswGraphSearcher graphSearcher =
+ new HnswGraphSearcher(
+ similarityFunction,
+ new NeighborQueue(topK, similarityFunction.reversed == false),
+ new SparseFixedBitSet(vectors.size()));
+ NeighborQueue results;
+ int[] eps = new int[] {graphValues.entryNode()};
+ for (int level = graphValues.numLevels() - 1; level >= 1; level--) {
+ results = graphSearcher.searchLevel(query, 1, level, eps, vectors, graphValues, null);
+ eps[0] = results.pop();
+ }
+ results = graphSearcher.searchLevel(query, topK, 0, eps, vectors, graphValues, acceptOrds);
+ return results;
+ }
+
+ /**
+ * Searches for the nearest neighbors of a query vector in a given level
+ *
+ * @param query search query vector
+ * @param topK the number of nearest to query results to return
+ * @param level level to search
+ * @param eps the entry points for search at this level expressed as level 0th ordinals
+ * @param vectors vector values
+ * @param graphValues the graph values
+ * @param acceptOrds {@link Bits} that represents the allowed document ordinals to match, or
+ * {@code null} if they are all allowed to match.
+ * @return a priority queue holding the closest neighbors found
+ */
+ NeighborQueue searchLevel(
+ float[] query,
+ int topK,
+ int level,
+ final int[] eps,
+ RandomAccessVectorValues vectors,
+ KnnGraphValues graphValues,
+ Bits acceptOrds)
+ throws IOException {
+ int size = graphValues.size();
+ NeighborQueue results = new NeighborQueue(topK, similarityFunction.reversed);
+ clearScratchState();
+
+ for (int ep : eps) {
+ if (visited.getAndSet(ep) == false) {
+ float score = similarityFunction.compare(query, vectors.vectorValue(ep));
+ candidates.add(ep, score);
+ if (acceptOrds == null || acceptOrds.get(ep)) {
+ results.add(ep, score);
+ }
+ }
+ }
+
+ // A bound that holds the minimum similarity to the query vector that a candidate vector must
+ // have to be considered.
+ BoundsChecker bound = BoundsChecker.create(similarityFunction.reversed);
+ if (results.size() >= topK) {
+ bound.set(results.topScore());
+ }
+ while (candidates.size() > 0) {
+ // get the best candidate (closest or best scoring)
+ float topCandidateScore = candidates.topScore();
+ if (bound.check(topCandidateScore)) {
+ break;
+ }
+ int topCandidateNode = candidates.pop();
+ graphValues.seek(level, topCandidateNode);
+ int friendOrd;
+ while ((friendOrd = graphValues.nextNeighbor()) != NO_MORE_DOCS) {
+ assert friendOrd < size : "friendOrd=" + friendOrd + "; size=" + size;
+ if (visited.getAndSet(friendOrd)) {
+ continue;
+ }
+
+ float score = similarityFunction.compare(query, vectors.vectorValue(friendOrd));
+ if (bound.check(score) == false) {
+ candidates.add(friendOrd, score);
+ if (acceptOrds == null || acceptOrds.get(friendOrd)) {
+ if (results.insertWithOverflow(friendOrd, score) && results.size() >= topK) {
+ bound.set(results.topScore());
+ }
+ }
+ }
+ }
+ }
+ while (results.size() > topK) {
+ results.pop();
+ }
+ results.setVisitedCount(visited.approximateCardinality());
+ return results;
+ }
+
+ private void clearScratchState() {
+ candidates.clear();
+ visited.clear(0, visited.length());
+ }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborQueue.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborQueue.java
index 6f08a7b..e06d120 100644
--- a/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborQueue.java
+++ b/lucene/core/src/java/org/apache/lucene/util/hnsw/NeighborQueue.java
@@ -115,6 +115,11 @@ public class NeighborQueue {
return NumericUtils.sortableIntToFloat((int) (order.apply(heap.top()) >> 32));
}
+ public void clear() {
+ heap.clear();
+ visitedCount = 0;
+ }
+
public int visitedCount() {
return visitedCount;
}
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswGraph.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswGraph.java
index de21bef..5142b78 100644
--- a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswGraph.java
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestHnswGraph.java
@@ -162,7 +162,7 @@ public class TestHnswGraph extends LuceneTestCase {
HnswGraph hnsw = builder.build(vectors);
// run some searches
NeighborQueue nn =
- HnswGraph.search(
+ HnswGraphSearcher.search(
new float[] {1, 0},
10,
vectors.randomAccess(),
@@ -201,7 +201,7 @@ public class TestHnswGraph extends LuceneTestCase {
// the first 10 docs must not be deleted to ensure the expected recall
Bits acceptOrds = createRandomAcceptOrds(10, vectors.size);
NeighborQueue nn =
- HnswGraph.search(
+ HnswGraphSearcher.search(
new float[] {1, 0},
10,
vectors.randomAccess(),
@@ -234,7 +234,7 @@ public class TestHnswGraph extends LuceneTestCase {
acceptOrds.set(i);
}
NeighborQueue nn =
- HnswGraph.search(
+ HnswGraphSearcher.search(
new float[] {1, 0},
10,
vectors.randomAccess(),
@@ -383,7 +383,7 @@ public class TestHnswGraph extends LuceneTestCase {
for (int i = 0; i < 100; i++) {
float[] query = randomVector(random(), dim);
NeighborQueue actual =
- HnswGraph.search(query, 100, vectors, similarityFunction, hnsw, acceptOrds);
+ HnswGraphSearcher.search(query, 100, vectors, similarityFunction, hnsw, acceptOrds);
while (actual.size() > topK) {
actual.pop();
}
diff --git a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighbors.java b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborQueue.java
similarity index 91%
rename from lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighbors.java
rename to lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborQueue.java
index ef207c9..ef432b3 100644
--- a/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighbors.java
+++ b/lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborQueue.java
@@ -19,7 +19,7 @@ package org.apache.lucene.util.hnsw;
import org.apache.lucene.tests.util.LuceneTestCase;
-public class TestNeighbors extends LuceneTestCase {
+public class TestNeighborQueue extends LuceneTestCase {
public void testNeighborsProduct() {
// make sure we have the sign correct
@@ -67,6 +67,17 @@ public class TestNeighbors extends LuceneTestCase {
assertEquals(100, nn.visitedCount());
}
+ public void testClear() {
+ NeighborQueue nn = new NeighborQueue(2, false);
+ nn.add(1, 1.1f);
+ nn.add(2, -2.2f);
+ nn.setVisitedCount(42);
+ nn.clear();
+
+ assertEquals(0, nn.size());
+ assertEquals(0, nn.visitedCount());
+ }
+
public void testMaxSizeQueue() {
NeighborQueue nn = new NeighborQueue(2, false);
nn.add(1, 1);