Howto run Nearest Neighbour Search with Lucene HnswGraph

user_644326 0
Viewed 362

I would like to use Lucene to run a nearest neighbour search. I'm using Lucene 9.0.0 on JVM 11. I did not find much documentation and mainly tried to piece things together using the existing tests.

I wrote a small test which prepares a HnswGraph but so far the search does not yield the expected result. I setup a set of random vectors and add a final vector (0.99f,0.01f) which is very close to my search target. The search unfortunately never returns the expected value. I'm not sure where my error is. I assume it may be related with the insert and document id order.

Maybe someone who is more familar with lucene might be able to provide some feedback. Is my approach correct? I'm using Documents only for persistence.

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()));

The whole test source can be found here: https://gist.github.com/Jotschi/cea21a72412bcba80c46b967e9c52b0f

lucene
2 Answers

I managed to get this working.

Instead of using the HnswGraph API directly I now use LeafReader#searchNearestVectors. While debugging I noticed that the Lucene90HnswVectorsWriter for example invokes extra steps using the HnswGraph API. I assume this is done to create a correlation between inserted vectors and document Ids. The nodeIds I retrieved using a HnswGraph#search never matched up with the matched up with the vector Ids. I don't know whether extra steps are needed to setup the graph or whether the correlation needs to be created afterwards somehow.

The good news is that the LeafReader#searchNearestVectors method works. I have updated the example which now also makes use of the Lucene documents.

@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);
    }
}

Vector 7 with [0.97|0.02] is very close to the search query target [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]

Full example: https://gist.github.com/Jotschi/d8a91758c84203d172f818c8be4964e4

user_6443260

Another way to solve this is to use the 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));
    }
}

Full example: https://gist.github.com/Jotschi/7d599dff331d75a3bdd02e62f65abfba

user_6443260
Related