Reducing the number of allocations in GraphSearcher (#501)

marianotepper · web-flow · commit 66399923474c · 2025-08-19T14:17:32.000-07:00
* Reduces the number of allocations in GraphSearcher

* Minor improvements to code readability
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphSearcher.java
@@ -32,7 +32,6 @@
 import io.github.jbellis.jvector.util.Bits;
 import io.github.jbellis.jvector.util.BoundedLongHeap;
 import io.github.jbellis.jvector.util.GrowableLongHeap;
-import io.github.jbellis.jvector.util.SparseBits;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import org.agrona.collections.Int2ObjectHashMap;
@@ -47,8 +46,6 @@
  * search algorithm, see {@link GraphIndex}.
  */
 public class GraphSearcher implements Closeable {
-    private boolean pruneSearch;
-
     private GraphIndex.View view;
 
     // Scratch data structures that are used in each {@link #searchInternal} call. These can be expensive
@@ -64,6 +61,9 @@ public class GraphSearcher implements Closeable {
     private SearchScoreProvider scoreProvider;
     private CachingReranker cachingReranker;
 
+    private boolean pruneSearch;
+    private final ScoreTracker.ScoreTrackerFactory scoreTrackerFactory;
+
     private int visitedCount;
     private int expandedCount;
     private int expandedCountBaseLayer;
@@ -75,14 +75,31 @@ public GraphSearcher(GraphIndex graph) {
         this(graph.getView());
     }
 
-    private GraphSearcher(GraphIndex.View view) {
+    /**
+     * Creates a new graph searcher from the given GraphIndex.View
+     */
+    protected GraphSearcher(GraphIndex.View view) {
         this.view = view;
         this.candidates = new NodeQueue(new GrowableLongHeap(100), NodeQueue.Order.MAX_HEAP);
         this.evictedResults = new NodesUnsorted(100);
         this.approximateResults = new NodeQueue(new BoundedLongHeap(100), NodeQueue.Order.MIN_HEAP);
         this.rerankedResults = new NodeQueue(new BoundedLongHeap(100), NodeQueue.Order.MIN_HEAP);
         this.visited = new IntHashSet();
+
         this.pruneSearch = true;
+        this.scoreTrackerFactory = new ScoreTracker.ScoreTrackerFactory();
+    }
+
+    protected int getVisitedCount() {
+        return visitedCount;
+    }
+
+    protected int getExpandedCount() {
+        return expandedCount;
+    }
+
+    protected int getExpandedCountBaseLayer() {
+        return expandedCountBaseLayer;
     }
 
     private void initializeScoreProvider(SearchScoreProvider scoreProvider) {
@@ -208,6 +225,35 @@ public SearchResult search(SearchScoreProvider scoreProvider,
             return new SearchResult(new SearchResult.NodeScore[0], 0, 0, 0, 0, Float.POSITIVE_INFINITY);
         }
 
+        internalSearch(scoreProvider, entry, topK, rerankK, threshold, acceptOrds);
+        return reranking(topK, rerankK, rerankFloor);
+    }
+
+    /**
+     * Performs a search, leaving the results in the internal member variable approximateResults.
+     * It does not perform reranking.
+     *
+     * @param scoreProvider   provides functions to return the similarity of a given node to the query vector
+     * @param entry           the entry point to the graph. Assumed to be not null.
+     * @param topK            the number of results to look for. With threshold=0, the search will continue until at least
+     *                        `topK` results have been found, or until the entire graph has been searched.
+     * @param rerankK         the number of (approximately-scored) results to rerank before returning the best `topK`.
+     * @param threshold       the minimum similarity (0..1) to accept; 0 will accept everything. May be used
+     *                        with a large topK to find (approximately) all nodes above the given threshold.
+     *                        If threshold > 0 then the search will stop when it is probabilistically unlikely
+     *                        to find more nodes above the threshold, even if `topK` results have not yet been found.
+     * @param acceptOrds      a Bits instance indicating which nodes are acceptable results.
+     *                        If {@link Bits#ALL}, all nodes are acceptable.
+     *                        It is caller's responsibility to ensure that there are enough acceptable nodes
+     *                        that we don't search the entire graph trying to satisfy topK.
+     */
+    protected void internalSearch(SearchScoreProvider scoreProvider,
+                                  NodeAtLevel entry,
+                                  int topK,
+                                  int rerankK,
+                                  float threshold,
+                                  Bits acceptOrds)
+    {
         initializeInternal(scoreProvider, entry, acceptOrds);
 
         // Move downward from entry.level to 1
@@ -219,7 +265,7 @@ public SearchResult search(SearchScoreProvider scoreProvider,
         }
 
         // Now do the main search at layer 0
-        return resume(topK, rerankK, threshold, rerankFloor);
+        searchLayer0(topK, rerankK, threshold);;
     }
 
     /**
@@ -276,6 +322,7 @@ void initializeInternal(SearchScoreProvider scoreProvider, NodeAtLevel entry, Bi
         this.acceptOrds = Bits.intersectionOf(rawAcceptOrds, view.liveNodes());
 
         // reset the scratch data structures
+        approximateResults.clear();
         evictedResults.clear();
         candidates.clear();
         visited.clear();
@@ -290,6 +337,19 @@ void initializeInternal(SearchScoreProvider scoreProvider, NodeAtLevel entry, Bi
         expandedCountBaseLayer = 0;
     }
 
+    private boolean stopSearch(NodeQueue localCandidates, ScoreTracker scoreTracker, int rerankK, float threshold) {
+        float topCandidateScore = localCandidates.topScore();
+        // we're done when we have K results and the best candidate is worse than the worst result so far
+        if (approximateResults.size() >= rerankK && topCandidateScore < approximateResults.topScore()) {
+            return true;
+        }
+        // when querying by threshold, also stop when we are probabilistically unlikely to find more qualifying results
+        if (threshold > 0 && scoreTracker.shouldStop()) {
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Performs a single-layer ANN search, expanding from the given candidates queue.
      *
@@ -335,24 +395,17 @@ void searchOneLayer(SearchScoreProvider scoreProvider,
             approximateResults.setMaxSize(rerankK);
 
             // track scores to predict when we are done with threshold queries
-            var scoreTracker = threshold > 0
-                    ? new ScoreTracker.TwoPhaseTracker(threshold)
-                    : pruneSearch ? new ScoreTracker.RelaxedMonotonicityTracker(rerankK) : new ScoreTracker.NoOpTracker();
+            var scoreTracker = scoreTrackerFactory.getScoreTracker(pruneSearch, rerankK, threshold);
             VectorFloat<?> similarities = null;
 
             // the main search loop
             while (candidates.size() > 0) {
-                // we're done when we have K results and the best candidate is worse than the worst result so far
-                float topCandidateScore = candidates.topScore();
-                if (approximateResults.size() >= rerankK && topCandidateScore < approximateResults.topScore()) {
-                    break;
-                }
-                // when querying by threshold, also stop when we are probabilistically unlikely to find more qualifying results
-                if (threshold > 0 && scoreTracker.shouldStop()) {
+                if (stopSearch(candidates, scoreTracker, rerankK, threshold)) {
                     break;
                 }
 
                 // process the top candidate
+                float topCandidateScore = candidates.topScore();
                 int topCandidateNode = candidates.pop();
                 if (acceptOrdsThisLayer.get(topCandidateNode) && topCandidateScore >= threshold) {
                     addTopCandidate(topCandidateNode, topCandidateScore, rerankK);
@@ -397,7 +450,7 @@ void searchOneLayer(SearchScoreProvider scoreProvider,
         }
     }
 
-    SearchResult resume(int topK, int rerankK, float threshold, float rerankFloor) {
+    private void searchLayer0(int topK, int rerankK, float threshold) {
         // rR is persistent to save on allocations
         rerankedResults.clear();
         rerankedResults.setMaxSize(topK);
@@ -407,7 +460,9 @@ SearchResult resume(int topK, int rerankK, float threshold, float rerankFloor) {
         evictedResults.clear();
 
         searchOneLayer(scoreProvider, rerankK, threshold, 0, acceptOrds);
+    }
 
+    SearchResult reranking(int topK, int rerankK, float rerankFloor) {
         // rerank results
         assert approximateResults.size() <= rerankK;
         NodeQueue popFromQueue;
@@ -445,6 +500,11 @@ SearchResult resume(int topK, int rerankK, float threshold, float rerankFloor) {
         return new SearchResult(nodes, visitedCount, expandedCount, expandedCountBaseLayer, reranked, worstApproximateInTopK);
     }
 
+    SearchResult resume(int topK, int rerankK, float threshold, float rerankFloor) {
+        searchLayer0(topK, rerankK, threshold);
+        return reranking(topK, rerankK, rerankFloor);
+    }
+
     @SuppressWarnings("StatementWithEmptyBody")
     private void addTopCandidate(int topCandidateNode, float topCandidateScore, int rerankK) {
         // add the new node to the results queue, and any evicted node to evictedResults in case we resume later
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ScoreTracker.java b/jvector-base/src/main/java/io/github/jbellis/jvector/graph/ScoreTracker.java
@@ -16,13 +16,54 @@
 
 package io.github.jbellis.jvector.graph;
 
+import io.github.jbellis.jvector.util.ArrayUtil;
 import io.github.jbellis.jvector.util.BoundedLongHeap;
 import org.apache.commons.math3.stat.StatUtils;
 
 import static io.github.jbellis.jvector.util.NumericUtils.floatToSortableInt;
 import static io.github.jbellis.jvector.util.NumericUtils.sortableIntToFloat;
 
-interface ScoreTracker {
+public interface ScoreTracker {
+    class ScoreTrackerFactory {
+        private TwoPhaseTracker twoPhaseTracker;
+        private RelaxedMonotonicityTracker relaxedMonotonicityTracker;
+        private NoOpTracker noOpTracker;
+
+        ScoreTrackerFactory() {
+            twoPhaseTracker = null;
+            relaxedMonotonicityTracker = null;
+            noOpTracker = null;
+        }
+
+        public ScoreTracker getScoreTracker(boolean pruneSearch, int rerankK, float threshold) {
+            // track scores to predict when we are done with threshold queries
+            final ScoreTracker scoreTracker;
+
+            if (threshold > 0) {
+                if (twoPhaseTracker == null) {
+                    twoPhaseTracker = new ScoreTracker.TwoPhaseTracker();
+                } else {
+                    twoPhaseTracker.reset(threshold);
+                }
+                scoreTracker = twoPhaseTracker;
+            } else {
+                if (pruneSearch) {
+                    if (relaxedMonotonicityTracker == null) {
+                        relaxedMonotonicityTracker = new ScoreTracker.RelaxedMonotonicityTracker();
+                    } else {
+                        relaxedMonotonicityTracker.reset(rerankK);
+                    }
+                    scoreTracker = relaxedMonotonicityTracker;
+                } else {
+                    if (noOpTracker == null) {
+                        noOpTracker = new ScoreTracker.NoOpTracker();
+                    }
+                    scoreTracker = noOpTracker;
+                }
+            }
+            return scoreTracker;
+        }
+    }
 
     ScoreTracker NO_OP = new NoOpTracker();
 
@@ -59,11 +100,22 @@ class TwoPhaseTracker implements ScoreTracker {
         // observation count
         private int observationCount;
 
-        private final double threshold;
+        private double threshold;
 
         TwoPhaseTracker(double threshold) {
             this.recentScores = new double[RECENT_SCORES_TRACKED];
             this.bestScores = new BoundedLongHeap(BEST_SCORES_TRACKED);
+            this.observationCount = 0;
+            this.threshold = threshold;
+        }
+
+        TwoPhaseTracker() {
+            this(0);
+        }
+
+        void reset(double threshold) {
+            this.bestScores.clear();
+            this.observationCount = 0;
             this.threshold = threshold;
         }
 
@@ -108,10 +160,13 @@ public boolean shouldStop() {
      * (approximately the 96th percentile of the Normal distribution).
      */
     class RelaxedMonotonicityTracker implements ScoreTracker {
-        static final double SIGMA_FACTOR = 1.75;
+        private static final double SIGMA_FACTOR = 1.75;
+
+        private static final int BASE_RECENT_SCORES_SIZE = 200;
 
         // a sliding window of recent scores
-        private final double[] recentScores;
+        private double[] recentScores;
+        private int recentScoresSize;
         private int recentEntryIndex;
 
         // Heap of the best scores seen so far
@@ -132,11 +187,32 @@ class RelaxedMonotonicityTracker implements ScoreTracker {
          *                          the results anymore. An empirical rule of thumb is bestScoresTracked=rerankK.
          */
         RelaxedMonotonicityTracker(int bestScoresTracked) {
+            this.recentScoresSize = getRecentScoresSize(bestScoresTracked);
+            this.recentScores = new double[this.recentScoresSize];
+            this.bestScores = new BoundedLongHeap(bestScoresTracked);
+            this.observationCount = 0;
+            this.mean = 0;
+            this.dSquared = 0;
+        }
+
+        RelaxedMonotonicityTracker() {
+            this(100);
+        }
+
+        private static int getRecentScoresSize(int bestScoresTracked) {
             // A quick empirical study yields that the number of recent scores
             // that we need to consider grows by a factor of ~sqrt(bestScoresTracked / 2)
             int factor = (int) Math.round(Math.sqrt(bestScoresTracked / 2.0));
-            this.recentScores = new double[200 * factor];
-            this.bestScores = new BoundedLongHeap(bestScoresTracked);
+            return BASE_RECENT_SCORES_SIZE * factor;
+        }
+
+        void reset(int bestScoresTracked) {
+            this.recentScoresSize = getRecentScoresSize(bestScoresTracked);
+            if (this.recentScoresSize > recentScores.length) {
+                recentScores = ArrayUtil.grow(recentScores, this.recentScoresSize);
+            }
+            this.bestScores.clear();
+            this.observationCount = 0;
             this.mean = 0;
             this.dSquared = 0;
         }
@@ -150,7 +226,7 @@ public void track(float score) {
             // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online
             // and
             // https://nestedsoftware.com/2019/09/26/incremental-average-and-standard-deviation-with-sliding-window-470k.176143.html
-            if (observationCount <= this.recentScores.length) {
+            if (observationCount <= this.recentScoresSize) {
                 // if the buffer is not full yet, use standard Welford method
                 var meanDelta = (score - this.mean) / observationCount;
                 var newMean = this.mean + meanDelta;
@@ -163,7 +239,7 @@ public void track(float score) {
             } else {
                 // once the buffer is full, adjust Welford method for window size
                 var oldScore = recentScores[recentEntryIndex];
-                var meanDelta = (score - oldScore) / this.recentScores.length;
+                var meanDelta = (score - oldScore) / this.recentScoresSize;
                 var newMean = this.mean + meanDelta;
 
                 var dSquaredDelta = ((score - oldScore) * (score - newMean + oldScore - this.mean));
@@ -173,21 +249,21 @@ public void track(float score) {
                 this.dSquared = newDSquared;
             }
             recentScores[recentEntryIndex] = score;
-            recentEntryIndex = (recentEntryIndex + 1) % this.recentScores.length;
+            recentEntryIndex = (recentEntryIndex + 1) % this.recentScoresSize;
         }
 
         @Override
         public boolean shouldStop() {
             // don't stop if we don't have enough data points
-            if (observationCount < this.recentScores.length) {
+            if (observationCount < this.recentScoresSize) {
                 return false;
             }
 
             // We're in phase 2 if the q-th percentile of the recent scores evaluated,
             //     mean + SIGMA_FACTOR * sqrt(variance),
             // is lower than the worst of the best scores seen.
             // (paper suggests using the median of recent scores, but experimentally that is too prone to false positives)
-            double std = Math.sqrt(this.dSquared / (this.recentScores.length - 1));
+            double std = Math.sqrt(this.dSquared / (this.recentScoresSize - 1));
             double windowPercentile = this.mean + SIGMA_FACTOR * std;
             double worstBestScore = sortableIntToFloat((int) bestScores.top());
             return windowPercentile < worstBestScore;
