integrate into new bulk scorers (#15021)

mccullocht · web-flow · commit 0421a8de771e · 2025-09-10T09:16:46.000-04:00
Return a little extra information from the bulkScore() as this will allow callers to skip results that are below the minimum similarity needed to appear in the result set. Utilize this in hnsw and exhaustive search to skip any topN heap manipulation for any bulk scored batches that do not meet the minimum score. See #14013
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java
@@ -361,21 +361,23 @@ private void search(
           }
           ords[numOrds++] = i;
           if (numOrds == ords.length) {
-            scorer.bulkScore(ords, scores, numOrds);
-            for (int j = 0; j < numOrds; j++) {
-              knnCollector.incVisitedCount(1);
-              knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);
+            knnCollector.incVisitedCount(numOrds);
+            if (scorer.bulkScore(ords, scores, numOrds) > knnCollector.minCompetitiveSimilarity()) {
+              for (int j = 0; j < numOrds; j++) {
+                knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);
+              }
             }
             numOrds = 0;
           }
         }
       }
 
       if (numOrds > 0) {
-        scorer.bulkScore(ords, scores, numOrds);
-        for (int j = 0; j < numOrds; j++) {
-          knnCollector.incVisitedCount(1);
-          knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);
+        knnCollector.incVisitedCount(numOrds);
+        if (scorer.bulkScore(ords, scores, numOrds) > knnCollector.minCompetitiveSimilarity()) {
+          for (int j = 0; j < numOrds; j++) {
+            knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);
+          }
         }
       }
     }
diff --git a/lucene/core/src/java/org/apache/lucene/search/AbstractKnnCollector.java b/lucene/core/src/java/org/apache/lucene/search/AbstractKnnCollector.java
@@ -47,7 +47,7 @@ public final boolean earlyTerminated() {
 
   @Override
   public final void incVisitedCount(int count) {
-    assert count > 0;
+    assert count >= 0;
     this.visitedCount += count;
   }
 
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/HnswGraphSearcher.java
@@ -335,10 +335,11 @@ void searchLevel(
         bulkNodes[numNodes++] = friendOrd;
       }
 
-      if (numNodes > 0) {
-        numNodes = (int) Math.min((long) numNodes, results.visitLimit() - results.visitedCount());
-        scorer.bulkScore(bulkNodes, bulkScores, numNodes);
-        results.incVisitedCount(numNodes);
+      numNodes = (int) Math.min((long) numNodes, results.visitLimit() - results.visitedCount());
+      results.incVisitedCount(numNodes);
+      if (numNodes > 0
+          && scorer.bulkScore(bulkNodes, bulkScores, numNodes)
+              > results.minCompetitiveSimilarity()) {
         for (int i = 0; i < numNodes; i++) {
           int node = bulkNodes[i];
           float score = bulkScores[i];
diff --git a/lucene/core/src/java/org/apache/lucene/util/hnsw/RandomVectorScorer.java b/lucene/core/src/java/org/apache/lucene/util/hnsw/RandomVectorScorer.java
@@ -42,11 +42,15 @@ public interface RandomVectorScorer {
    * @param nodes array of nodes to score.
    * @param scores output array of scores corresponding to each node.
    * @param numNodes number of nodes to score. Must not exceed length of nodes or scores arrays.
+   * @return the maximum scored value of any node, or Float.NEGATIVE_INFINITY if numNodes == 0.
    */
-  default void bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException {
+  default float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException {
+    float max = Float.NEGATIVE_INFINITY;
     for (int i = 0; i < numNodes; i++) {
       scores[i] = score(nodes[i]);
+      max = Math.max(max, scores[i]);
     }
+    return max;
   }
 
   /**
diff --git a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorer.java
@@ -80,19 +80,24 @@ final void checkOrdinal(int ord) {
   }
 
   @Override
-  public void bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException {
+  public float bulkScore(int[] nodes, float[] scores, int numNodes) throws IOException {
     int i = 0;
     final int limit = numNodes & ~3;
+    float maxScore = Float.NEGATIVE_INFINITY;
     for (; i < limit; i += 4) {
       long offset1 = (long) nodes[i] * vectorByteSize;
       long offset2 = (long) nodes[i + 1] * vectorByteSize;
       long offset3 = (long) nodes[i + 2] * vectorByteSize;
       long offset4 = (long) nodes[i + 3] * vectorByteSize;
       vectorOp(seg, scratchScores, offset1, offset2, offset3, offset4, query.length);
       scores[i + 0] = normalizeRawScore(scratchScores[0]);
+      maxScore = Math.max(maxScore, scores[i + 0]);
       scores[i + 1] = normalizeRawScore(scratchScores[1]);
+      maxScore = Math.max(maxScore, scores[i + 1]);
       scores[i + 2] = normalizeRawScore(scratchScores[2]);
+      maxScore = Math.max(maxScore, scores[i + 2]);
       scores[i + 3] = normalizeRawScore(scratchScores[3]);
+      maxScore = Math.max(maxScore, scores[i + 3]);
     }
     // Handle remaining 1–3 nodes in bulk (if any)
     int remaining = numNodes - i;
@@ -102,9 +107,17 @@ public void bulkScore(int[] nodes, float[] scores, int numNodes) throws IOExcept
       long addr3 = (remaining > 2) ? (long) nodes[i + 2] * vectorByteSize : addr1;
       vectorOp(seg, scratchScores, addr1, addr2, addr3, addr3, query.length);
       scores[i] = normalizeRawScore(scratchScores[0]);
-      if (remaining > 1) scores[i + 1] = normalizeRawScore(scratchScores[1]);
-      if (remaining > 2) scores[i + 2] = normalizeRawScore(scratchScores[2]);
+      maxScore = Math.max(maxScore, scores[i]);
+      if (remaining > 1) {
+        scores[i + 1] = normalizeRawScore(scratchScores[1]);
+        maxScore = Math.max(maxScore, scores[i + 1]);
+      }
+      if (remaining > 2) {
+        scores[i + 2] = normalizeRawScore(scratchScores[2]);
+        maxScore = Math.max(maxScore, scores[i + 1]);
+      }
     }
+    return maxScore;
   }
 
   abstract void vectorOp(
diff --git a/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java b/lucene/core/src/java24/org/apache/lucene/internal/vectorization/Lucene99MemorySegmentFloatVectorScorerSupplier.java
@@ -289,9 +289,10 @@ public float score(int node) {
     }
 
     @Override
-    public void bulkScore(int[] nodes, float[] scores, int numNodes) {
+    public float bulkScore(int[] nodes, float[] scores, int numNodes) {
       int i = 0;
       long queryAddr = (long) queryOrd * vectorByteSize;
+      float maxScore = Float.NEGATIVE_INFINITY;
       final int limit = numNodes & ~3;
       for (; i < limit; i += 4) {
         long offset1 = (long) nodes[i] * vectorByteSize;
@@ -300,9 +301,13 @@ public void bulkScore(int[] nodes, float[] scores, int numNodes) {
         long offset4 = (long) nodes[i + 3] * vectorByteSize;
         vectorOp(seg, scratchScores, queryAddr, offset1, offset2, offset3, offset4, dims);
         scores[i + 0] = normalizeRawScore(scratchScores[0]);
+        maxScore = Math.max(maxScore, scores[i + 0]);
         scores[i + 1] = normalizeRawScore(scratchScores[1]);
+        maxScore = Math.max(maxScore, scores[i + 1]);
         scores[i + 2] = normalizeRawScore(scratchScores[2]);
+        maxScore = Math.max(maxScore, scores[i + 2]);
         scores[i + 3] = normalizeRawScore(scratchScores[3]);
+        maxScore = Math.max(maxScore, scores[i + 3]);
       }
       // Handle remaining 1–3 nodes in bulk (if any)
       int remaining = numNodes - i;
@@ -312,9 +317,17 @@ public void bulkScore(int[] nodes, float[] scores, int numNodes) {
         long addr3 = (remaining > 2) ? (long) nodes[i + 2] * vectorByteSize : addr1;
         vectorOp(seg, scratchScores, queryAddr, addr1, addr2, addr3, addr1, dims);
         scores[i] = normalizeRawScore(scratchScores[0]);
-        if (remaining > 1) scores[i + 1] = normalizeRawScore(scratchScores[1]);
-        if (remaining > 2) scores[i + 2] = normalizeRawScore(scratchScores[2]);
+        maxScore = Math.max(maxScore, scores[i]);
+        if (remaining > 1) {
+          scores[i + 1] = normalizeRawScore(scratchScores[1]);
+          maxScore = Math.max(maxScore, scores[i + 1]);
+        }
+        if (remaining > 2) {
+          scores[i + 2] = normalizeRawScore(scratchScores[2]);
+          maxScore = Math.max(maxScore, scores[i + 2]);
+        }
       }
+      return maxScore;
     }
 
     @Override
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/hnsw/TestFlatVectorScorer.java b/lucene/core/src/test/org/apache/lucene/codecs/hnsw/TestFlatVectorScorer.java
@@ -260,11 +260,13 @@ void assertBulkEqualsNonBulk(KnnVectorValues values, VectorSimilarityFunction si
             : flatVectorsScorer.getRandomVectorScorer(sim, values, randomFloatVector(dims));
     int[] indices = randomIndices(size);
     float[] expectedScores = new float[size];
+    float expectedMaxScore = Float.NEGATIVE_INFINITY;
     for (int i = 0; i < size; i++) {
       expectedScores[i] = scorer.score(indices[i]);
+      expectedMaxScore = Math.max(expectedMaxScore, expectedScores[i]);
     }
     float[] bulkScores = new float[size];
-    scorer.bulkScore(indices, bulkScores, size);
+    assertEquals(expectedMaxScore, scorer.bulkScore(indices, bulkScores, size), 0.001);
     assertArrayEquals(expectedScores, bulkScores, delta);
     assertNoScoreBeyondNumNodes(scorer, size);
   }
@@ -302,16 +304,18 @@ void assertScoresAgainstDefaultFlatScorer(KnnVectorValues values, VectorSimilari
         DefaultFlatVectorScorer.INSTANCE.getRandomVectorScorerSupplier(sim, values).scorer();
     defaultScorer.setScoringOrdinal(targetNode);
     float[] expectedScores = new float[size];
+    float expectedMaxScore = Float.NEGATIVE_INFINITY;
     for (int i = 0; i < size; i++) {
       expectedScores[i] = defaultScorer.score(indices[i]);
+      expectedMaxScore = Math.max(expectedMaxScore, expectedScores[i]);
     }
 
     var supplier = flatVectorsScorer.getRandomVectorScorerSupplier(sim, values);
     for (var ss : List.of(supplier, supplier.copy())) {
       var updatableScorer = ss.scorer();
       updatableScorer.setScoringOrdinal(targetNode);
       float[] bulkScores = new float[size];
-      updatableScorer.bulkScore(indices, bulkScores, size);
+      assertEquals(expectedMaxScore, updatableScorer.bulkScore(indices, bulkScores, size), 0.001);
       assertArrayEquals(expectedScores, bulkScores, delta);
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -361,21 +361,23 @@ private void search(`
`361`	`361`	`}`
`362`	`362`	`ords[numOrds++] = i;`
`363`	`363`	`if (numOrds == ords.length) {`
`364`		`- scorer.bulkScore(ords, scores, numOrds);`
`365`		`- for (int j = 0; j < numOrds; j++) {`
`366`		`- knnCollector.incVisitedCount(1);`
`367`		`- knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);`
	`364`	`+ knnCollector.incVisitedCount(numOrds);`
	`365`	`+ if (scorer.bulkScore(ords, scores, numOrds) > knnCollector.minCompetitiveSimilarity()) {`
	`366`	`+ for (int j = 0; j < numOrds; j++) {`
	`367`	`+ knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);`
	`368`	`+ }`
`368`	`369`	`}`
`369`	`370`	`numOrds = 0;`
`370`	`371`	`}`
`371`	`372`	`}`
`372`	`373`	`}`
`373`	`374`
`374`	`375`	`if (numOrds > 0) {`
`375`		`- scorer.bulkScore(ords, scores, numOrds);`
`376`		`- for (int j = 0; j < numOrds; j++) {`
`377`		`- knnCollector.incVisitedCount(1);`
`378`		`- knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);`
	`376`	`+ knnCollector.incVisitedCount(numOrds);`
	`377`	`+ if (scorer.bulkScore(ords, scores, numOrds) > knnCollector.minCompetitiveSimilarity()) {`
	`378`	`+ for (int j = 0; j < numOrds; j++) {`
	`379`	`+ knnCollector.collect(scorer.ordToDoc(ords[j]), scores[j]);`
	`380`	`+ }`
`379`	`381`	`}`
`380`	`382`	`}`
`381`	`383`	`}`
Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ public final boolean earlyTerminated() {`
`47`	`47`
`48`	`48`	`@Override`
`49`	`49`	`public final void incVisitedCount(int count) {`
`50`		`- assert count > 0;`
	`50`	`+ assert count >= 0;`
`51`	`51`	`this.visitedCount += count;`
`52`	`52`	`}`
`53`	`53`