From 5232d6aa81673328745f3fd9414c7f409e877aa9 Mon Sep 17 00:00:00 2001
From: Mehmet Dogan <mehmet@dogan.io>
Date: Fri, 9 Feb 2024 15:43:24 +0300
Subject: [PATCH] Fix GraphIndexBuilder.reconnectOrphanedNodes

`connectToClosestNeighbor()` method looks like connecting a `neighborNode` to itself
with the similarity score between `node` and `neighborNode`.

If I'm getting this logic correctly, it should connect `neighborNode` to `node` instead.

Closes #214
---
 .../jvector/graph/GraphIndexBuilder.java      |  2 +-
 .../jvector/graph/GraphIndexBuilderTest.java  | 79 +++++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderTest.java

diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java
index 1c045012c..1c47cc7b5 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java
@@ -291,7 +291,7 @@ private boolean connectToClosestNeighbor(int node, NodeArray neighbors, Set<Inte
             var neighborNode = neighbors.node[i];
             var neighborScore = neighbors.score[i];
             if (connectionTargets.add(neighborNode)) {
-                graph.getNeighbors(neighborNode).insertNotDiverse(neighborNode, neighborScore, true);
+                graph.getNeighbors(neighborNode).insertNotDiverse(node, neighborScore, true);
                 return true;
             }
         }
diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderTest.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderTest.java
new file mode 100644
index 000000000..6375b6972
--- /dev/null
+++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/graph/GraphIndexBuilderTest.java
@@ -0,0 +1,79 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 io.github.jbellis.jvector.graph;
+
+import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
+import io.github.jbellis.jvector.LuceneTestCase;
+import io.github.jbellis.jvector.TestUtil;
+import io.github.jbellis.jvector.disk.SimpleMappedReader;
+import io.github.jbellis.jvector.vector.VectorEncoding;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.function.Supplier;
+
+import static io.github.jbellis.jvector.graph.GraphIndexTestCase.createRandomFloatVectors;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
+public class GraphIndexBuilderTest extends LuceneTestCase {
+
+    private Path testDirectory;
+
+    @Before
+    public void setup() throws IOException {
+        testDirectory = Files.createTempDirectory(getClass().getSimpleName());
+    }
+
+    @After
+    public void tearDown() {
+        TestUtil.deleteQuietly(testDirectory);
+    }
+
+    @Test
+    public void testSaveAndLoad() throws IOException {
+        int dimension = randomIntBetween(2, 32);
+        var ravv = MockVectorValues.fromValues(createRandomFloatVectors(randomIntBetween(10, 100), dimension, getRandom()));
+        Supplier<GraphIndexBuilder<float[]>> newBuilder = () ->
+            new GraphIndexBuilder<>(ravv, VectorEncoding.FLOAT32, VectorSimilarityFunction.COSINE, 2, 10, 1.0f, 1.0f);
+
+        var indexDataPath = testDirectory.resolve("index_builder.data");
+        var builder = newBuilder.get();
+
+        try (var graph = TestUtil.buildSequentially(builder, ravv);
+             var out = TestUtil.openFileForWriting(indexDataPath)) {
+            graph.save(out);
+            out.flush();
+        }
+
+        builder = newBuilder.get();
+        try(var reader = new SimpleMappedReader(indexDataPath)) {
+            builder.load(reader);
+        }
+
+        assertEquals(ravv.size(), builder.graph.size());
+        for (int i = 0; i < ravv.size(); i++) {
+            assertTrue(builder.graph.containsNode(i));
+        }
+    }
+}