harden tests for heap graph reconstruction

fix minor bug in construction

Signed-off-by: Samuel Herman <sherman8915@gmail.com>

Author: sam-herman

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/RecallWithRandomVectorsBenchmark.java

@@ -247,11 +247,8 @@ private double calculateRecall(Set<Integer> predicted, int[] groundTruth, int k)
         int actualK = Math.min(k, Math.min(predicted.size(), groundTruth.length));
 
         for (int i = 0; i < actualK; i++) {
-            for (int j = 0; j < actualK; j++) {
-                if (predicted.contains(groundTruth[j])) {
-                    hits++;
-                    break;
-                }
+            if (predicted.contains(groundTruth[i])) {
+                hits++;
             }
         }
 

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java

@@ -439,12 +439,17 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi
     }
 
     public ImmutableGraphIndex build(RandomAccessVectorValues ravv) {
+        return build(ravv, null);
+    }
+
+    public ImmutableGraphIndex build(RandomAccessVectorValues ravv, int[] graphToRavvOrdMap) {
         var vv = ravv.threadLocalSupplier();
         int size = ravv.size();
 
         simdExecutor.submit(() -> {
             IntStream.range(0, size).parallel().forEach(node -> {
-                addGraphNode(node, vv.get().getVector(node));
+                int ravvOrdinal = (graphToRavvOrdMap != null) ? graphToRavvOrdMap[node] : node;
+                addGraphNode(node, vv.get().getVector(ravvOrdinal));
             });
         }).join();
 

jvector-tests/src/test/java/io/github/jbellis/jvector/graph/OnHeapGraphIndexTest.java

@@ -45,6 +45,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
+import static org.apache.commons.lang3.ArrayUtils.shuffle;
 import static org.junit.Assert.assertEquals;
 
 @ThreadLeakScope(ThreadLeakScope.Scope.NONE)
@@ -61,6 +62,7 @@ public class OnHeapGraphIndexTest extends RandomizedTest  {
     private static final float NEIGHBOR_OVERFLOW = 1.2f;
     private static final boolean ADD_HIERARCHY = false;
     private static final int TOP_K = 10;
+    private static VectorSimilarityFunction SIMILARITY_FUNCTION = VectorSimilarityFunction.EUCLIDEAN;
 
     private Path testDirectory;
 
@@ -71,6 +73,7 @@ public class OnHeapGraphIndexTest extends RandomizedTest  {
     private RandomAccessVectorValues newVectorsRavv;
     private RandomAccessVectorValues allVectorsRavv;
     private VectorFloat<?> queryVector;
+    private int[] groundTruthBaseVectors;
     private int[] groundTruthAllVectors;
     private BuildScoreProvider baseBuildScoreProvider;
     private BuildScoreProvider allBuildScoreProvider;
@@ -100,11 +103,12 @@ public void setup() throws IOException {
         allVectorsRavv = new ListRandomAccessVectorValues(allVectors, DIMENSION);
 
         queryVector = createRandomVector(DIMENSION);
-        groundTruthAllVectors = getGroundTruth(allVectorsRavv, queryVector, TOP_K, VectorSimilarityFunction.EUCLIDEAN);
+        groundTruthBaseVectors = getGroundTruth(baseVectorsRavv, queryVector, TOP_K, SIMILARITY_FUNCTION);
+        groundTruthAllVectors = getGroundTruth(allVectorsRavv, queryVector, TOP_K, SIMILARITY_FUNCTION);
 
         // score provider using the raw, in-memory vectors
-        baseBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(baseVectorsRavv, VectorSimilarityFunction.EUCLIDEAN);
-        allBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(allVectorsRavv, VectorSimilarityFunction.EUCLIDEAN);
+        baseBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(baseVectorsRavv, SIMILARITY_FUNCTION);
+        allBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(allVectorsRavv, SIMILARITY_FUNCTION);
         var baseGraphIndexBuilder = new GraphIndexBuilder(baseBuildScoreProvider,
                 baseVectorsRavv.dimension(),
                 M, // graph degree
@@ -129,14 +133,38 @@ public void tearDown() {
         TestUtil.deleteQuietly(testDirectory);
     }
 
+    /**
+     * Test that we can build a graph with a non-identity mapping from graph node id to ravv ordinal
+     * and that the recall is the same as the identity mapping (meaning the graphs are equivalent)
+     * @throws IOException exception
+     */
+    @Test
+    public void testGraphConstructionWithNonIdentityOrdinalMapping() throws IOException {
+        // create reversed mapping from graph node id to ravv ordinal
+        int[] graphToRavvOrdMap = IntStream.range(0, baseVectorsRavv.size()).map(i -> baseVectorsRavv.size() - 1 - i).toArray();
+        var bsp = BuildScoreProvider.randomAccessScoreProvider(baseVectorsRavv, graphToRavvOrdMap, SIMILARITY_FUNCTION);
+        try (var baseGraphIndexBuilder = new GraphIndexBuilder(bsp,
+                baseVectorsRavv.dimension(),
+                M, // graph degree
+                BEAM_WIDTH, // construction search depth
+                NEIGHBOR_OVERFLOW, // allow degree overflow during construction by this factor
+                ALPHA, // relax neighbor diversity requirement by this factor
+                ADD_HIERARCHY); // add the hierarchy) {
+             var baseGraphIndexFromShuffledVectors = baseGraphIndexBuilder.build(baseVectorsRavv, graphToRavvOrdMap)) {
+            float recallFromBaseGraphIndexFromShuffledVectors = calculateRecall(baseGraphIndexFromShuffledVectors, bsp, queryVector, groundTruthBaseVectors, TOP_K, graphToRavvOrdMap);
+            float recallFromBaseGraphIndex = calculateRecall(baseGraphIndex, baseBuildScoreProvider, queryVector, groundTruthBaseVectors, TOP_K);
+            Assert.assertEquals(recallFromBaseGraphIndex, recallFromBaseGraphIndexFromShuffledVectors, 0.01f);
+        }
+    }
 
     /**
      * Create an {@link OnHeapGraphIndex} persist it as a {@link OnDiskGraphIndex} and reconstruct back to a mutable {@link OnHeapGraphIndex}
+     * Using identity mapping from graph node id to ravv ordinal
      * Make sure that both graphs are equivalent
      * @throws IOException
      */
     @Test
-    public void testReconstructionOfOnHeapGraphIndex() throws IOException {
+    public void testReconstructionOfOnHeapGraphIndex_withIdentityOrdinalMapping() throws IOException {
         var graphOutputPath = testDirectory.resolve("testReconstructionOfOnHeapGraphIndex_" + baseGraphIndex.getClass().getSimpleName());
         var heapGraphOutputPath = testDirectory.resolve("testReconstructionOfOnHeapGraphIndex_" + baseGraphIndex.getClass().getSimpleName() + "_onHeap");
 
@@ -166,6 +194,44 @@ public void testReconstructionOfOnHeapGraphIndex() throws IOException {
         }
     }
 
+    /**
+     * Create an {@link OnHeapGraphIndex} persist it as a {@link OnDiskGraphIndex} and reconstruct back to a mutable {@link OnHeapGraphIndex}
+     * Using random mapping from graph node id to ravv ordinal
+     * Make sure that both graphs are equivalent
+     * @throws IOException
+     */
+    @Test
+    public void testReconstructionOfOnHeapGraphIndex_withNonIdentityOrdinalMapping() throws IOException {
+        var graphOutputPath = testDirectory.resolve("testReconstructionOfOnHeapGraphIndex_" + baseGraphIndex.getClass().getSimpleName());
+        var heapGraphOutputPath = testDirectory.resolve("testReconstructionOfOnHeapGraphIndex_" + baseGraphIndex.getClass().getSimpleName() + "_onHeap");
+
+        // create reversed mapping from graph node id to ravv ordinal
+        int[] graphToRavvOrdMap = IntStream.range(0, baseVectorsRavv.size()).map(i -> baseVectorsRavv.size() - 1 - i).toArray();
+        var bsp = BuildScoreProvider.randomAccessScoreProvider(baseVectorsRavv, graphToRavvOrdMap, SIMILARITY_FUNCTION);
+        try (var baseGraphIndexBuilder = new GraphIndexBuilder(bsp,
+                baseVectorsRavv.dimension(),
+                M, // graph degree
+                BEAM_WIDTH, // construction search depth
+                NEIGHBOR_OVERFLOW, // allow degree overflow during construction by this factor
+                ALPHA, // relax neighbor diversity requirement by this factor
+                ADD_HIERARCHY); // add the hierarchy) {
+             var baseGraphIndex = baseGraphIndexBuilder.build(baseVectorsRavv, graphToRavvOrdMap)) {
+            log.info("Writing graph to {}", graphOutputPath);
+            TestUtil.writeGraph(baseGraphIndex, baseVectorsRavv, graphOutputPath);
+
+            log.info("Writing on-heap graph to {}", heapGraphOutputPath);
+            try (SimpleWriter writer = new SimpleWriter(heapGraphOutputPath.toAbsolutePath())) {
+                ((OnHeapGraphIndex) baseGraphIndex).save(writer);
+            }
+
+            log.info("Reading on-heap graph from {}", heapGraphOutputPath);
+            try (var readerSupplier = new SimpleMappedReader.Supplier(heapGraphOutputPath.toAbsolutePath())) {
+                MutableGraphIndex reconstructedOnHeapGraphIndex = OnHeapGraphIndex.load(readerSupplier.get(), baseVectorsRavv.dimension(), NEIGHBOR_OVERFLOW, new VamanaDiversityProvider(bsp, ALPHA));
+                TestUtil.assertGraphEquals(baseGraphIndex, reconstructedOnHeapGraphIndex);
+            }
+        }
+    }
+
     /**
      * Create {@link OnDiskGraphIndex} then append to it via {@link GraphIndexBuilder#buildAndMergeNewNodes}
      * Verify that the resulting OnHeapGraphIndex is equivalent to the graph that would have been alternatively generated by bulk index into a new {@link OnDiskGraphIndex}
@@ -230,10 +296,27 @@ private static int[] getGroundTruth(RandomAccessVectorValues ravv, VectorFloat<?
     }
 
     private static float calculateRecall(ImmutableGraphIndex graphIndex, BuildScoreProvider buildScoreProvider, VectorFloat<?> queryVector, int[] groundTruth, int k) throws IOException {
+        return calculateRecall(graphIndex, buildScoreProvider, queryVector, groundTruth, k, null);
+    }
+
+    private static float calculateRecall(ImmutableGraphIndex graphIndex, BuildScoreProvider buildScoreProvider, VectorFloat<?> queryVector, int[] groundTruth, int k, int[] graphToRavvOrdMap) throws IOException {
         try (GraphSearcher graphSearcher = new GraphSearcher(graphIndex)){
             SearchScoreProvider ssp = buildScoreProvider.searchProviderFor(queryVector);
             var searchResults = graphSearcher.search(ssp, k, Bits.ALL);
-            var predicted = Arrays.stream(searchResults.getNodes()).mapToInt(nodeScore -> nodeScore.node).boxed().collect(Collectors.toSet());
+            Set<Integer> predicted;
+            if (graphToRavvOrdMap != null) {
+                // Convert graph node IDs to RAVV ordinals for comparison with ground truth
+                predicted = Arrays.stream(searchResults.getNodes())
+                        .mapToInt(nodeScore -> graphToRavvOrdMap[nodeScore.node])
+                        .boxed()
+                        .collect(Collectors.toSet());
+            } else {
+                // Identity mapping: graph node IDs == RAVV ordinals
+                predicted = Arrays.stream(searchResults.getNodes())
+                        .mapToInt(nodeScore -> nodeScore.node)
+                        .boxed()
+                        .collect(Collectors.toSet());
+            }
             return calculateRecall(predicted, groundTruth, k);
         }
     }
@@ -249,11 +332,8 @@ private static float calculateRecall(Set<Integer> predicted, int[] groundTruth,
         int actualK = Math.min(k, Math.min(predicted.size(), groundTruth.length));
 
         for (int i = 0; i < actualK; i++) {
-            for (int j = 0; j < actualK; j++) {
-                if (predicted.contains(groundTruth[j])) {
-                    hits++;
-                    break;
-                }
+            if (predicted.contains(groundTruth[i])) {
+                hits++;
             }
         }