update graph to RAVV mapping for buildScoreProvider

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

Author: sam-herman

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

@@ -826,18 +826,12 @@ public static OnHeapGraphIndex buildAndMergeNewNodes(OnDiskGraphIndex onDiskGrap
                 PhysicalCoreExecutor.pool(),
                 ForkJoinPool.commonPool())) {
 
-            // Add each new vector incrementally
-            final List<ForkJoinTask<?>> forkJoinTask = new ArrayList<>(newVectors.size());
-            for (int i = startingNodeOffset; i < newVectors.size(); i++) {
-                final int nodeId = i;
-                final VectorFloat<?> vector = newVectors.getVector(graphToRavvOrdMap[nodeId]);
-
-                // The GraphIndexBuilder can add nodes to an existing index
-                forkJoinTask.add(PhysicalCoreExecutor.pool().submit(() -> builder.addGraphNode(nodeId, vector)));
-            }
-            for (ForkJoinTask<?> task : forkJoinTask) {
-                task.join();
-            }
+            var vv = newVectors.threadLocalSupplier();
+
+            // parallel graph construction from the merge documents Ids
+            PhysicalCoreExecutor.pool().submit(() -> IntStream.range(startingNodeOffset, newVectors.size()).parallel().forEach(ord -> {
+                builder.addGraphNode(ord, vv.get().getVector(graphToRavvOrdMap[ord]));
+            })).join();
 
             builder.cleanup();
             return builder.getGraph();

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

@@ -133,6 +133,58 @@ public SearchScoreProvider diversityProviderFor(int node1) {
         };
     }
 
+    /**
+     * Returns a BSP that performs exact score comparisons using the given RandomAccessVectorValues and VectorSimilarityFunction.
+     */
+    static BuildScoreProvider randomAccessScoreProvider(RandomAccessVectorValues ravv, int[] graphToRavvOrdMap, VectorSimilarityFunction similarityFunction) {
+        // We need two sources of vectors in order to perform diversity check comparisons without
+        // colliding.  ThreadLocalSupplier makes this a no-op if the RAVV is actually un-shared.
+        var vectors = ravv.threadLocalSupplier();
+        var vectorsCopy = ravv.threadLocalSupplier();
+
+        return new BuildScoreProvider() {
+            @Override
+            public boolean isExact() {
+                return true;
+            }
+
+            @Override
+            public VectorFloat<?> approximateCentroid() {
+                var vv = vectors.get();
+                var centroid = vts.createFloatVector(vv.dimension());
+                for (int i = 0; i < vv.size(); i++) {
+                    var v = vv.getVector(i);
+                    if (v != null) { // MapRandomAccessVectorValues is not necessarily dense
+                        VectorUtil.addInPlace(centroid, v);
+                    }
+                }
+                VectorUtil.scale(centroid, 1.0f / vv.size());
+                return centroid;
+            }
+
+            @Override
+            public SearchScoreProvider searchProviderFor(VectorFloat<?> vector) {
+                var vc = vectorsCopy.get();
+                return DefaultSearchScoreProvider.exact(vector, graphToRavvOrdMap, similarityFunction, vc);
+            }
+
+            @Override
+            public SearchScoreProvider searchProviderFor(int node1) {
+                RandomAccessVectorValues randomAccessVectorValues = vectors.get();
+                var v = randomAccessVectorValues.getVector(graphToRavvOrdMap[node1]);
+                return searchProviderFor(v);
+            }
+
+            @Override
+            public SearchScoreProvider diversityProviderFor(int node1) {
+                RandomAccessVectorValues randomAccessVectorValues = vectors.get();
+                var v = randomAccessVectorValues.getVector(graphToRavvOrdMap[node1]);
+                var vc = vectorsCopy.get();
+                return DefaultSearchScoreProvider.exact(v, graphToRavvOrdMap, similarityFunction, vc);
+            }
+        };
+    }
+
     /**
      * Returns a BSP that performs approximate score comparisons using the given PQVectors,
      * with reranking performed using RandomAccessVectorValues (which is intended to be

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

@@ -78,4 +78,20 @@ public float similarityTo(int node2) {
         };
         return new DefaultSearchScoreProvider(sf);
     }
+
+    /**
+     * A SearchScoreProvider for a single-pass search based on exact similarity.
+     * Generally only suitable when your RandomAccessVectorValues is entirely in-memory,
+     * e.g. during construction.
+     */
+    public static DefaultSearchScoreProvider exact(VectorFloat<?> v, int[] graphToRavvOrdMap ,VectorSimilarityFunction vsf, RandomAccessVectorValues ravv) {
+        // don't use ESF.reranker, we need thread safety here
+        var sf = new ScoreFunction.ExactScoreFunction() {
+            @Override
+            public float similarityTo(int node2) {
+                return vsf.compare(v, ravv.getVector(graphToRavvOrdMap[node2]));
+            }
+        };
+        return new DefaultSearchScoreProvider(sf);
+    }
 }

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

@@ -24,13 +24,15 @@
 import io.github.jbellis.jvector.graph.disk.NeighborsScoreCache;
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider;
+import io.github.jbellis.jvector.graph.similarity.SearchScoreProvider;
 import io.github.jbellis.jvector.util.Bits;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorizationProvider;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
 import org.apache.logging.log4j.Logger;
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -50,15 +52,16 @@
 public class OnHeapGraphIndexTest extends RandomizedTest  {
     private final static Logger log = org.apache.logging.log4j.LogManager.getLogger(OnHeapGraphIndexTest.class);
     private static final VectorTypeSupport VECTOR_TYPE_SUPPORT = VectorizationProvider.getInstance().getVectorTypeSupport();
-    private static final int numBaseVectors = 100;
-    private static final int numNewVectors = 100;
-    private static final int numAllVectors = numBaseVectors + numNewVectors;
-    private static final int dimension = 16;
+    private static final int NUM_BASE_VECTORS = 100;
+    private static final int NUM_NEW_VECTORS = 100;
+    private static final int NUM_ALL_VECTORS = NUM_BASE_VECTORS + NUM_NEW_VECTORS;
+    private static final int DIMENSION = 16;
     private static final int M = 8;
-    private static final int beamWidth = 100;
-    private static final float alpha = 1.2f;
-    private static final float neighborOverflow = 1.2f;
-    private static final boolean addHierarchy = false;
+    private static final int BEAM_WIDTH = 100;
+    private static final float ALPHA = 1.2f;
+    private static final float NEIGHBOR_OVERFLOW = 1.2f;
+    private static final boolean ADD_HIERARCHY = false;
+    private static final int TOP_K = 10;
 
     private Path testDirectory;
 
@@ -68,6 +71,9 @@ public class OnHeapGraphIndexTest extends RandomizedTest  {
     private RandomAccessVectorValues baseVectorsRavv;
     private RandomAccessVectorValues newVectorsRavv;
     private RandomAccessVectorValues allVectorsRavv;
+    private VectorFloat<?> queryVector;
+    private int[] groundTruthBaseVectors;
+    private int[] groundTruthAllVectors;
     private BuildScoreProvider baseBuildScoreProvider;
     private BuildScoreProvider newBuildScoreProvider;
     private BuildScoreProvider allBuildScoreProvider;
@@ -78,42 +84,47 @@ public class OnHeapGraphIndexTest extends RandomizedTest  {
     @Before
     public void setup() throws IOException {
         testDirectory = Files.createTempDirectory(this.getClass().getSimpleName());
-        baseVectors = new ArrayList<>(numBaseVectors);
-        newVectors = new ArrayList<>(numNewVectors);
-        allVectors = new ArrayList<>(numAllVectors);
-        for (int i = 0; i < numBaseVectors; i++) {
-            VectorFloat<?> vector = createRandomVector(dimension);
+        baseVectors = new ArrayList<>(NUM_BASE_VECTORS);
+        newVectors = new ArrayList<>(NUM_NEW_VECTORS);
+        allVectors = new ArrayList<>(NUM_ALL_VECTORS);
+        for (int i = 0; i < NUM_BASE_VECTORS; i++) {
+            VectorFloat<?> vector = createRandomVector(DIMENSION);
             baseVectors.add(vector);
             allVectors.add(vector);
         }
-        for (int i = 0; i < numNewVectors; i++) {
-            VectorFloat<?> vector = createRandomVector(dimension);
+        for (int i = 0; i < NUM_NEW_VECTORS; i++) {
+            VectorFloat<?> vector = createRandomVector(DIMENSION);
             newVectors.add(vector);
             allVectors.add(vector);
         }
 
         // wrap the raw vectors in a RandomAccessVectorValues
-        baseVectorsRavv = new ListRandomAccessVectorValues(baseVectors, dimension);
-        newVectorsRavv = new ListRandomAccessVectorValues(newVectors, dimension);
-        allVectorsRavv = new ListRandomAccessVectorValues(allVectors, dimension);
+        baseVectorsRavv = new ListRandomAccessVectorValues(baseVectors, DIMENSION);
+        newVectorsRavv = new ListRandomAccessVectorValues(newVectors, DIMENSION);
+        allVectorsRavv = new ListRandomAccessVectorValues(allVectors, DIMENSION);
 
+        queryVector = createRandomVector(DIMENSION);
+        groundTruthBaseVectors = getGroundTruth(baseVectorsRavv, queryVector, TOP_K, VectorSimilarityFunction.EUCLIDEAN);
+        groundTruthAllVectors = getGroundTruth(allVectorsRavv, queryVector, TOP_K, VectorSimilarityFunction.EUCLIDEAN);
+
+        // score provider using the raw, in-memory vectors
         baseBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(baseVectorsRavv, VectorSimilarityFunction.EUCLIDEAN);
         newBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(newVectorsRavv, VectorSimilarityFunction.EUCLIDEAN);
         allBuildScoreProvider = BuildScoreProvider.randomAccessScoreProvider(allVectorsRavv, VectorSimilarityFunction.EUCLIDEAN);
         var baseGraphIndexBuilder = new GraphIndexBuilder(baseBuildScoreProvider,
                 baseVectorsRavv.dimension(),
                 M, // graph degree
-                beamWidth, // construction search depth
-                neighborOverflow, // allow degree overflow during construction by this factor
-                alpha, // relax neighbor diversity requirement by this factor
-                addHierarchy); // add the hierarchy
+                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 allGraphIndexBuilder = new GraphIndexBuilder(allBuildScoreProvider,
                 allVectorsRavv.dimension(),
                 M, // graph degree
-                beamWidth, // construction search depth
-                neighborOverflow, // allow degree overflow during construction by this factor
-                alpha, // relax neighbor diversity requirement by this factor
-                addHierarchy); // add the hierarchy
+                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
 
         baseGraphIndex = baseGraphIndexBuilder.build(baseVectorsRavv);
         allGraphIndex = allGraphIndexBuilder.build(allVectorsRavv);
@@ -156,7 +167,7 @@ public void testReconstructionOfOnHeapGraphIndex() throws IOException {
                 validateVectors(onDiskView, baseVectorsRavv);
             }
 
-            OnHeapGraphIndex reconstructedOnHeapGraphIndex = OnHeapGraphIndex.convertToHeap(onDiskGraph, neighborsScoreCacheRead, baseBuildScoreProvider, neighborOverflow, alpha);
+            OnHeapGraphIndex reconstructedOnHeapGraphIndex = OnHeapGraphIndex.convertToHeap(onDiskGraph, neighborsScoreCacheRead, baseBuildScoreProvider, NEIGHBOR_OVERFLOW, ALPHA);
             TestUtil.assertGraphEquals(baseGraphIndex, reconstructedOnHeapGraphIndex);
             TestUtil.assertGraphEquals(onDiskGraph, reconstructedOnHeapGraphIndex);
 
@@ -178,12 +189,18 @@ public void testIncrementalInsertionFromOnDiskIndex() throws IOException {
             TestUtil.assertGraphEquals(baseGraphIndex, onDiskGraph);
             // We will create a trivial 1:1 mapping between the new graph and the ravv
             final int[] graphToRavvOrdMap = IntStream.range(0, allVectorsRavv.size()).toArray();
-            OnHeapGraphIndex reconstructedAllNodeOnHeapGraphIndex = GraphIndexBuilder.buildAndMergeNewNodes(onDiskGraph, neighborsScoreCache, allVectorsRavv, allBuildScoreProvider, numBaseVectors, graphToRavvOrdMap, beamWidth, neighborOverflow, alpha, addHierarchy);
+            OnHeapGraphIndex reconstructedAllNodeOnHeapGraphIndex = GraphIndexBuilder.buildAndMergeNewNodes(onDiskGraph, neighborsScoreCache, allVectorsRavv, allBuildScoreProvider, NUM_BASE_VECTORS, graphToRavvOrdMap, BEAM_WIDTH, NEIGHBOR_OVERFLOW, ALPHA, ADD_HIERARCHY);
+
+            // Verify that the recall is similar
+            float recallFromReconstructedAllNodeOnHeapGraphIndex = calculateRecall(reconstructedAllNodeOnHeapGraphIndex, allBuildScoreProvider, queryVector, groundTruthAllVectors, TOP_K);
+            float recallFromAllGraphIndex = calculateRecall(allGraphIndex, allBuildScoreProvider, queryVector, groundTruthAllVectors, TOP_K);
+            Assert.assertEquals(recallFromReconstructedAllNodeOnHeapGraphIndex, recallFromAllGraphIndex, 0.01f);
 
+            // Verify that the result sets overlap
             try (GraphSearcher reconstructedAllGraphSearcher = new GraphSearcher(reconstructedAllNodeOnHeapGraphIndex);
                  GraphSearcher allGraphSearcher = new GraphSearcher(allGraphIndex)) {
-                final int topK = 10;
-                VectorFloat<?> queryVector = createRandomVector(dimension);
+                final int topK = TOP_K;
+                VectorFloat<?> queryVector = createRandomVector(DIMENSION);
                 var resultFromReconstructed = reconstructedAllGraphSearcher.search(allBuildScoreProvider.searchProviderFor(queryVector), topK, Bits.ALL);
                 var resultFromAll = allGraphSearcher.search(allBuildScoreProvider.searchProviderFor(queryVector), topK, Bits.ALL);
                 log.info("Reconstructed result: {}, all result: {}", resultFromReconstructed, resultFromAll);
@@ -210,4 +227,54 @@ private VectorFloat<?> createRandomVector(int dimension) {
         }
         return vector;
     }
+
+    /**
+     * Get the ground truth for a query vector
+     * @param ravv the vectors to search
+     * @param queryVector the query vector
+     * @param topK the number of results to return
+     * @param similarityFunction the similarity function to use
+
+     * @return the ground truth
+     */
+    private static int[] getGroundTruth(RandomAccessVectorValues ravv, VectorFloat<?> queryVector, int topK, VectorSimilarityFunction similarityFunction) {
+        var exactResults = new ArrayList<SearchResult.NodeScore>();
+        for (int i = 0; i < ravv.size(); i++) {
+            float similarityScore = similarityFunction.compare(queryVector, ravv.getVector(i));
+            exactResults.add(new SearchResult.NodeScore(i, similarityScore));
+        }
+        exactResults.sort((a, b) -> Float.compare(b.score, a.score));
+        return exactResults.stream().limit(topK).mapToInt(nodeScore -> nodeScore.node).toArray();
+    }
+
+    private static float calculateRecall(OnHeapGraphIndex graphIndex, BuildScoreProvider buildScoreProvider, VectorFloat<?> queryVector, int[] groundTruth, int k) 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());
+            return calculateRecall(predicted, groundTruth, k);
+        }
+    }
+    /**
+     * Calculate the recall for a set of predicted results
+     * @param predicted the predicted results
+     * @param groundTruth the ground truth
+     * @param k the number of results to consider
+     * @return the recall
+     */
+    private static float calculateRecall(Set<Integer> predicted, int[] groundTruth, int k) {
+        int hits = 0;
+        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;
+                }
+            }
+        }
+
+        return ((float) hits) / (float) actualK;
+    }
 }