Bulk nextDocsAndScores for constant-score disjunction clauses

lucene/CHANGES.txt

@@ -338,6 +338,11 @@ Improvements
 
 Optimizations
 ---------------------
+* GITHUB#16394: Add DocIdSetIterator#intoArray, which DisjunctionDISIApproximation implements by
+  loading a window of doc IDs through #intoBitSet rather than one nextDoc() call per match.
+  ConstantScoreScorer#nextDocsAndScores now builds on it, which avoids a priority-queue update per
+  matching doc under MaxScoreBulkScorer's term-at-a-time scoring. (Chengxi Shi)
+
 * GITHUG#16280: Single-pass writeString fast path for short strings in ByteBuffersDataOutput (neoremind)
 
 * GITHUB#16146: Speed up no-score filtered-optional Boolean queries by routing dense conjunctions with cached FixedBitSet filters

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

@@ -58,9 +58,27 @@ public long cost() {
 
     @Override
     public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+      if (doc != delegate.docID()) {
+        // The delegate was swapped for an empty iterator (see setMinCompetitiveScore); the
+        // default implementation terminates via nextDoc() without touching the stale delegate
+        // position.
+        super.intoBitSet(upTo, bitSet, offset);
+        return;
+      }
       delegate.intoBitSet(upTo, bitSet, offset);
       doc = delegate.docID();
     }
+
+    @Override
+    public int intoArray(int upTo, int[] docs) throws IOException {
+      if (doc != delegate.docID()) {
+        // See #intoBitSet.
+        return super.intoArray(upTo, docs);
+      }
+      int size = delegate.intoArray(upTo, docs);
+      doc = delegate.docID();
+      return size;
+    }
   }
 
   private final float score;
@@ -154,20 +172,29 @@ public float score() throws IOException {
     return score;
   }
 
+  // Number of doc IDs that a single #nextDocsAndScores call loads at a time. It matches
+  // MaxScoreBulkScorer#INNER_WINDOW_SIZE, so that iterators which load doc IDs in bulk, such as
+  // disjunctions, can cover a full inner scoring window in a single call.
+  private static final int BATCH_SIZE = 4096;
+
   @Override
   public void nextDocsAndScores(int upTo, Bits liveDocs, DocAndFloatFeatureBuffer buffer)
       throws IOException {
-    int batchSize = 64;
-    buffer.growNoCopy(batchSize);
-    int size = 0;
     DocIdSetIterator iterator = iterator();
-    for (int doc = iterator.docID(); doc < upTo && size < batchSize; doc = iterator.nextDoc()) {
-      if (liveDocs == null || liveDocs.get(doc)) {
-        buffer.docs[size] = doc;
-        ++size;
+    buffer.growNoCopy(BATCH_SIZE);
+    for (; ; ) {
+      buffer.size = iterator.intoArray(upTo, buffer.docs);
+      Arrays.fill(buffer.features, 0, buffer.size, score);
+      if (liveDocs == null || buffer.size == 0) {
+        break;
+      }
+      buffer.apply(liveDocs);
+      // An empty buffer indicates that there are no docs left before upTo. We may be unlucky, and
+      // there are docs left, but all docs from the current batch happen to be marked as deleted.
+      // So we need to iterate until we find a batch that has at least one non-deleted doc.
+      if (buffer.size != 0) {
+        break;
       }
     }
-    Arrays.fill(buffer.features, 0, size, score);
-    buffer.size = size;
   }
 }

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

@@ -31,6 +31,11 @@
  */
 public final class DisjunctionDISIApproximation extends AbstractDocIdSetIterator {
 
+  // Doc ID window that a single #intoArray call loads at a time. It matches
+  // MaxScoreBulkScorer#INNER_WINDOW_SIZE and DenseConjunctionBulkScorer#WINDOW_SIZE so that a
+  // single window covers a full scoring window, with a bit set that stays cache-resident.
+  private static final int WINDOW_SIZE = 4096;
+
   public static DisjunctionDISIApproximation of(
       Collection<? extends DisiWrapper> subIterators, long leadCost) {
 
@@ -44,6 +49,7 @@ public static DisjunctionDISIApproximation of(
   private final long cost;
   private DisiWrapper leadTop;
   private int minOtherDoc;
+  private FixedBitSet window; // lazily allocated by #intoArray
 
   public DisjunctionDISIApproximation(
       Collection<? extends DisiWrapper> subIterators, long leadCost) {
@@ -161,6 +167,32 @@ public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOExcept
     doc = Math.min(leadTop.doc, minOtherDoc);
   }
 
+  @Override
+  public int intoArray(int upTo, int[] docs) throws IOException {
+    assert docs.length > 0;
+    // Loading a window of doc IDs through #intoBitSet performs one bulk load and one heap update
+    // per sub-iterator, while advancing one doc at a time performs a heap update per matching doc.
+    // The current doc always falls into the first window, so the loop below normally runs a single
+    // iteration. It only guards against a window that yields no doc, which callers would read as
+    // "no doc left below upTo".
+    int windowSize = Math.min(docs.length, WINDOW_SIZE);
+    while (doc < upTo) {
+      if (window == null) {
+        window = new FixedBitSet(WINDOW_SIZE);
+      } else {
+        window.clear();
+      }
+      int windowMin = doc;
+      int windowMax = (int) Math.min(upTo, (long) windowMin + windowSize);
+      intoBitSet(windowMax, window, windowMin);
+      int size = window.intoArray(0, windowMax - windowMin, windowMin, docs);
+      if (size != 0) {
+        return size;
+      }
+    }
+    return 0;
+  }
+
   /** Return the linked list of iterators positioned on the current doc. */
   public DisiWrapper topList() {
     if (leadTop.doc < minOtherDoc) {

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

@@ -209,6 +209,39 @@ public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOExcept
     }
   }
 
+  /**
+   * Copy doc IDs of this iterator into the given array, starting at the current {@link #docID()}
+   * and stopping before {@code upTo}, and return the number of copied doc IDs. At most {@code
+   * docs.length} doc IDs are copied, so callers must call this method repeatedly in order to
+   * consume all doc IDs below {@code upTo}. Upon return, this iterator is positioned on the first
+   * doc ID that has not been copied.
+   *
+   * <p>A return value of {@code 0} means that this iterator has no doc ID left below {@code upTo}.
+   * Implementations must never return {@code 0} while doc IDs below {@code upTo} remain.
+   *
+   * <p>This should behave exactly as if implemented as below, which is the default implementation:
+   *
+   * <pre><code class="language-java">
+   * int size = 0;
+   * for (int doc = docID(); doc &lt; upTo &amp;&amp; size &lt; docs.length; doc = nextDoc()) {
+   *   docs[size++] = doc;
+   * }
+   * return size;
+   * </code></pre>
+   *
+   * <p><b>Note</b>: The given array must not be empty. Behaviour is undefined if this iterator is
+   * unpositioned.
+   *
+   * @lucene.internal
+   */
+  public int intoArray(int upTo, int[] docs) throws IOException {
+    int size = 0;
+    for (int doc = docID(); doc < upTo && size < docs.length; doc = nextDoc()) {
+      docs[size++] = doc;
+    }
+    return size;
+  }
+
   /**
    * Returns the end of the run of consecutive doc IDs that match this {@link DocIdSetIterator} and
    * that contains the current {@link #docID()}, that is: one plus the last doc ID of the run.

lucene/core/src/test/org/apache/lucene/search/TestConstantScoreScorer.java

@@ -20,6 +20,7 @@
 import static org.apache.lucene.search.DocIdSetIterator.NO_MORE_DOCS;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.document.Document;
@@ -32,7 +33,12 @@
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.analysis.MockAnalyzer;
 import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.apache.lucene.tests.search.AssertingScorer;
 import org.apache.lucene.tests.util.LuceneTestCase;
+import org.apache.lucene.tests.util.TestUtil;
+import org.apache.lucene.util.BitSetIterator;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.FixedBitSet;
 
 public class TestConstantScoreScorer extends LuceneTestCase {
   private static final String FIELD = "f";
@@ -221,6 +227,158 @@ public void close() throws IOException {
     }
   }
 
+  /**
+   * An empty {@link DocAndFloatFeatureBuffer} tells callers that the iterator has no doc left
+   * before {@code upTo}, so batches whose docs are all deleted must not be reported as such.
+   */
+  public void testNextDocsAndScoresSkipsFullyDeletedBatches() throws IOException {
+    int maxDoc = 10_000;
+    int firstLiveDoc = 9_000;
+    Bits liveDocs =
+        new Bits() {
+          @Override
+          public boolean get(int index) {
+            return index >= firstLiveDoc;
+          }
+
+          @Override
+          public int length() {
+            return maxDoc;
+          }
+        };
+
+    for (ScoreMode scoreMode : new ScoreMode[] {ScoreMode.COMPLETE, ScoreMode.TOP_SCORES}) {
+      for (boolean disjunction : new boolean[] {false, true}) {
+        DocIdSetIterator disi;
+        if (disjunction) {
+          disi =
+              DisjunctionDISIApproximation.of(
+                  List.of(
+                      new DisiWrapper(
+                          new ConstantScoreScorer(
+                              1f, ScoreMode.COMPLETE_NO_SCORES, DocIdSetIterator.range(0, 5_000)),
+                          false),
+                      new DisiWrapper(
+                          new ConstantScoreScorer(
+                              1f,
+                              ScoreMode.COMPLETE_NO_SCORES,
+                              DocIdSetIterator.range(5_000, maxDoc)),
+                          false)),
+                  maxDoc);
+        } else {
+          disi = DocIdSetIterator.all(maxDoc);
+        }
+
+        ConstantScoreScorer scorer = new ConstantScoreScorer(2f, scoreMode, disi);
+        assertEquals(0, scorer.iterator().nextDoc());
+
+        DocAndFloatFeatureBuffer buffer = new DocAndFloatFeatureBuffer();
+        List<Integer> collected = new ArrayList<>();
+        for (scorer.nextDocsAndScores(maxDoc, liveDocs, buffer);
+            buffer.size > 0;
+            scorer.nextDocsAndScores(maxDoc, liveDocs, buffer)) {
+          for (int i = 0; i < buffer.size; ++i) {
+            collected.add(buffer.docs[i]);
+            assertEquals(2f, buffer.features[i], 0f);
+          }
+        }
+
+        List<Integer> expected = new ArrayList<>();
+        for (int doc = firstLiveDoc; doc < maxDoc; ++doc) {
+          expected.add(doc);
+        }
+        assertEquals(expected, collected);
+      }
+    }
+  }
+
+  /**
+   * Randomized companion to the test above. Reaching a fully deleted batch takes a doc ID range
+   * that spans more than one batch together with deletions that cluster, so uniformly random
+   * deletions over the small indexes that most randomized tests build never get there. Scorers are
+   * wrapped in {@link AssertingScorer} so that the buffer contract is checked as well.
+   */
+  public void testNextDocsAndScoresRandomClusteredDeletions() throws IOException {
+    int iters = atLeast(10);
+    for (int iter = 0; iter < iters; ++iter) {
+      int maxDoc = TestUtil.nextInt(random(), 12_000, 30_000);
+      float density = 1f / (1 << random().nextInt(3));
+
+      FixedBitSet matches = new FixedBitSet(maxDoc);
+      for (int doc = 0; doc < maxDoc; ++doc) {
+        if (random().nextFloat() < density) {
+          matches.set(doc);
+        }
+      }
+      if (matches.cardinality() == 0) {
+        continue;
+      }
+
+      // A run of deleted docs wide enough to swallow a whole batch whatever its alignment, plus
+      // scattered deletions.
+      FixedBitSet liveDocs = new FixedBitSet(maxDoc);
+      liveDocs.set(0, maxDoc);
+      int runLength = TestUtil.nextInt(random(), 2 * 4096, 3 * 4096);
+      int runStart = TestUtil.nextInt(random(), 0, maxDoc - runLength);
+      liveDocs.clear(runStart, runStart + runLength);
+      for (int i = 0, scattered = random().nextInt(100); i < scattered; ++i) {
+        liveDocs.clear(random().nextInt(maxDoc));
+      }
+
+      List<Integer> expected = new ArrayList<>();
+      for (int doc = 0; doc < maxDoc; ++doc) {
+        if (matches.get(doc) && liveDocs.get(doc)) {
+          expected.add(doc);
+        }
+      }
+
+      DocIdSetIterator disi;
+      if (random().nextBoolean()) {
+        int numSubs = TestUtil.nextInt(random(), 2, 4);
+        FixedBitSet[] subSets = new FixedBitSet[numSubs];
+        for (int i = 0; i < numSubs; ++i) {
+          subSets[i] = new FixedBitSet(maxDoc);
+        }
+        for (int doc = 0; doc < maxDoc; ++doc) {
+          if (matches.get(doc)) {
+            subSets[random().nextInt(numSubs)].set(doc);
+          }
+        }
+        List<DisiWrapper> subs = new ArrayList<>();
+        for (FixedBitSet subSet : subSets) {
+          subs.add(
+              new DisiWrapper(
+                  new ConstantScoreScorer(
+                      1f,
+                      ScoreMode.COMPLETE_NO_SCORES,
+                      new BitSetIterator(subSet, subSet.cardinality())),
+                  false));
+        }
+        disi = DisjunctionDISIApproximation.of(subs, maxDoc);
+      } else {
+        disi = new BitSetIterator(matches, matches.cardinality());
+      }
+
+      ScoreMode scoreMode = random().nextBoolean() ? ScoreMode.COMPLETE : ScoreMode.TOP_SCORES;
+      Scorer scorer =
+          AssertingScorer.wrap(new ConstantScoreScorer(2f, scoreMode, disi), true, false);
+      scorer.iterator().nextDoc();
+
+      DocAndFloatFeatureBuffer buffer = new DocAndFloatFeatureBuffer();
+      List<Integer> collected = new ArrayList<>();
+      for (scorer.nextDocsAndScores(maxDoc, liveDocs, buffer);
+          buffer.size > 0;
+          scorer.nextDocsAndScores(maxDoc, liveDocs, buffer)) {
+        for (int i = 0; i < buffer.size; ++i) {
+          collected.add(buffer.docs[i]);
+          assertEquals(2f, buffer.features[i], 0f);
+        }
+      }
+
+      assertEquals(expected, collected);
+    }
+  }
+
   public void testEarlyTermination() throws IOException {
     Analyzer analyzer = new MockAnalyzer(random());
     Directory dir = newDirectory();