Speed up advancing within a block. (#13692)

Advancing within a block consists of finding the first index within an array of
128 values whose value is greater than or equal a target. Given the small size,
it's not obvious whether it's better to perform a linear search, a binary
search or something else... It is surprisingly hard to beat the linear search
that we are using today.

Experiments suggested that the following approach works in practice:
 - First check if the next item in the array is greater than or equal to the
   target.
 - Then find the first 4-values interval that contains our target.
 - Then perform a branchless binary search within this interval of 4 values.

This approach still biases heavily towards the case when the target is very
close to the current index, only a bit less than a linear search.

Author: jpountz

lucene/core/src/java/org/apache/lucene/codecs/lucene912/Lucene912PostingsReader.java

@@ -67,6 +67,7 @@
 public final class Lucene912PostingsReader extends PostingsReaderBase {
 
   static final VectorizationProvider VECTORIZATION_PROVIDER = VectorizationProvider.getInstance();
+  static int BINARY_SEARCH_WINDOW_SIZE = 4;
 
   private final IndexInput docIn;
   private final IndexInput posIn;
@@ -212,13 +213,74 @@ static void prefixSum(long[] buffer, int count, long base) {
     }
   }
 
-  static int findFirstGreater(long[] buffer, int target, int from) {
-    for (int i = from; i < BLOCK_SIZE; ++i) {
-      if (buffer[i] >= target) {
-        return i;
+  private static boolean assertDocBuffer(long[] values, int start) {
+    assert values.length == BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE;
+    assert values[BLOCK_SIZE] == DocIdSetIterator.NO_MORE_DOCS;
+    assert start < BLOCK_SIZE;
+
+    int endOffset;
+    if (values[0] == DocIdSetIterator.NO_MORE_DOCS) {
+      endOffset = 0;
+    } else {
+      endOffset = -1;
+      for (int i = 1; i < values.length; ++i) {
+        assert values[i] > values[i - 1] : Arrays.toString(values);
+        if (values[i] == DocIdSetIterator.NO_MORE_DOCS) {
+          endOffset = i;
+          break;
+        }
       }
     }
-    return BLOCK_SIZE;
+    for (int i = 0; i < BINARY_SEARCH_WINDOW_SIZE; ++i) {
+      assert values[endOffset + i] == DocIdSetIterator.NO_MORE_DOCS;
+    }
+    return true;
+  }
+
+  /**
+   * Return the first index in sorted array {@code values} whose value is greater than or equal to
+   * {@code target}. For correctness, it requires the last 4 values to be set to {@code
+   * NO_MORE_DOCS}.
+   */
+  static int findNextGEQ(long[] values, long target, int start) {
+    assert assertDocBuffer(values, start);
+
+    if (values[start] >= target) {
+      // Surprisingly this is a likely condition in practice, so optimizing for it helps.
+      return start;
+    }
+
+    // We just ruled out that our target index is at `start`.
+    start += 1;
+
+    // Now find the first interval of 4 values that contains our target.
+    for (int i = start;
+        i + BINARY_SEARCH_WINDOW_SIZE <= values.length;
+        i += BINARY_SEARCH_WINDOW_SIZE) {
+      if (values[i + BINARY_SEARCH_WINDOW_SIZE - 1] >= target) {
+        start = i;
+        break;
+      }
+    }
+
+    // Binary search in this interval of 4 values.
+    return binarySearch4(values, target, start);
+  }
+
+  /**
+   * Return the first index whose value is greater than or equal to {@code target} among the 4
+   * values starting at {@code start}. If none of the values is greater than or equal to {@code
+   * target}, this returns {@code start+3}.
+   */
+  private static int binarySearch4(long[] values, long target, int start) {
+    // This code is organized in a way that compiles to a branchless binary search.
+    if (values[start + 1] < target) {
+      start += 2;
+    }
+    if (values[start] < target) {
+      start += 1;
+    }
+    return start;
   }
 
   @Override
@@ -348,7 +410,7 @@ final class BlockDocsEnum extends PostingsEnum {
     final ForDeltaUtil forDeltaUtil = new ForDeltaUtil();
     final PForUtil pforUtil = new PForUtil(forUtil);
 
-    private final long[] docBuffer = new long[BLOCK_SIZE + 1];
+    private final long[] docBuffer = new long[BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE];
     private final long[] freqBuffer = new long[BLOCK_SIZE];
 
     private int docBufferUpto;
@@ -390,9 +452,9 @@ public BlockDocsEnum(FieldInfo fieldInfo) throws IOException {
                       .compareTo(IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS)
                   >= 0
               || fieldInfo.hasPayloads();
-      // We set the last element of docBuffer to NO_MORE_DOCS, it helps save conditionals in
+      // We set the last elements of docBuffer to NO_MORE_DOCS, it helps save conditionals in
       // advance()
-      docBuffer[BLOCK_SIZE] = NO_MORE_DOCS;
+      Arrays.fill(docBuffer, BLOCK_SIZE, BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
     }
 
     public boolean canReuse(IndexInput docIn, FieldInfo fieldInfo) {
@@ -495,7 +557,7 @@ private void refillFullBlock() throws IOException {
       docCountUpto += BLOCK_SIZE;
       prevDocID = docBuffer[BLOCK_SIZE - 1];
       docBufferUpto = 0;
-      assert docBuffer[BLOCK_SIZE] == NO_MORE_DOCS;
+      assert assertDocBuffer(docBuffer, 0);
     }
 
     private void refillRemainder() throws IOException {
@@ -506,15 +568,14 @@ private void refillRemainder() throws IOException {
       if (docFreq == 1) {
         docBuffer[0] = singletonDocID;
         freqBuffer[0] = totalTermFreq;
-        docBuffer[1] = NO_MORE_DOCS;
         docCountUpto++;
       } else {
         // Read vInts:
         PostingsUtil.readVIntBlock(docIn, docBuffer, freqBuffer, left, indexHasFreq, needsFreq);
         prefixSum(docBuffer, left, prevDocID);
-        docBuffer[left] = NO_MORE_DOCS;
         docCountUpto += left;
       }
+      Arrays.fill(docBuffer, left, left + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
       docBufferUpto = 0;
       freqFP = -1;
     }
@@ -612,7 +673,7 @@ public int advance(int target) throws IOException {
         }
       }
 
-      int next = findFirstGreater(docBuffer, target, docBufferUpto);
+      int next = findNextGEQ(docBuffer, target, docBufferUpto);
       this.doc = (int) docBuffer[next];
       docBufferUpto = next + 1;
       return doc;
@@ -630,8 +691,8 @@ final class EverythingEnum extends PostingsEnum {
     final ForDeltaUtil forDeltaUtil = new ForDeltaUtil();
     final PForUtil pforUtil = new PForUtil(forUtil);
 
-    private final long[] docBuffer = new long[BLOCK_SIZE + 1];
-    private final long[] freqBuffer = new long[BLOCK_SIZE + 1];
+    private final long[] docBuffer = new long[BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE];
+    private final long[] freqBuffer = new long[BLOCK_SIZE];
     private final long[] posDeltaBuffer = new long[BLOCK_SIZE];
 
     private final long[] payloadLengthBuffer;
@@ -752,9 +813,9 @@ public EverythingEnum(FieldInfo fieldInfo) throws IOException {
         payload = null;
       }
 
-      // We set the last element of docBuffer to NO_MORE_DOCS, it helps save conditionals in
+      // We set the last elements of docBuffer to NO_MORE_DOCS, it helps save conditionals in
       // advance()
-      docBuffer[BLOCK_SIZE] = NO_MORE_DOCS;
+      Arrays.fill(docBuffer, BLOCK_SIZE, BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
     }
 
     public boolean canReuse(IndexInput docIn, FieldInfo fieldInfo) {
@@ -843,18 +904,18 @@ private void refillDocs() throws IOException {
       } else if (docFreq == 1) {
         docBuffer[0] = singletonDocID;
         freqBuffer[0] = totalTermFreq;
-        docBuffer[1] = NO_MORE_DOCS;
+        Arrays.fill(docBuffer, 1, 1 + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
         docCountUpto++;
       } else {
         // Read vInts:
         PostingsUtil.readVIntBlock(docIn, docBuffer, freqBuffer, left, indexHasFreq, true);
         prefixSum(docBuffer, left, prevDocID);
-        docBuffer[left] = NO_MORE_DOCS;
+        Arrays.fill(docBuffer, left, left + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
         docCountUpto += left;
       }
       prevDocID = docBuffer[BLOCK_SIZE - 1];
       docBufferUpto = 0;
-      assert docBuffer[BLOCK_SIZE] == NO_MORE_DOCS;
+      assert assertDocBuffer(docBuffer, 0);
     }
 
     private void skipLevel1To(int target) throws IOException {
@@ -1016,7 +1077,7 @@ public int advance(int target) throws IOException {
         refillDocs();
       }
 
-      int next = findFirstGreater(docBuffer, target, docBufferUpto);
+      int next = findNextGEQ(docBuffer, target, docBufferUpto);
       for (int i = docBufferUpto; i <= next; ++i) {
         posPendingCount += freqBuffer[i];
       }
@@ -1217,7 +1278,7 @@ final class BlockImpactsDocsEnum extends ImpactsEnum {
     final ForDeltaUtil forDeltaUtil = new ForDeltaUtil();
     final PForUtil pforUtil = new PForUtil(forUtil);
 
-    private final long[] docBuffer = new long[BLOCK_SIZE + 1];
+    private final long[] docBuffer = new long[BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE];
     private final long[] freqBuffer = new long[BLOCK_SIZE];
 
     private int docBufferUpto;
@@ -1265,9 +1326,9 @@ public BlockImpactsDocsEnum(FieldInfo fieldInfo, IntBlockTermState termState)
                       .compareTo(IndexOptions.DOCS_AND_FREQS_AND_POSITIONS_AND_OFFSETS)
                   >= 0
               || fieldInfo.hasPayloads();
-      // We set the last element of docBuffer to NO_MORE_DOCS, it helps save conditionals in
+      // We set the last elements of docBuffer to NO_MORE_DOCS, it helps save conditionals in
       // advance()
-      docBuffer[BLOCK_SIZE] = NO_MORE_DOCS;
+      Arrays.fill(docBuffer, BLOCK_SIZE, BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
 
       docFreq = termState.docFreq;
       if (docFreq > 1) {
@@ -1357,13 +1418,13 @@ private void refillDocs() throws IOException {
         // Read vInts:
         PostingsUtil.readVIntBlock(docIn, docBuffer, freqBuffer, left, indexHasFreq, true);
         prefixSum(docBuffer, left, prevDocID);
-        docBuffer[left] = NO_MORE_DOCS;
+        Arrays.fill(docBuffer, left, left + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
         freqFP = -1;
         docCountUpto += left;
       }
       prevDocID = docBuffer[BLOCK_SIZE - 1];
       docBufferUpto = 0;
-      assert docBuffer[BLOCK_SIZE] == NO_MORE_DOCS;
+      assert assertDocBuffer(docBuffer, 0);
     }
 
     private void skipLevel1To(int target) throws IOException {
@@ -1489,7 +1550,7 @@ public int advance(int target) throws IOException {
         needsRefilling = false;
       }
 
-      int next = findFirstGreater(docBuffer, target, docBufferUpto);
+      int next = findNextGEQ(docBuffer, target, docBufferUpto);
       this.doc = (int) docBuffer[next];
       docBufferUpto = next + 1;
       return doc;
@@ -1572,7 +1633,7 @@ final class BlockImpactsPostingsEnum extends ImpactsEnum {
     final ForDeltaUtil forDeltaUtil = new ForDeltaUtil();
     final PForUtil pforUtil = new PForUtil(forUtil);
 
-    private final long[] docBuffer = new long[BLOCK_SIZE + 1];
+    private final long[] docBuffer = new long[BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE];
     private final long[] freqBuffer = new long[BLOCK_SIZE];
     private final long[] posDeltaBuffer = new long[BLOCK_SIZE];
 
@@ -1652,9 +1713,9 @@ public BlockImpactsPostingsEnum(FieldInfo fieldInfo, IntBlockTermState termState
       this.posIn = Lucene912PostingsReader.this.posIn.clone();
       posInUtil = VECTORIZATION_PROVIDER.newPostingDecodingUtil(posIn);
 
-      // We set the last element of docBuffer to NO_MORE_DOCS, it helps save conditionals in
+      // We set the last elements of docBuffer to NO_MORE_DOCS, it helps save conditionals in
       // advance()
-      docBuffer[BLOCK_SIZE] = NO_MORE_DOCS;
+      Arrays.fill(docBuffer, BLOCK_SIZE, BLOCK_SIZE + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
 
       docFreq = termState.docFreq;
       posTermStartFP = termState.posStartFP;
@@ -1724,18 +1785,18 @@ private void refillDocs() throws IOException {
       } else if (docFreq == 1) {
         docBuffer[0] = singletonDocID;
         freqBuffer[0] = totalTermFreq;
-        docBuffer[1] = NO_MORE_DOCS;
+        Arrays.fill(docBuffer, left, left + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
         docCountUpto++;
       } else {
         // Read vInts:
         PostingsUtil.readVIntBlock(docIn, docBuffer, freqBuffer, left, indexHasFreq, true);
         prefixSum(docBuffer, left, prevDocID);
-        docBuffer[left] = NO_MORE_DOCS;
+        Arrays.fill(docBuffer, left, left + BINARY_SEARCH_WINDOW_SIZE, NO_MORE_DOCS);
         docCountUpto += left;
       }
       prevDocID = docBuffer[BLOCK_SIZE - 1];
       docBufferUpto = 0;
-      assert docBuffer[BLOCK_SIZE] == NO_MORE_DOCS;
+      assert assertDocBuffer(docBuffer, 0);
     }
 
     private void skipLevel1To(int target) throws IOException {
@@ -1932,7 +1993,7 @@ public int advance(int target) throws IOException {
         needsRefilling = false;
       }
 
-      int next = findFirstGreater(docBuffer, target, docBufferUpto);
+      int next = findNextGEQ(docBuffer, target, docBufferUpto);
       for (int i = docBufferUpto; i <= next; ++i) {
         posPendingCount += freqBuffer[i];
       }

lucene/core/src/test/org/apache/lucene/codecs/lucene912/TestLucene912PostingsFormat.java

@@ -31,6 +31,7 @@
 import org.apache.lucene.index.Impact;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.store.ByteArrayDataInput;
 import org.apache.lucene.store.ByteArrayDataOutput;
 import org.apache.lucene.store.Directory;
@@ -154,4 +155,19 @@ private void doTestImpactSerialization(List<Impact> impacts) throws IOException
       }
     }
   }
+
+  public void testFindNextGEQ() {
+    long[] values =
+        new long[ForUtil.BLOCK_SIZE + Lucene912PostingsReader.BINARY_SEARCH_WINDOW_SIZE];
+    for (int i = 0; i < ForUtil.BLOCK_SIZE; ++i) {
+      values[i] = i * 2;
+    }
+    Arrays.fill(values, ForUtil.BLOCK_SIZE, values.length, DocIdSetIterator.NO_MORE_DOCS);
+    for (int i = 0; i < ForUtil.BLOCK_SIZE; ++i) {
+      for (int start = 0; start <= i; ++start) {
+        assertEquals(i, Lucene912PostingsReader.findNextGEQ(values, i * 2, start));
+        assertEquals(i + 1, Lucene912PostingsReader.findNextGEQ(values, i * 2 + 1, start));
+      }
+    }
+  }
 }