如何使用Lucene HnswGraph运行最近的邻居搜索

Question

我想用Lucene进行最近的邻居搜索。我在JVM 11上使用Lucene9.0.0，我没有找到太多的文档，我主要尝试使用现有的测试将所有东西拼凑在一起。

我编写了一个小测试，它准备了一个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

Answer 1

解决这一问题的另一种方法是使用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

Answer 2

我成功地完成了这个任务。

我现在不用直接使用HnswGraph API，而是使用LeafReader#searchNearestVectors。在调试时，我注意到，例如，Lucene90HnswVectorsWriter使用HnswGraph API调用额外的步骤。我假设这样做是为了在插入的向量和文档Ids之间创建关联。我使用nodeIds检索的HnswGraph#search从未与匹配的向量Ids匹配。我不知道是否需要额外的步骤来设置图形，或者相关性是否需要在某种程度上被创建。

好消息是，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);
    }
}

向量7的0.97\0.0 2非常接近搜索查询目标0.98\0.0 1。

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