fix single node scenario (#4)

apache · Jan 12, 2020 · e4d348b · e4d348b
1 parent b4a8979
commit e4d348b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HNSWGraphReader.java
@@ -30,6 +30,7 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.VectorValues;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.util.IntsRef;
 
 /** Executes approximate nearest neighbor search on per-reader {@link HNSWGraph}.
  * This also caches built {@link HNSWGraph}s for repeated use. */
@@ -117,8 +118,14 @@ public static HNSWGraph load(String field, VectorValues.DistanceFunction distFun
  int maxLevel = graphValues.getMaxLevel();
  hnsw.ensureLevel(maxLevel);
  for (int l = 0; l <= maxLevel; l++) {
- for (int friend : graphValues.getFriends(l).ints) {
- hnsw.connectNodes(l, doc, friend);
+ final IntsRef friends = graphValues.getFriends(l);
+ if (friends.length > 0) {
+ for (int friend : friends.ints) {
+ hnsw.connectNodes(l, doc, friend);
+ }
+ } else {
+ /// if there's single one node (thus no friends), the node should be added to all layers
+ hnsw.addNode(l, doc);
  }
  }
  }

diff --git a/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java b/lucene/core/src/test/org/apache/lucene/index/TestKnnGraph.java
@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Set;
 
+import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
 import org.apache.lucene.document.VectorField;
@@ -81,6 +82,21 @@ public void testSingleDocument() throws Exception {
  }
  }
 
+ public void testSingleDocRecall() throws Exception {
+ try (Directory dir = newDirectory();
+ IndexWriter iw = new IndexWriter(dir, newIndexWriterConfig(null)
+ .setCodec(Codec.forName("Lucene90")))) {
+ float[][] values = new float[][]{new float[]{0, 1, 2}};
+ add(iw, 0, values[0]);
+ assertConsistentGraph(iw, values);
+
+ iw.commit();
+ assertConsistentGraph(iw, values);
+
+ assertRecall(dir, 0, values[0]);
+ }
+ }
+
  /**
  * Verify that the graph properties are preserved when merging
  */
@@ -294,4 +310,23 @@ private void add(IndexWriter iw, int id, float[] vector) throws IOException {
  iw.addDocument(doc);
  }
 
+ private void assertRecall(Directory dir, int expectDocId, float[] value) throws IOException {
+ try (IndexReader reader = DirectoryReader.open(dir)) {
+ IndexSearcher searcher = new IndexSearcher(reader);
+ KnnGraphQuery query = new KnnGraphQuery(KNN_GRAPH_FIELD, value);
+ TopDocs result = searcher.search(query, 1);
+
+ int recallCnt = 0;
+ for (LeafReaderContext ctx : reader.leaves()) {
+ VectorValues vector = ctx.reader().getVectorValues(KNN_GRAPH_FIELD);
+ if (vector.seek(result.scoreDocs[0].doc)) {
+ ++recallCnt;
+ assertEquals(expectDocId, vector.docID());
+ assertEquals(0, Arrays.compare(value, vector.vectorValue()));
+ }
+ }
+ assertEquals(1, recallCnt);
+ }
+ }
+
 }