Help collectors take advantage of bulk-retrieval of doc values. (#15173)

We added `NumericDocValues#longValues` to retrieve multiple values in a single
call and better amortize the cost of virtual function calls in #15149. But it's
not easy today to take advantage of this API in collectors, the best option is
to buffer doc IDs in `LeafCollector#collect`, but then it's not great since
this `LeafCollector#collect` call itself may be virtual, so we may still be
doing 1+ virtual call per collected doc.

This PR adds `DocIdStream#intoArray` to help collectors collect batches of doc
IDs at once and retrieve numeric doc values for them. This way, you can compute
facets/aggregations with less than one virtual call per collected document.

I benchmarked on the geonames dataset, computing the average of the elevation
field with the following collector:

```java
class AverageCollector implements Collector {

  long sum;
  int count;

  @OverRide
  public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
    SortedNumericDocValues sortedValues = context.reader().getSortedNumericDocValues("elevation");
    if (sortedValues == null) {
      throw new CollectionTerminatedException();
    }
    NumericDocValues values = DocValues.unwrapSingleton(sortedValues);
    if (values == null) { // the field is single-valued
      throw new Error();
    }
    return new LeafCollector() {

      int[] docBuffer = new int[64];
      long[] valueBuffer = new long[docBuffer.length];

      @OverRide
      public void setScorer(Scorable scorer) throws IOException {

      }

      @OverRide
      public void collect(DocIdStream stream) throws IOException {
        for (int count = stream.intoArray(docBuffer); count != 0; count = stream.intoArray(docBuffer)) {
          values.longValues(count, docBuffer, valueBuffer, Long.MIN_VALUE);

          for (int i = 0; i < count; ++i) {
            long v = valueBuffer[i];
            sum += v;
            if (v != Long.MIN_VALUE) {
              AverageCollector.this.count++;
            }
          }
        }
      }

      @OverRide
      public void collect(int doc) throws IOException {
        if (values.advanceExact(doc)) {
          sum += values.longValue();
          count += 1;
        }
      }
    };
  }

  @OverRide
  public ScoreMode scoreMode() {
    return ScoreMode.COMPLETE_NO_SCORES;
  }

}
```

The benchmark runs a few queries with different collectors first, so that
`LeafCollector#collect` is polymorphic. I get the following numbers:

| Query |  | Latency without DocIdStream#intoArray (ms) | Latency with DocIdStream#intoArray (ms) |
| - | - | - | - |
| `MatchAllDocsQuery` | Uses `RangeDocIdStream` under the hood | 81 | 65 |
| `featureClass:(S P)` | Matches spots or cities (5M docs), uses `BitSetDocIdStream` under the hood | 54 | 43 |

Note that `NumericDocValues#advanceExact` and `NumericDocValues#longValue` are
not polymorphic in this benchmark since all segments use the same impl on the
`elevation` field. The difference would be bigger if they were.

My profiler suggests that retrieving doc values is the bottleneck when
computing the average on a `MatchAllDocsQuery`, while the bottleneck is a mix
of retrieving doc values and extracting set bits from the bit set when
computing the average on `featureClass:(S P)`.

Author: jpountz

lucene/CHANGES.txt

@@ -115,6 +115,10 @@ API Changes
 * GITHUB#15149: Introduce NumericDocValues#longValues to help speed up the
   retrieval of many doc values at once. (Adrien Grand)
 
+* GITHUB#15173: Introduce DocIdStream#intoArray to help collect batches of doc
+  IDs at once, in particular by taking advantage of the new
+  NumericDocValues#longValues API. (Adrien Grand)
+
 New Features
 ---------------------
 * GITHUB#14565: Add ParentsChildrenBlockJoinQuery that supports parent and child filter in the same query

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/BitSetUtil.java

@@ -1039,7 +1039,8 @@ public void nextPostings(int upTo, DocAndFloatFeatureBuffer buffer) throws IOExc
       }
 
       // Only return docs from the current block
-      // +1 to make BitSetUtil#denseBitsetToArray work, see its java doc.
+      // +1 to make FixedBitSet#intoArray perform a bit faster, it uses a slower path if it doesn't
+      // have a free slot after the last element
       buffer.growNoCopy(BLOCK_SIZE + 1);
       upTo = (int) Math.min(upTo, level0LastDocID + 1L);
 
@@ -1055,8 +1056,9 @@ public void nextPostings(int upTo, DocAndFloatFeatureBuffer buffer) throws IOExc
           break;
         case UNARY:
           buffer.size =
-              BitSetUtil.denseBitsetToArray(
-                  docBitSet, doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+              docBitSet.intoArray(
+                  doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+          assert buffer.size < buffer.docs.length; // buffer's capacity is BLOCK_SIZE+1
           break;
       }
 

lucene/backward-codecs/src/java/org/apache/lucene/backward_codecs/lucene103/Lucene103PostingsReader.java

@@ -1047,7 +1047,8 @@ public void nextPostings(int upTo, DocAndFloatFeatureBuffer buffer) throws IOExc
       }
 
       // Only return docs from the current block
-      // +1 to make BitSetUtil#denseBitsetToArray work, see its java doc.
+      // +1 to make FixedBitSet#intoArray perform a bit faster, it uses a slower path if it doesn't
+      // have a free slot after the last element
       buffer.growNoCopy(BLOCK_SIZE + 1);
       upTo = (int) Math.min(upTo, level0LastDocID + 1L);
 
@@ -1063,8 +1064,9 @@ public void nextPostings(int upTo, DocAndFloatFeatureBuffer buffer) throws IOExc
           break;
         case UNARY:
           buffer.size =
-              BitSetUtil.denseBitsetToArray(
-                  docBitSet, doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+              docBitSet.intoArray(
+                  doc - docBitSetBase, upTo - docBitSetBase, docBitSetBase, buffer.docs);
+          assert buffer.size < buffer.docs.length; // buffer's capacity is BLOCK_SIZE+1
           break;
       }
 

lucene/backward-codecs/src/test/org/apache/lucene/backward_codecs/lucene103/TestBitSetUtil.java

@@ -59,4 +59,18 @@ public int count(int upTo) throws IOException {
       return 0;
     }
   }
+
+  @Override
+  public int intoArray(int upTo, int[] array) {
+    if (upTo > this.upTo) {
+      upTo = Math.min(upTo, max);
+      int count = bitSet.intoArray(this.upTo - offset, upTo - offset, offset, array);
+      if (count == array.length) { // The whole range of doc IDs may not have been copied
+        upTo = array[array.length - 1] + 1;
+      }
+      this.upTo = upTo;
+      return count;
+    }
+    return 0;
+  }
 }

lucene/core/src/java/org/apache/lucene/codecs/lucene104/BitSetUtil.java

@@ -266,6 +266,11 @@ public void setScorer(Scorable scorer) throws IOException {
               public void collect(int doc) throws IOException {
                 collector.collect(doc);
               }
+
+              @Override
+              public void collect(DocIdStream stream) throws IOException {
+                collector.collect(stream);
+              }
             };
         return scorer.score(noScoreCollector, acceptDocs, min, max);
       }

lucene/core/src/java/org/apache/lucene/codecs/lucene104/Lucene104PostingsReader.java

@@ -58,6 +58,22 @@ public int count() throws IOException {
   // so would defeat the purpose of collecting hits via a DocIdStream.
   public abstract int count(int upTo) throws IOException;
 
+  /**
+   * Copy some matching doc IDs into the provided array and return the number of copied elements. A
+   * return value of {@code 0} indicates that there are no remaining doc IDs. The given array must
+   * not be empty.
+   */
+  public int intoArray(int[] array) {
+    return intoArray(DocIdSetIterator.NO_MORE_DOCS, array);
+  }
+
+  /**
+   * Copy some matching doc IDs under {@code upTo} (exclusive) into the provided array and return
+   * the number of copied elements. A return value of {@code 0} indicates that there are no matching
+   * doc IDs under {@code upTo} anymore. The given array must not be empty.
+   */
+  public abstract int intoArray(int upTo, int[] array);
+
   /**
    * Return {@code true} if this stream may have remaining doc IDs. This must eventually return
    * {@code false} when the stream is exhausted.