我想使用 Lucene 来运行最近邻搜索。我在 JVM 11 上使用 Lucene 9.0.0。我没有找到太多文档,主要尝试使用现有的 tests 将内容拼凑在一起。
我编写了一个小测试,准备了
HnswGraph
但到目前为止,搜索没有产生预期的结果。我设置了一组随机向量并添加了一个非常接近我的搜索目标的最终向量(0.99f,0.01f)。
不幸的是,搜索从未返回预期值。我不确定我的错误在哪里。我认为这可能与插入和文档 ID 顺序有关。
也许更熟悉 lucene 的人可以提供一些反馈。我的做法正确吗?我使用文档只是为了持久化。
HnswGraphBuilder builder = new HnswGraphBuilder(vectors, similarityFunction, maxConn, beamWidth, seed);
HnswGraph hnsw = builder.build(vectors);
// Run a search
NeighborQueue nn = HnswGraph.search(
new float[] { 1, 0 },
10,
10,
vectors.randomAccess(), // ? Why do I need to specify the graph values again?
similarityFunction, // ? Why can I specify a different similarityFunction for search. Should that not be the same that was used for graph creation?
hnsw,
null,
new SplittableRandom(RandomUtils.nextLong()));
完整的测试源可以在这里找到: https://gist.github.com/Jotschi/cea21a72412bcba80c46b967e9c52b0f
由于对于像我这样寻找 lucene hnsw 示例的人来说这是一个热门搜索(那里没有太多),因此从 Lucene 9.6 开始,它看起来像这样,没有引入更高级别的 Lucene 类:
// Create a random vector universe
var vectorDimensions = 1500;
var universeSize = 2_000;
var universe = new ArrayList<float[]>(universeSize);
for (var i = 0; i < universeSize; i++) {
universe.add(randomVector(vectorDimensions));
}
// construct a HNSW graph of the universe
System.out.println("Constructing HNSW graph...");
var ravv = new ListRandomAccessVectorValues(universe, vectorDimensions);
var builder = HnswGraphBuilder.create(ravv, VectorEncoding.FLOAT32, similarityFunction, 16, 100, random.nextInt());
var hnsw = builder.build(ravv.copy());
// search for the nearest neighbors of a random vector
var queryVector = randomVector(vectorDimensions);
System.out.println("Searching for top 10 neighbors of a random vector");
var nn = HnswGraphSearcher.search(queryVector, 10, ravv.copy(), VectorEncoding.FLOAT32, similarityFunction, hnsw, null, Integer.MAX_VALUE);
for (var i : nn.nodes()) {
var neighbor = universe.get(i);
var similarity = similarityFunction.compare(queryVector, neighbor);
System.out.printf(" ordinal %d (similarity: %s)%n", i, similarity);
}
ListRandomAccessVectorValues 的实现非常简单,可以在这里找到:https://github.com/jbellis/hnswdemo
正如您所注意到的,API 有点不稳定,虽然您可以指定与构建器不同的编码或搜索相似性,但您显然会得到无意义的结果。
必须将向量值传递给搜索函数的原因是,为了提高效率,索引本身不存储向量本身的副本,而只存储值提供程序中的 int 偏移量。
我设法让它工作。
我现在使用
HnswGraph
,而不是直接使用 LeafReader#searchNearestVectors
API。在调试时,我注意到 Lucene90HnswVectorsWriter
例如使用 HnswGraph
API 调用额外的步骤。我认为这样做是为了创建插入向量和文档 ID 之间的相关性。我使用 HnswGraph#search
检索到的 nodeId 从未与向量 Id 匹配。我不知道是否需要额外的步骤来设置图表,或者之后是否需要以某种方式创建相关性。
好消息是
LeafReader#searchNearestVectors
方法有效。我已经更新了示例,现在也使用了 Lucene 文档。
@Test
public void testWriteAndQueryIndex() throws IOException {
// Persist and read the data
try (MMapDirectory dir = new MMapDirectory(indexPath)) {
// Write index
int indexedDoc = writeIndex(dir, vectors);
// Read index
readAndQuery(dir, vectors, indexedDoc);
}
}
[0.97|0.02]的向量7非常接近搜索查询目标[0.98|0.01]。
Test vectors:
0 => [0.13|0.37]
1 => [0.99|0.49]
2 => [0.98|0.57]
3 => [0.23|0.64]
4 => [0.72|0.92]
5 => [0.08|0.74]
6 => [0.50|0.27]
7 => [0.97|0.02]
8 => [0.90|0.21]
9 => [0.89|0.09]
10 => [0.11|0.95]
Doc Based Search:
Searching for NN of [0.98 | 0.01]
TotalHits: 11
7 => [0.97|0.02]
9 => [0.89|0.09]
完整示例: https://gist.github.com/Jotschi/d8a91758c84203d172f818c8be4964e4
解决此问题的另一种方法是使用 KnnVectorQuery。
try (IndexReader reader = DirectoryReader.open(dir)) {
IndexSearcher searcher = new IndexSearcher(reader);
System.out.println("Query: [" + String.format("%.2f", queryVector[0]) + ", " + String.format("%.2f", queryVector[1]) + "]");
TopDocs results = searcher.search(new KnnVectorQuery("field", queryVector, 3), 10);
System.out.println("Hits: " + results.totalHits);
for (ScoreDoc sdoc : results.scoreDocs) {
Document doc = reader.document(sdoc.doc);
StoredField idField = (StoredField) doc.getField("id");
System.out.println("Found: " + idField.numericValue() + " = " + String.format("%.1f", sdoc.score));
}
}
完整示例: https://gist.github.com/Jotschi/7d599dff331d75a3bdd02e62f65abfba