Batch virtual calls from DefaultBulkScorer#scoreRange

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

@@ -123,6 +123,31 @@ default void collect(DocIdStream stream) throws IOException {
     stream.forEach(this::collect);
   }
 
+  /**
+   * Bulk-collect doc IDs.
+   *
+   * <p>Note: The provided int[] may be reused across calls and should be consumed immediately.
+   *
+   * <p>Note: The provided int[] typically only holds a small subset of query matches. This method
+   * may be called multiple times per segment.
+   *
+   * <p>Like {@link #collect(int)}, it is guaranteed that doc IDs get collected in order, ie. doc
+   * IDs are collected in order within a int[], and if called twice, all doc IDs from the second
+   * int[] will be greater than all doc IDs from the first int[].
+   *
+   * <p>It is legal for callers to mix calls to {@link #collect(int[], int)}, {@link
+   * #collect(DocIdStream)} and {@link #collect(int)}.
+   *
+   * <p>The default implementation calls {@code for(int i = 0; i < count; ++i) { collect(docs[i]);
+   * }; }.
+   */
+  default void collect(int[] docs, int count) throws IOException {
+    for (int i = 0; i < count; ++i) {
+      collect(docs[i]);
+    }
+    ;
+  }
+
   /**
    * Optionally returns an iterator over competitive documents.
    *

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

@@ -127,6 +127,11 @@ public void collect(int doc) throws IOException {
         }
       }
 
+      @Override
+      public void collect(int[] docs, int count) {
+        collect(docs, count);
+      }
+
       private void collectCompetitiveHit(int doc, float score) throws IOException {
         final long code = DocScoreEncoder.encode(doc + docBase, score);
         topCode = heap.updateTop(code);

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

@@ -275,52 +275,65 @@ public int score(LeafCollector collector, Bits acceptDocs, int min, int max)
       // collect() because only a subset of collectors produce a competitive iterator, and the set
       // of implementing classes for two-phase approximations is smaller than the set of doc id set
       // iterator implementations.
+
+      // Is it better to initialize the buffer within each iterator implementation?
+      int[] docBuffer = new int[64];
       if (twoPhase == null && competitiveIterator == null) {
         // Optimize simple iterators with collectors that can't skip
-        scoreIterator(collector, acceptDocs, iterator, max);
+        scoreIterator(collector, acceptDocs, iterator, max, docBuffer);
       } else if (competitiveIterator == null) {
-        scoreTwoPhaseIterator(collector, acceptDocs, iterator, twoPhase, max);
+        scoreTwoPhaseIterator(collector, acceptDocs, iterator, twoPhase, max, docBuffer);
       } else if (twoPhase == null) {
-        scoreCompetitiveIterator(collector, acceptDocs, iterator, competitiveIterator, max);
+        scoreCompetitiveIterator(
+            collector, acceptDocs, iterator, competitiveIterator, max, docBuffer);
       } else {
         scoreTwoPhaseOrCompetitiveIterator(
-            collector, acceptDocs, iterator, twoPhase, competitiveIterator, max);
+            collector, acceptDocs, iterator, twoPhase, competitiveIterator, max, docBuffer);
       }
 
       return iterator.docID();
     }
 
     private static void scoreIterator(
-        LeafCollector collector, Bits acceptDocs, DocIdSetIterator iterator, int max)
+        LeafCollector collector, Bits acceptDocs, DocIdSetIterator iterator, int max, int[] docs)
         throws IOException {
+      int count = 0;
       for (int doc = iterator.docID(); doc < max; doc = iterator.nextDoc()) {
         if (acceptDocs == null || acceptDocs.get(doc)) {
-          collector.collect(doc);
+          count = collect(collector, docs, count, doc);
         }
       }
+
+      collect(collector, docs, count, -1);
     }
 
     private static void scoreTwoPhaseIterator(
         LeafCollector collector,
         Bits acceptDocs,
         DocIdSetIterator iterator,
         TwoPhaseIterator twoPhase,
-        int max)
+        int max,
+        int[] docs)
         throws IOException {
+      int count = 0;
       for (int doc = iterator.docID(); doc < max; doc = iterator.nextDoc()) {
         if ((acceptDocs == null || acceptDocs.get(doc)) && twoPhase.matches()) {
-          collector.collect(doc);
+          count = collect(collector, docs, count, doc);
         }
       }
+
+      collect(collector, docs, count, -1);
     }
 
     private static void scoreCompetitiveIterator(
         LeafCollector collector,
         Bits acceptDocs,
         DocIdSetIterator iterator,
         DocIdSetIterator competitiveIterator,
-        int max)
+        int max,
+        int[] docs)
         throws IOException {
+      int count = 0;
       for (int doc = iterator.docID(); doc < max; ) {
         assert competitiveIterator.docID() <= doc; // invariant
         if (competitiveIterator.docID() < doc) {
@@ -332,11 +345,13 @@ private static void scoreCompetitiveIterator(
         }
 
         if ((acceptDocs == null || acceptDocs.get(doc))) {
-          collector.collect(doc);
+          count = collect(collector, docs, count, doc);
         }
 
         doc = iterator.nextDoc();
       }
+
+      collect(collector, docs, count, -1);
     }
 
     private static void scoreTwoPhaseOrCompetitiveIterator(
@@ -345,8 +360,10 @@ private static void scoreTwoPhaseOrCompetitiveIterator(
         DocIdSetIterator iterator,
         TwoPhaseIterator twoPhase,
         DocIdSetIterator competitiveIterator,
-        int max)
+        int max,
+        int[] docs)
         throws IOException {
+      int count = 0;
       for (int doc = iterator.docID(); doc < max; ) {
         assert competitiveIterator.docID() <= doc; // invariant
         if (competitiveIterator.docID() < doc) {
@@ -358,11 +375,26 @@ private static void scoreTwoPhaseOrCompetitiveIterator(
         }
 
         if ((acceptDocs == null || acceptDocs.get(doc)) && twoPhase.matches()) {
-          collector.collect(doc);
+          count = collect(collector, docs, count, doc);
         }
 
         doc = iterator.nextDoc();
       }
+
+      collect(collector, docs, count, -1);
+    }
+
+    private static int collect(LeafCollector collector, int[] docs, int count, int docId)
+        throws IOException {
+      if (count == docs.length || docId == -1) {
+        collector.collect(docs, count);
+        count = 0;
+      }
+
+      // count is always expected to be less than docs.length
+      docs[count++] = docId;
+
+      return count;
     }
   }
 }