Implement #intoBitSet on IntArrayDocIdSet and RoaringDocIdSet. (#…

…14135) These doc id sets can implement `#intoBitSet` in a way that auto-vectorizes. For reference, `RoaringDocIdSet` is used by the query cache, and `IntArrayDocIdSet` is used by point queries.
apache · Jan 21, 2025 · f737c00 · f737c00
1 parent 1a3db4c
commit f737c00
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 0 deletions.
diff --git a/lucene/core/src/java/org/apache/lucene/search/DocIdSetIterator.java b/lucene/core/src/java/org/apache/lucene/search/DocIdSetIterator.java
@@ -228,6 +228,8 @@ protected final int slowAdvance(int target) throws IOException {
    *
    * <p><b>Note</b>: It is important not to clear bits from {@code bitSet} that may be already set.
    *
+   * <p><b>Note</b>: {@code offset} may be negative.
+   *
    * @lucene.internal
    */
   public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {

diff --git a/lucene/core/src/java/org/apache/lucene/util/IntArrayDocIdSet.java b/lucene/core/src/java/org/apache/lucene/util/IntArrayDocIdSet.java
@@ -95,6 +95,21 @@ public int advance(int target) throws IOException {
       return doc = docs[i++];
     }
 
+    @Override
+    public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+      if (doc >= upTo) {
+        return;
+      }
+
+      int from = i - 1;
+      int to = VectorUtil.findNextGEQ(docs, upTo, from, length);
+      for (int i = from; i < to; ++i) {
+        bitSet.set(docs[i] - offset);
+      }
+      doc = docs[to];
+      i = to + 1;
+    }
+
     @Override
     public long cost() {
       return length;

diff --git a/lucene/core/src/java/org/apache/lucene/util/RoaringDocIdSet.java b/lucene/core/src/java/org/apache/lucene/util/RoaringDocIdSet.java
@@ -217,6 +217,20 @@ public int advance(int target) throws IOException {
             return doc = docId(i);
           }
         }
+
+        @Override
+        public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+          if (doc >= upTo) {
+            return;
+          }
+
+          int from = i;
+          advance(upTo);
+          int to = i;
+          for (int i = from; i < to; ++i) {
+            bitSet.set(docId(i) - offset);
+          }
+        }
       };
     }
   }
@@ -312,6 +326,26 @@ private int firstDocFromNextBlock() throws IOException {
       }
     }
 
+    @Override
+    public void intoBitSet(int upTo, FixedBitSet bitSet, int offset) throws IOException {
+      for (; ; ) {
+        int subUpto = upTo - (block << 16);
+        if (subUpto < 0) {
+          break;
+        }
+        int subOffset = offset - (block << 16);
+        sub.intoBitSet(subUpto, bitSet, subOffset);
+        if (sub.docID() == NO_MORE_DOCS) {
+          if (firstDocFromNextBlock() == NO_MORE_DOCS) {
+            break;
+          }
+        } else {
+          doc = (block << 16) | sub.docID();
+          break;
+        }
+      }
+    }
+
     @Override
     public long cost() {
       return cardinality;