Add bulk scoring option to VectorScorer interface (#15171)

Now that we have bulk scoring for our RandomScorer's, let's enable it for our VectorScorer interface. 

I only implemented it for float32 right now, as that is the only place where we are actually doing bulk scoring. But once the interface is in place, I imagine it can be easily implemented elsewhere.

This shows a marginal performance improvement:

baseline
```
recall  latency(ms)  netCPU  avgCpuCount     nDoc  topK  fanout  maxConn  beamWidth  quantized  index(s)  index_docs/s  force_merge(s)  num_segments  index_size(MB)  vec_disk(MB)  vec_RAM(MB)  indexType
 1.000        0.880   0.760        0.864  1000000   500     550       16        100         no      0.00      Infinity            0.06             0            0.00         0.000        0.000       HNSW
```
candidate:
```
recall  latency(ms)  netCPU  avgCpuCount     nDoc  topK  fanout  maxConn  beamWidth  quantized  index(s)  index_docs/s  force_merge(s)  num_segments  index_size(MB)  vec_disk(MB)  vec_RAM(MB)  indexType
 1.000        0.780   0.640        0.821  1000000   500     550       16        100         no      0.00      Infinity            0.06             0            0.00         0.000        0.000       HNSW
```

Author: benwtrent

lucene/CHANGES.txt

@@ -141,6 +141,9 @@ Optimizations
 * GITHUB#15160: Increased the size used for blocks of postings from 128 to 256.
   This gives a noticeable speedup to many queries. (Adrien Grand)
 
+* GITHUB#15171: Add `VectorScorer.Bulk` for bulk iteration and scoring of vectors. This new interface is now
+  used in AbstractKnnVectorQuery to make exact matches faster. (Ben Trent)
+
 Bug Fixes
 ---------------------
 * GITHUB#14161: PointInSetQuery's constructor now throws IllegalArgumentException

lucene/core/src/java/org/apache/lucene/codecs/lucene95/OffHeapFloatVectorValues.java

@@ -18,15 +18,19 @@
 package org.apache.lucene.codecs.lucene95;
 
 import java.io.IOException;
+import java.util.List;
 import org.apache.lucene.codecs.hnsw.FlatVectorsScorer;
 import org.apache.lucene.codecs.lucene90.IndexedDISI;
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.index.VectorEncoding;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.search.ConjunctionUtils;
+import org.apache.lucene.search.DocAndFloatFeatureBuffer;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.VectorScorer;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.RandomAccessInput;
+import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.Bits;
 import org.apache.lucene.util.hnsw.RandomVectorScorer;
 import org.apache.lucene.util.packed.DirectMonotonicReader;
@@ -173,6 +177,30 @@ public float score() throws IOException {
         public DocIdSetIterator iterator() {
           return iterator;
         }
+
+        @Override
+        public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) {
+          final DocIdSetIterator matches =
+              matchingDocs == null
+                  ? iterator
+                  : ConjunctionUtils.createConjunction(List.of(matchingDocs, iterator), List.of());
+          return (nextCount, liveDocs, buffer) -> {
+            if (matches.docID() == -1) {
+              matches.nextDoc();
+            }
+            buffer.growNoCopy(nextCount);
+            int size = 0;
+            for (int doc = matches.docID();
+                doc != DocIdSetIterator.NO_MORE_DOCS && size < nextCount;
+                doc = matches.nextDoc()) {
+              if (liveDocs == null || liveDocs.get(doc)) {
+                buffer.docs[size++] = doc;
+              }
+            }
+            buffer.size = size;
+            return randomVectorScorer.bulkScore(buffer.docs, buffer.features, size);
+          };
+        }
       };
     }
   }
@@ -266,6 +294,43 @@ public float score() throws IOException {
         public DocIdSetIterator iterator() {
           return iterator;
         }
+
+        @Override
+        public VectorScorer.Bulk bulk(DocIdSetIterator matchingDocs) {
+          return new Bulk() {
+            final DocIdSetIterator matches =
+                matchingDocs == null
+                    ? iterator
+                    : ConjunctionUtils.createConjunction(
+                        List.of(matchingDocs, iterator), List.of());
+            int[] docIds = new int[0];
+
+            @Override
+            public float nextDocsAndScores(
+                int nextCount, Bits liveDocs, DocAndFloatFeatureBuffer buffer) throws IOException {
+              if (matches.docID() == -1) {
+                matches.nextDoc();
+              }
+              buffer.growNoCopy(nextCount);
+              docIds = ArrayUtil.growNoCopy(docIds, nextCount);
+              int size = 0;
+              for (int doc = matches.docID();
+                  doc != DocIdSetIterator.NO_MORE_DOCS && size < nextCount;
+                  doc = matches.nextDoc()) {
+                if (liveDocs == null || liveDocs.get(doc)) {
+                  buffer.docs[size] = iterator.index();
+                  docIds[size] = doc;
+                  ++size;
+                }
+              }
+              buffer.size = size;
+              float maxScore = randomVectorScorer.bulkScore(buffer.docs, buffer.features, size);
+              // copy back the real doc IDs
+              System.arraycopy(docIds, 0, buffer.docs, 0, size);
+              return maxScore;
+            }
+          };
+        }
       };
     }
   }

lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99HnswVectorsReader.java

@@ -339,7 +339,8 @@ private void search(
     int filteredDocCount = Math.min(acceptDocs.cost(), graph.size());
     Bits accepted = acceptDocs.bits();
     final Bits acceptedOrds = scorer.getAcceptOrds(accepted);
-    boolean doHnsw = knnCollector.k() < scorer.maxOrd();
+    int numVectors = scorer.maxOrd();
+    boolean doHnsw = knnCollector.k() < numVectors;
     // The approximate number of vectors that would be visited if we did not filter
     int unfilteredVisit = HnswGraphSearcher.expectedVisitedNodes(knnCollector.k(), graph.size());
     if (unfilteredVisit >= filteredDocCount) {
@@ -354,7 +355,7 @@ private void search(
       int[] ords = new int[EXHAUSTIVE_BULK_SCORE_ORDS];
       float[] scores = new float[EXHAUSTIVE_BULK_SCORE_ORDS];
       int numOrds = 0;
-      for (int i = 0; i < scorer.maxOrd(); i++) {
+      for (int i = 0; i < numVectors; i++) {
         if (acceptedOrds == null || acceptedOrds.get(i)) {
           if (knnCollector.earlyTerminated()) {
             break;

lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java

@@ -309,21 +309,28 @@ protected TopDocs exactSearch(
     TotalHits.Relation relation = TotalHits.Relation.EQUAL_TO;
     ScoreDoc topDoc = queue.top();
     DocIdSetIterator vectorIterator = vectorScorer.iterator();
-    DocIdSetIterator conjunction =
-        ConjunctionDISI.createConjunction(List.of(vectorIterator, acceptIterator), List.of());
-    int doc;
-    while ((doc = conjunction.nextDoc()) != DocIdSetIterator.NO_MORE_DOCS) {
+    DocAndFloatFeatureBuffer buffer = new DocAndFloatFeatureBuffer();
+    VectorScorer.Bulk bulkScorer = vectorScorer.bulk(acceptIterator);
+    while (vectorIterator.docID() != DocIdSetIterator.NO_MORE_DOCS) {
       // Mark results as partial if timeout is met
       if (queryTimeout != null && queryTimeout.shouldExit()) {
         relation = TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO;
         break;
       }
-      assert vectorIterator.docID() == doc;
-      float score = vectorScorer.score();
-      if (score > topDoc.score) {
-        topDoc.score = score;
-        topDoc.doc = doc;
-        topDoc = queue.updateTop();
+      // iterator already takes live docs into account
+      float maxScore = bulkScorer.nextDocsAndScores(64, null, buffer);
+      if (maxScore < topDoc.score) {
+        // all the scores in this batch are too low, skip
+        continue;
+      }
+      for (int i = 0; i < buffer.size; i++) {
+        float score = buffer.features[i];
+        int doc = buffer.docs[i];
+        if (score > topDoc.score) {
+          topDoc.score = score;
+          topDoc.doc = doc;
+          topDoc = queue.updateTop();
+        }
       }
     }
 

lucene/core/src/java/org/apache/lucene/search/VectorScorer.java

@@ -17,6 +17,8 @@
 package org.apache.lucene.search;
 
 import java.io.IOException;
+import java.util.List;
+import org.apache.lucene.util.Bits;
 
 /**
  * Computes the similarity score between a given query vector and different document vectors. This
@@ -38,4 +40,64 @@ public interface VectorScorer {
    * @return a {@link DocIdSetIterator} over the documents.
    */
   DocIdSetIterator iterator();
+
+  /**
+   * An optional bulk scorer implementation that allows bulk scoring over the provided matching
+   * docs. The iterator of this instance of VectorScorer should be used and iterated in conjunction
+   * with the provided matchingDocs iterator to score only the documents that are present in both
+   * iterators. If the provided matchingDocs iterator is null, then all documents should be scored.
+   * Additionally, if the iterators are unpositioned (docID() == -1), this method should position
+   * them to the first document.
+   *
+   * @param matchingDocs the documents to score
+   * @return a {@link Bulk} scorer
+   * @throws IOException if an exception occurs during bulk scorer creation
+   * @lucene.experimental
+   */
+  default Bulk bulk(DocIdSetIterator matchingDocs) throws IOException {
+    final DocIdSetIterator iterator =
+        matchingDocs == null
+            ? iterator()
+            : ConjunctionUtils.createConjunction(List.of(matchingDocs, iterator()), List.of());
+    if (iterator.docID() == -1) {
+      iterator.nextDoc();
+    }
+    return (nextCount, liveDocs, buffer) -> {
+      buffer.growNoCopy(nextCount);
+      int size = 0;
+      float maxScore = Float.NEGATIVE_INFINITY;
+      for (int doc = iterator.docID();
+          doc != DocIdSetIterator.NO_MORE_DOCS && size < nextCount;
+          doc = iterator.nextDoc()) {
+        if (liveDocs == null || liveDocs.get(doc)) {
+          buffer.docs[size] = doc;
+          buffer.features[size] = score();
+          maxScore = Math.max(maxScore, buffer.features[size]);
+          ++size;
+        }
+      }
+      buffer.size = size;
+      return maxScore;
+    };
+  }
+
+  /**
+   * Bulk scorer interface to score multiple vectors at once
+   *
+   * @lucene.experimental
+   */
+  interface Bulk {
+    /**
+     * Score up to nextCount documents, store the results in the provided buffer. Behaves similarly
+     * to {@link Scorer#nextDocsAndScores(int, Bits, DocAndFloatFeatureBuffer)}
+     *
+     * @param nextCount the maximum number of documents to score
+     * @param liveDocs the live docs, or null if all docs are live
+     * @param buffer the buffer to store the results
+     * @return the max score of the scored documents
+     * @throws IOException if an exception occurs during scoring
+     */
+    float nextDocsAndScores(int nextCount, Bits liveDocs, DocAndFloatFeatureBuffer buffer)
+        throws IOException;
+  }
 }