Impl intoBitset for IndexedDISI and Docvalues (#14529)

Author: gf2121

lucene/core/src/java/org/apache/lucene/codecs/lucene90/IndexedDISI.java

@@ -163,9 +163,9 @@ private static byte[] createRank(FixedBitSet buffer, byte denseRankPower) {
    *
    * @param it the document IDs.
    * @param out destination for the blocks.
-   * @throws IOException if there was an error writing to out.
    * @return the number of jump-table entries following the blocks, -1 for no entries. This should
    *     be stored in meta and used when creating an instance of IndexedDISI.
+   * @throws IOException if there was an error writing to out.
    */
   static short writeBitSet(DocIdSetIterator it, IndexOutput out) throws IOException {
     return writeBitSet(it, out, DEFAULT_DENSE_RANK_POWER);
@@ -184,9 +184,9 @@ static short writeBitSet(DocIdSetIterator it, IndexOutput out) throws IOExceptio
    *     disables DENSE rank. Recommended values are 8-12: Every 256-4096 docIDs or 4-64 longs.
    *     {@link #DEFAULT_DENSE_RANK_POWER} is 9: Every 512 docIDs. This should be stored in meta and
    *     used when creating an instance of IndexedDISI.
-   * @throws IOException if there was an error writing to out.
    * @return the number of jump-table entries following the blocks, -1 for no entries. This should
    *     be stored in meta and used when creating an instance of IndexedDISI.
+   * @throws IOException if there was an error writing to out.
    */
   public static short writeBitSet(DocIdSetIterator it, IndexOutput out, byte denseRankPower)
       throws IOException {
@@ -436,6 +436,7 @@ public static RandomAccessInput createJumpTable(
   int numberOfOnes;
   // Used with rank for jumps inside of DENSE as they are absolute instead of relative
   int denseOrigoIndex;
+  FixedBitSet bitSet;
 
   // ALL variables
   int gap;
@@ -491,6 +492,16 @@ public int advance(int target) throws IOException {
     return doc;
   }
 
+  @Override
+  public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+    assert doc >= offset;
+    while (doc < upTo && method.intoBitSetWithinBlock(this, upTo, bitSet, offset) == false) {
+      readBlockHeader();
+      boolean found = method.advanceWithinBlock(this, block);
+      assert found;
+    }
+  }
+
   public boolean advanceExact(int target) throws IOException {
     final int targetBlock = target & 0xFFFF0000;
     if (block < targetBlock) {
@@ -625,6 +636,25 @@ boolean advanceExactWithinBlock(IndexedDISI disi, int target) throws IOException
         disi.exists = false;
         return false;
       }
+
+      @Override
+      boolean intoBitSetWithinBlock(IndexedDISI disi, int upTo, FixedBitSet bitSet, int offset)
+          throws IOException {
+        bitSet.set(disi.doc - offset);
+        for (; disi.index < disi.nextBlockIndex; ) {
+          int docInBlock = disi.slice.readShort() & 0xFFFF;
+          int doc = disi.block | docInBlock;
+          disi.index++;
+          if (doc >= upTo) {
+            disi.doc = doc;
+            disi.exists = true;
+            disi.nextExistDocInBlock = docInBlock;
+            return true;
+          }
+          bitSet.set(doc - offset);
+        }
+        return false;
+      }
     },
     DENSE {
       @Override
@@ -693,6 +723,34 @@ boolean advanceExactWithinBlock(IndexedDISI disi, int target) throws IOException
         disi.index = disi.numberOfOnes - Long.bitCount(leftBits);
         return (leftBits & 1L) != 0;
       }
+
+      @Override
+      boolean intoBitSetWithinBlock(IndexedDISI disi, int upTo, FixedBitSet bitSet, int offset)
+          throws IOException {
+        if (disi.bitSet == null) {
+          disi.bitSet = new FixedBitSet(BLOCK_SIZE);
+        }
+
+        int sourceFrom = disi.doc & 0xFFFF;
+        int sourceTo = Math.min(upTo - disi.block, BLOCK_SIZE);
+        int destFrom = disi.doc - offset;
+
+        long fp = disi.slice.getFilePointer();
+        disi.slice.seek(fp - Long.BYTES); // seek back a long to include current word (disi.word).
+        int numWords = FixedBitSet.bits2words(sourceTo) - disi.wordIndex;
+        disi.slice.readLongs(disi.bitSet.getBits(), disi.wordIndex, numWords);
+        FixedBitSet.orRange(disi.bitSet, sourceFrom, bitSet, destFrom, sourceTo - sourceFrom);
+
+        int blockEnd = disi.block | 0xFFFF;
+        if (upTo > blockEnd) {
+          disi.slice.seek(disi.blockEnd);
+          disi.index += disi.bitSet.cardinality(sourceFrom, sourceTo);
+          return false;
+        } else {
+          disi.slice.seek(fp);
+          return advanceWithinBlock(disi, upTo);
+        }
+      }
     },
     ALL {
       @Override
@@ -707,6 +765,19 @@ boolean advanceExactWithinBlock(IndexedDISI disi, int target) {
         disi.index = target - disi.gap;
         return true;
       }
+
+      @Override
+      boolean intoBitSetWithinBlock(IndexedDISI disi, int upTo, FixedBitSet bitSet, int offset) {
+        final int blockEnd = disi.block | 0xFFFF;
+        if (upTo <= blockEnd) {
+          bitSet.set(disi.doc - offset, upTo - offset);
+          advanceWithinBlock(disi, upTo);
+          return true;
+        } else {
+          bitSet.set(disi.doc - offset, blockEnd - offset + 1);
+          return false;
+        }
+      }
     };
 
     /**
@@ -720,6 +791,19 @@ boolean advanceExactWithinBlock(IndexedDISI disi, int target) {
      * return whether this document exists.
      */
     abstract boolean advanceExactWithinBlock(IndexedDISI disi, int target) throws IOException;
+
+    /**
+     * Similar to {@link DocIdSetIterator#intoBitSet}, load docs in this block into a bitset. This
+     * method returns true if there are remaining docs (gte upTo) in the block, otherwise false.
+     * When false return, fp of {@link IndexedDISI#slice} is at {@link IndexedDISI#blockEnd} and
+     * {@link IndexedDISI#index} is correct but other status vars are undefined. Caller should
+     * decode the header of next block by {@link #readBlockHeader()}.
+     *
+     * <p>Caller need to make sure {@link IndexedDISI#doc} greater than or equals to {@link
+     * IndexedDISI#block} and less than {@code upTo} when calling this.
+     */
+    abstract boolean intoBitSetWithinBlock(
+        IndexedDISI disi, int upTo, FixedBitSet bitSet, int offset) throws IOException;
   }
 
   /**

lucene/core/src/java/org/apache/lucene/codecs/lucene90/Lucene90DocValuesProducer.java

@@ -50,6 +50,7 @@
 import org.apache.lucene.store.RandomAccessInput;
 import org.apache.lucene.store.ReadAdvice;
 import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.FixedBitSet;
 import org.apache.lucene.util.IOUtils;
 import org.apache.lucene.util.LongValues;
 import org.apache.lucene.util.compress.LZ4;
@@ -500,6 +501,11 @@ public int docID() {
       return disi.docID();
     }
 
+    @Override
+    public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+      disi.intoBitSet(upTo, bitSet, offset);
+    }
+
     @Override
     public long cost() {
       return disi.cost();
@@ -781,6 +787,11 @@ public int advance(int target) throws IOException {
     public boolean advanceExact(int target) throws IOException {
       return disi.advanceExact(target);
     }
+
+    @Override
+    public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+      disi.intoBitSet(upTo, bitSet, offset);
+    }
   }
 
   @Override
@@ -987,6 +998,11 @@ public int advance(int target) throws IOException {
             return disi.advance(target);
           }
 
+          @Override
+          public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+            disi.intoBitSet(upTo, bitSet, offset);
+          }
+
           @Override
           public long cost() {
             return disi.cost();
@@ -1491,6 +1507,12 @@ public int docValueCount() {
           return count;
         }
 
+        @Override
+        public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+          set = false;
+          disi.intoBitSet(upTo, bitSet, offset);
+        }
+
         private void set() {
           if (set == false) {
             final int index = disi.index();
@@ -1641,6 +1663,12 @@ public int advance(int target) throws IOException {
             return disi.advance(target);
           }
 
+          @Override
+          public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+            set = false;
+            disi.intoBitSet(upTo, bitSet, offset);
+          }
+
           @Override
           public long cost() {
             return disi.cost();

lucene/core/src/java/org/apache/lucene/index/DocValuesIterator.java

@@ -26,6 +26,9 @@ abstract class DocValuesIterator extends DocIdSetIterator {
    * {@code target} must be greater than or equal to the current {@link #docID() doc ID} and must be
    * a valid doc ID, ie. ≥ 0 and < {@code maxDoc}. After this method returns, {@link #docID()}
    * returns {@code target}.
+   *
+   * <p>Note: it is illegal to call {@link DocIdSetIterator#intoBitSet} or {@link
+   * DocIdSetIterator#docIDRunEnd()} when this method returns false.
    */
   public abstract boolean advanceExact(int target) throws IOException;
 }

lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestIndexedDISI.java

@@ -526,6 +526,16 @@ private void doTest(BitSet set, Directory dir) throws IOException {
       }
     }
 
+    for (int step : new int[] {100, 1000, 10000, 100000}) {
+      try (IndexInput in = dir.openInput("foo", IOContext.DEFAULT)) {
+        IndexedDISI disi =
+            new IndexedDISI(in, 0L, length, jumpTableentryCount, denseRankPower, cardinality);
+        BitSetIterator disi2 = new BitSetIterator(set, cardinality);
+        int disi2length = set.length();
+        assertIntoBitsetRandomized(disi, disi2, disi2length, step);
+      }
+    }
+
     dir.deleteFile("foo");
   }
 
@@ -555,6 +565,55 @@ private void assertAdvanceExactRandomized(
     }
   }
 
+  private void assertIntoBitsetRandomized(
+      IndexedDISI disi, BitSetIterator disi2, int disi2length, int step) throws IOException {
+    int index = -1;
+    FixedBitSet set1 = new FixedBitSet(step);
+    FixedBitSet set2 = new FixedBitSet(step);
+
+    for (int upTo = 0; upTo < disi2length; ) {
+      int lastUpTo = upTo;
+      upTo += TestUtil.nextInt(random(), 0, step);
+      int offset = TestUtil.nextInt(random(), lastUpTo, upTo);
+
+      if (disi.docID() < offset) {
+        disi.advance(offset);
+      }
+      int doc = disi2.docID();
+      while (doc < offset) {
+        index++;
+        doc = disi2.nextDoc();
+      }
+      while (doc < upTo) {
+        set2.set(doc - offset);
+        index++;
+        doc = disi2.nextDoc();
+      }
+
+      disi.intoBitSet(upTo, set1, offset);
+      assertEquals(index, disi.index());
+      assertEquals(disi2.docID(), disi.docID());
+
+      BitSetIterator expected = new BitSetIterator(set2, set2.cardinality());
+      BitSetIterator actual = new BitSetIterator(set1, set1.cardinality());
+      for (int expectedDoc = expected.nextDoc();
+          expectedDoc != DocIdSetIterator.NO_MORE_DOCS;
+          expectedDoc = expected.nextDoc()) {
+        int actualDoc = actual.nextDoc();
+        assertEquals(expectedDoc + offset, actualDoc + offset); // plus offset for better message.
+      }
+      assertEquals(DocIdSetIterator.NO_MORE_DOCS, actual.nextDoc());
+
+      if (disi2.docID() != DocIdSetIterator.NO_MORE_DOCS) {
+        assertEquals(disi2.nextDoc(), disi.nextDoc());
+        assertEquals(++index, disi.index());
+      }
+
+      set1.clear();
+      set2.clear();
+    }
+  }
+
   private void assertSingleStepEquality(IndexedDISI disi, BitSetIterator disi2) throws IOException {
     int i = 0;
     for (int doc = disi2.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = disi2.nextDoc()) {