Not maintain docBufferUpTo when only docs needed (#14164)

apache · Jan 25, 2025 · 52d3809 · 52d3809
1 parent 8487718
commit 52d3809
Showing 1 changed file with 32 additions and 22 deletions.
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene101/Lucene101PostingsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene101/Lucene101PostingsReader.java
@@ -615,19 +615,21 @@ private void refillFullBlock() throws IOException {
           numLongs = -bitsPerValue;
           docIn.readLongs(docBitSet.getBits(), 0, numLongs);
         }
-        // Note: we know that BLOCK_SIZE bits are set, so no need to compute the cumulative pop
-        // count at the last index, it will be BLOCK_SIZE.
-        // Note: this for loop auto-vectorizes
-        for (int i = 0; i < numLongs - 1; ++i) {
-          docCumulativeWordPopCounts[i] = Long.bitCount(docBitSet.getBits()[i]);
-        }
-        for (int i = 1; i < numLongs - 1; ++i) {
-          docCumulativeWordPopCounts[i] += docCumulativeWordPopCounts[i - 1];
+        if (needsFreq) {
+          // Note: we know that BLOCK_SIZE bits are set, so no need to compute the cumulative pop
+          // count at the last index, it will be BLOCK_SIZE.
+          // Note: this for loop auto-vectorizes
+          for (int i = 0; i < numLongs - 1; ++i) {
+            docCumulativeWordPopCounts[i] = Long.bitCount(docBitSet.getBits()[i]);
+          }
+          for (int i = 1; i < numLongs - 1; ++i) {
+            docCumulativeWordPopCounts[i] += docCumulativeWordPopCounts[i - 1];
+          }
+          docCumulativeWordPopCounts[numLongs - 1] = BLOCK_SIZE;
+          assert docCumulativeWordPopCounts[numLongs - 2]
+                  + Long.bitCount(docBitSet.getBits()[numLongs - 1])
+              == BLOCK_SIZE;
         }
-        docCumulativeWordPopCounts[numLongs - 1] = BLOCK_SIZE;
-        assert docCumulativeWordPopCounts[numLongs - 2]
-                + Long.bitCount(docBitSet.getBits()[numLongs - 1])
-            == BLOCK_SIZE;
         encoding = DeltaEncoding.UNARY;
       }
       if (indexHasFreq) {
@@ -726,7 +728,7 @@ private void skipLevel1To(int target) throws IOException {
     }
 
     private void doMoveToNextLevel0Block() throws IOException {
-      assert docBufferUpto == BLOCK_SIZE;
+      assert doc == level0LastDocID;
       if (posIn != null) {
         if (level0PosEndFP >= posIn.getFilePointer()) {
           posIn.seek(level0PosEndFP);
@@ -912,7 +914,7 @@ private void doAdvanceShallow(int target) throws IOException {
 
     @Override
     public int nextDoc() throws IOException {
-      if (docBufferUpto == BLOCK_SIZE) {
+      if (doc == level0LastDocID) {
         moveToNextLevel0Block();
       }
 
@@ -954,13 +956,21 @@ public int advance(int target) throws IOException {
             int next = docBitSet.nextSetBit(target - docBitSetBase);
             assert next != NO_MORE_DOCS;
             this.doc = docBitSetBase + next;
-            int wordIndex = next >> 6;
-            // Take the cumulative pop count for the given word, and subtract bits on the left of
-            // the current doc.
-            docBufferUpto =
-                1
-                    + docCumulativeWordPopCounts[wordIndex]
-                    - Long.bitCount(docBitSet.getBits()[wordIndex] >>> next);
+            if (needsFreq) {
+              int wordIndex = next >> 6;
+              // Take the cumulative pop count for the given word, and subtract bits on the left of
+              // the current doc.
+              docBufferUpto =
+                  1
+                      + docCumulativeWordPopCounts[wordIndex]
+                      - Long.bitCount(docBitSet.getBits()[wordIndex] >>> next);
+            } else {
+              // When only docs needed and block is UNARY encoded, we do not need to maintain
+              // docBufferUpTo to record the iteration position in the block.
+              // docBufferUpTo == 0 means the block has not been iterated.
+              // docBufferUpTo != 0 means the block has been iterated.
+              docBufferUpto = 1;
+            }
           }
           break;
       }
@@ -978,7 +988,7 @@ public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOExcept
       bitSet.set(doc - offset);
 
       for (; ; ) {
-        if (docBufferUpto == BLOCK_SIZE) {
+        if (doc == level0LastDocID) {
           // refill
           moveToNextLevel0Block();
         }