Repository: apache/lucene
Pull request #14273: Speed up histogram collection in a similar way as disjunction counts.
Author: login: jpountz
Labels: name: module:core/search
name: module:test-framework
name: module:sandbox
Created: 2025-02-21T15:41:36Z | Merged: 2025-03-28T14:51:54Z

This generalizes the `IndexSearcher#count` optimization from PR #12415 to histogram facets by introducing specialization for counting the number of matching docs in a range of doc IDs.

Currently, disjunctions and dense conjunctions both internally collect `DocIdStream`s backed by a bitset. In the future, we could make more queries collect whole `DocIdStream`s at once to speed up collection, e.g. `MatchAllDocsQuery` or doc-value range queries that take advantage of a sparse index.


Diff:
diff --git a/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
new file mode 100644
index 000000000000..42a75073c814
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class BitSetDocIdStream extends DocIdStream {
+
+  private final FixedBitSet bitSet;
+  private final int offset, max;
+  private int upTo;
+
+  BitSetDocIdStream(FixedBitSet bitSet, int offset) {
+    this.bitSet = bitSet;
+    this.offset = offset;
+    upTo = offset;
+    max = (int) Math.min(Integer.MAX_VALUE, (long) offset + bitSet.length());
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo > this.upTo) {
+      upTo = Math.min(upTo, max);
+      bitSet.forEach(this.upTo - offset, upTo - offset, offset, consumer);
+      this.upTo = upTo;
+    }
+  }
+
+  @Override
+  public int count(int upTo) throws IOException {
+    if (upTo > this.upTo) {
+      upTo = Math.min(upTo, max);
+      int count = bitSet.cardinality(this.upTo - offset, upTo - offset);
+      this.upTo = upTo;
+      return count;
+    } else {
+      return 0;
+    }
+  }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java b/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
index a6599a57fd25..0be42a2a3e7d 100644
--- a/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
@@ -82,49 +82,6 @@ public DisiWrapper get(int i) {
   final long cost;
   final boolean needsScores;
 
-  final class DocIdStreamView extends DocIdStream {
-
-    int base;
-
-    @Override
-    public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
-      FixedBitSet matching = BooleanScorer.this.matching;
-      Bucket[] buckets = BooleanScorer.this.buckets;
-      int base = this.base;
-      long[] bitArray = matching.getBits();
-      for (int idx = 0; idx < bitArray.length; idx++) {
-        long bits = bitArray[idx];
-        while (bits != 0L) {
-          int ntz = Long.numberOfTrailingZeros(bits);
-          if (buckets != null) {
-            final int indexInWindow = (idx << 6) | ntz;
-            final Bucket bucket = buckets[indexInWindow];
-            if (bucket.freq >= minShouldMatch) {
-              score.score = (float) bucket.score;
-              consumer.accept(base | indexInWindow);
-            }
-            bucket.freq = 0;
-            bucket.score = 0;
-          } else {
-            consumer.accept(base | (idx << 6) | ntz);
-          }
-          bits ^= 1L << ntz;
-        }
-      }
-    }
-
-    @Override
-    public int count() throws IOException {
-      if (minShouldMatch > 1) {
-        // We can't just count bits in that case
-        return super.count();
-      }
-      return matching.cardinality();
-    }
-  }
-
-  private final DocIdStreamView docIdStreamView = new DocIdStreamView();
-
   BooleanScorer(Collection<Scorer> scorers, int minShouldMatch, boolean needsScores) {
     if (minShouldMatch < 1 || minShouldMatch > scorers.size()) {
       throw new IllegalArgumentException(
@@ -202,14 +159,33 @@ private void scoreWindowIntoBitSetAndReplay(
       w.doc = it.docID();
     }
 
-    if (buckets == null && acceptDocs != null) {
-      // In this case, live docs have not been applied yet.
-      acceptDocs.applyMask(matching, base);
+    if (buckets == null) {
+      if (acceptDocs != null) {
+        // In this case, live docs have not been applied yet.
+        acceptDocs.applyMask(matching, base);
+      }
+      collector.collect(new BitSetDocIdStream(matching, base));
+    } else {
+      FixedBitSet matching = BooleanScorer.this.matching;
+      Bucket[] buckets = BooleanScorer.this.buckets;
+      long[] bitArray = matching.getBits();
+      for (int idx = 0; idx < bitArray.length; idx++) {
+        long bits = bitArray[idx];
+        while (bits != 0L) {
+          int ntz = Long.numberOfTrailingZeros(bits);
+          final int indexInWindow = (idx << 6) | ntz;
+          final Bucket bucket = buckets[indexInWindow];
+          if (bucket.freq >= minShouldMatch) {
+            score.score = (float) bucket.score;
+            collector.collect(base | indexInWindow);
+          }
+          bucket.freq = 0;
+          bucket.score = 0;
+          bits ^= 1L << ntz;
+        }
+      }
     }
 
-    docIdStreamView.base = base;
-    collector.collect(docIdStreamView);
-
     matching.clear();
   }
 
diff --git a/lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
new file mode 100644
index 000000000000..0541969d7e34
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class DISIDocIdStream extends DocIdStream {
+
+  private final DocIdSetIterator iterator;
+  private final int max;
+  private final FixedBitSet spare;
+
+  DISIDocIdStream(DocIdSetIterator iterator, int max, FixedBitSet spare) {
+    if (max - iterator.docID() > spare.length()) {
+      throw new IllegalArgumentException("Bit set is too small to hold all potential matches");
+    }
+    this.iterator = iterator;
+    this.max = max;
+    this.spare = spare;
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return iterator.docID() < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    // If there are no live docs to apply, loading matching docs into a bit set and then iterating
+    // bits is unlikely to beat iterating the iterator directly.
+    upTo = Math.min(upTo, max);
+    for (int doc = iterator.docID(); doc < upTo; doc = iterator.nextDoc()) {
+      consumer.accept(doc);
+    }
+  }
+
+  @Override
+  public int count(int upTo) throws IOException {
+    if (iterator.docID() >= upTo) {
+      return 0;
+    }
+    // If the collector is just interested in the count, loading in a bit set and counting bits is
+    // often faster than incrementing a counter on every call to nextDoc().
+    assert spare.scanIsEmpty();
+    upTo = Math.min(upTo, max);
+    int offset = iterator.docID();
+    iterator.intoBitSet(upTo, spare, offset);
+    int count = spare.cardinality(0, upTo - offset);
+    spare.clear(0, upTo - offset);
+    assert spare.scanIsEmpty();
+    return count;
+  }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/search/DenseConjunctionBulkScorer.java b/lucene/core/src/java/org/apache/lucene/search/DenseConjunctionBulkScorer.java
index 3f1bbe86146a..48190d6b5a14 100644
--- a/lucene/core/src/java/org/apache/lucene/search/DenseConjunctionBulkScorer.java
+++ b/lucene/core/src/java/org/apache/lucene/search/DenseConjunctionBulkScorer.java
@@ -68,10 +68,6 @@ int docIDRunEnd() throws IOException {
   private final FixedBitSet clauseWindowMatches = new FixedBitSet(WINDOW_SIZE);
   private final List<DocIdSetIterator> windowApproximations = new ArrayList<>();
   private final List<TwoPhaseIterator> windowTwoPhases = new ArrayList<>();
-  private final DocIdStreamView docIdStreamView = new DocIdStreamView();
-  private final RangeDocIdStream rangeDocIdStream = new RangeDocIdStream();
-  private final SingleIteratorDocIdStream singleIteratorDocIdStream =
-      new SingleIteratorDocIdStream();
 
   static DenseConjunctionBulkScorer of(List<Scorer> filters, int maxDoc, float constantScore) {
     List<DocIdSetIterator> iterators = new ArrayList<>();
@@ -184,9 +180,7 @@ private int scoreWindow(
 
       if (minDocIDRunEnd - min >= WINDOW_SIZE / 2) {
         // We have a large range of doc IDs that all match.
-        rangeDocIdStream.from = min;
-        rangeDocIdStream.to = minDocIDRunEnd;
-        collector.collect(rangeDocIdStream);
+        collector.collect(new RangeDocIdStream(min, minDocIDRunEnd));
         return minDocIDRunEnd;
       }
     }
@@ -206,10 +200,11 @@ private int scoreWindow(
       if (acceptDocs == null && windowApproximations.size() == 1) {
         // We have a range of doc IDs where all matches of an iterator are matches of the
         // conjunction.
-        singleIteratorDocIdStream.iterator = windowApproximations.get(0);
-        singleIteratorDocIdStream.from = min;
-        singleIteratorDocIdStream.to = bitsetWindowMax;
-        collector.collect(singleIteratorDocIdStream);
+        DocIdSetIterator iterator = windowApproximations.get(0);
+        if (iterator.docID() < min) {
+          iterator.advance(min);
+        }
+        collector.collect(new DISIDocIdStream(iterator, bitsetWindowMax, clauseWindowMatches));
       } else {
         scoreWindowUsingBitSet(collector, acceptDocs, windowApproximations, min, bitsetWindowMax);
       }
@@ -289,8 +284,7 @@ private void scoreWindowUsingBitSet(
         windowMatch = advance(windowMatches, windowMatch + 1);
       }
     } else {
-      docIdStreamView.windowBase = windowBase;
-      collector.collect(docIdStreamView);
+      collector.collect(new BitSetDocIdStream(windowMatches, windowBase));
     }
 
     windowMatches.clear();
@@ -370,78 +364,4 @@ private static void scoreWindowUsingLeapFrog(
   public long cost() {
     return iterators.get(0).approximation().cost();
   }
-
-  final class DocIdStreamView extends DocIdStream {
-
-    int windowBase;
-
-    @Override
-    public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
-      int windowBase = this.windowBase;
-      long[] bitArray = windowMatches.getBits();
-      for (int idx = 0; idx < bitArray.length; idx++) {
-        long bits = bitArray[idx];
-        while (bits != 0L) {
-          int ntz = Long.numberOfTrailingZeros(bits);
-          consumer.accept(windowBase + ((idx << 6) | ntz));
-          bits ^= 1L << ntz;
-        }
-      }
-    }
-
-    @Override
-    public int count() throws IOException {
-      return windowMatches.cardinality();
-    }
-  }
-
-  final class RangeDocIdStream extends DocIdStream {
-
-    int from, to;
-
-    @Override
-    public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
-      for (int i = from; i < to; ++i) {
-        consumer.accept(i);
-      }
-    }
-
-    @Override
-    public int count() throws IOException {
-      return to - from;
-    }
-  }
-
-  /** {@link DocIdStream} for a {@link DocIdSetIterator} with no live docs to apply. */
-  final class SingleIteratorDocIdStream extends DocIdStream {
-
-    int from, to;
-    DocIdSetIterator iterator;
-
-    @Override
-    public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
-      // If there are no live docs to apply, loading matching docs into a bit set and then iterating
-      // bits is unlikely to beat iterating the iterator directly.
-      if (iterator.docID() < from) {
-        iterator.advance(from);
-      }
-      for (int doc = iterator.docID(); doc < to; doc = iterator.nextDoc()) {
-        consumer.accept(doc);
-      }
-    }
-
-    @Override
-    public int count() throws IOException {
-      // If the collector is just interested in the count, loading in a bit set and counting bits is
-      // often faster than incrementing a counter on every call to nextDoc().
-      assert windowMatches.scanIsEmpty();
-      if (iterator.docID() < from) {
-        iterator.advance(from);
-      }
-      iterator.intoBitSet(to, clauseWindowMatches, from);
-      int count = clauseWindowMatches.cardinality();
-      clauseWindowMatches.clear();
-      return count;
-    }
-  }
 }
diff --git a/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
index afe73f3a56f2..b13e0ce6ec14 100644
--- a/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
+++ b/lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
@@ -19,8 +19,7 @@
 import java.io.IOException;
 
 /**
- * A stream of doc IDs. Most methods on {@link DocIdStream}s are terminal, meaning that the {@link
- * DocIdStream} may not be further used.
+ * A stream of doc IDs. Doc IDs may be consumed at most once.
  *
  * @see LeafCollector#collect(DocIdStream)
  * @lucene.experimental
@@ -34,12 +33,34 @@ protected DocIdStream() {}
    * Iterate over doc IDs contained in this stream in order, calling the given {@link
    * CheckedIntConsumer} on them. This is a terminal operation.
    */
-  public abstract void forEach(CheckedIntConsumer<IOException> consumer) throws IOException;
+  public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
+    forEach(DocIdSetIterator.NO_MORE_DOCS, consumer);
+  }
+
+  /**
+   * Iterate over doc IDs contained in this doc ID stream up to the given {@code upTo} exclusive,
+   * calling the given {@link CheckedIntConsumer} on them. It is not possible to iterate these doc
+   * IDs again later on.
+   */
+  public abstract void forEach(int upTo, CheckedIntConsumer<IOException> consumer)
+      throws IOException;
 
   /** Count the number of entries in this stream. This is a terminal operation. */
   public int count() throws IOException {
-    int[] count = new int[1];
-    forEach(_ -> count[0]++);
-    return count[0];
+    return count(DocIdSetIterator.NO_MORE_DOCS);
   }
+
+  /**
+   * Count the number of doc IDs in this stream that are below the given {@code upTo}. These doc IDs
+   * may not be consumed again later.
+   */
+  // Note: it's abstract rather than having a default impl that delegates to #forEach because doing
+  // so would defeat the purpose of collecting hits via a DocIdStream.
+  public abstract int count(int upTo) throws IOException;
+
+  /**
+   * Return {@code true} if this stream may have remaining doc IDs. This must eventually return
+   * {@code false} when the stream is exhausted.
+   */
+  public abstract boolean mayHaveRemaining();
 }
diff --git a/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java b/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
new file mode 100644
index 000000000000..0acba2955595
--- /dev/null
+++ b/lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+
+final class RangeDocIdStream extends DocIdStream {
+
+  private int upTo;
+  private final int max;
+
+  RangeDocIdStream(int min, int max) {
+    if (max < min) {
+      throw new IllegalArgumentException("min = " + min + " > max = " + max);
+    }
+    this.upTo = min;
+    this.max = max;
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo > this.upTo) {
+      upTo = Math.min(upTo, max);
+      for (int doc = this.upTo; doc < upTo; ++doc) {
+        consumer.accept(doc);
+      }
+      this.upTo = upTo;
+    }
+  }
+
+  @Override
+  public int count(int upTo) throws IOException {
+    if (upTo > this.upTo) {
+      upTo = Math.min(upTo, max);
+      int count = upTo - this.upTo;
+      this.upTo = upTo;
+      return count;
+    } else {
+      return 0;
+    }
+  }
+}
diff --git a/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java b/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
index 8ca4b6ae2736..22ef8032d2cc 100644
--- a/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
+++ b/lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
@@ -19,6 +19,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
+import org.apache.lucene.search.CheckedIntConsumer;
 import org.apache.lucene.search.DocIdSet;
 import org.apache.lucene.search.DocIdSetIterator;
 
@@ -204,6 +205,40 @@ public int cardinality() {
     return Math.toIntExact(tot);
   }
 
+  /**
+   * Return the number of set bits between indexes {@code from} inclusive and {@code to} exclusive.
+   */
+  public int cardinality(int from, int to) {
+    Objects.checkFromToIndex(from, to, length());
+
+    int cardinality = 0;
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        bits &= (1L << (to - from)) - 1L;
+        return Long.bitCount(bits);
+      }
+      cardinality += Long.bitCount(bits);
+      from += numBitsTilNextWord;
+      assert (from & 0x3F) == 0;
+    }
+
+    for (int i = from >> 6, end = to >> 6; i < end; ++i) {
+      cardinality += Long.bitCount(bits[i]);
+    }
+
+    // Now handle bits between the last complete word and to
+    if ((to & 0x3F) != 0) {
+      long bits = this.bits[to >> 6] << -to;
+      cardinality += Long.bitCount(bits);
+    }
+
+    return cardinality;
+  }
+
   @Override
   public int approximateCardinality() {
     // Naive sampling: compute the number of bits that are set on the first 16 longs every 1024
@@ -790,4 +825,48 @@ public void applyMask(FixedBitSet bitSet, int offset) {
       throw new IllegalArgumentException("Some bits are set beyond the end of live docs");
     }
   }
+
+  /**
+   * For each set bit from {@code from} inclusive to {@code to} exclusive, add {@code base} to the
+   * bit index and call {@code consumer} on it. This is internally used by queries that use bit sets
+   * as intermediate representations of their matches.
+   */
+  public void forEach(int from, int to, int base, CheckedIntConsumer<IOException> consumer)
+      throws IOException {
+    Objects.checkFromToIndex(from, to, length());
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        // All bits are in a single word
+        bits &= (1L << (to - from)) - 1L;
+        forEach(bits, from + base, consumer);
+        return;
+      }
+      forEach(bits, from + base, consumer);
+      from += numBitsTilNextWord;
+      assert (from & 0x3F) == 0;
+    }
+
+    for (int i = from >> 6, end = to >> 6; i < end; ++i) {
+      forEach(bits[i], base + (i << 6), consumer);
+    }
+
+    // Now handle remaining bits in the last partial word
+    if ((to & 0x3F) != 0) {
+      long bits = this.bits[to >> 6] & ((1L << to) - 1);
+      forEach(bits, base + (to & ~0x3F), consumer);
+    }
+  }
+
+  private static void forEach(long bits, int base, CheckedIntConsumer<IOException> consumer)
+      throws IOException {
+    while (bits != 0L) {
+      int ntz = Long.numberOfTrailingZeros(bits);
+      consumer.accept(base + ntz);
+      bits ^= 1L << ntz;
+    }
+  }
 }
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java b/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java
new file mode 100644
index 000000000000..51654cbef4c5
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/search/TestBitSetDocIdStream.java
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+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.FixedBitSet;
+
+public class TestBitSetDocIdStream extends LuceneTestCase {
+
+  private static FixedBitSet randomBitSet() {
+    FixedBitSet set = new FixedBitSet(100);
+    for (int i = TestUtil.nextInt(random(), 0, 5);
+        i < set.length();
+        i += TestUtil.nextInt(random(), 1, 5)) {
+      set.set(i);
+    }
+    return set;
+  }
+
+  public void testForEach() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    int base = 42;
+    BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+
+    stream.forEach(
+        doc -> {
+          assertEquals(base + iterator.nextDoc(), doc);
+        });
+    assertEquals(DocIdSetIterator.NO_MORE_DOCS, iterator.nextDoc());
+  }
+
+  public void testCount() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    int base = 42;
+    BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+    assertEquals(bitSet.cardinality(), stream.count());
+  }
+
+  public void testForEachUpTo() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    int base = 42;
+    BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(20, _ -> fail());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        100,
+        doc -> {
+          assertEquals(base + iterator.nextDoc(), doc);
+        });
+    assertEquals(bitSet.prevSetBit(99 - base), iterator.docID());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        200,
+        doc -> {
+          assertEquals(base + iterator.nextDoc(), doc);
+        });
+    assertEquals(bitSet.prevSetBit(bitSet.length() - 1), iterator.docID());
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+
+  public void testCountUpTo() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    int base = 42;
+    BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(0, stream.count(20));
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(0, 100 - 42), stream.count(100));
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(100 - 42, bitSet.length()), stream.count(200));
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+
+  public void testMixForEachCountUpTo() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    int base = 42;
+    BitSetDocIdStream stream = new BitSetDocIdStream(bitSet, base);
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(20, _ -> fail());
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(0, stream.count(30));
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        100,
+        doc -> {
+          assertEquals(base + iterator.nextDoc(), doc);
+        });
+    assertEquals(bitSet.prevSetBit(99 - base), iterator.docID());
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(100 - 42, 120 - 42), stream.count(120));
+
+    assertTrue(stream.mayHaveRemaining());
+    iterator.advance(bitSet.prevSetBit(120 - 1 - 42));
+    stream.forEach(
+        200,
+        doc -> {
+          assertEquals(base + iterator.nextDoc(), doc);
+        });
+    assertEquals(bitSet.prevSetBit(bitSet.length() - 1), iterator.docID());
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestDISIDocIdStream.java b/lucene/core/src/test/org/apache/lucene/search/TestDISIDocIdStream.java
new file mode 100644
index 000000000000..9a400c0237cf
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/search/TestDISIDocIdStream.java
@@ -0,0 +1,157 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+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.FixedBitSet;
+
+public class TestDISIDocIdStream extends LuceneTestCase {
+
+  private static FixedBitSet randomBitSet() {
+    FixedBitSet set = new FixedBitSet(100);
+    for (int i = TestUtil.nextInt(random(), 0, 5);
+        i < set.length();
+        i += TestUtil.nextInt(random(), 1, 5)) {
+      set.set(i);
+    }
+    return set;
+  }
+
+  public void testForEach() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    iterator.advance(42);
+    DocIdStream stream = new DISIDocIdStream(iterator, 80, new FixedBitSet(40));
+
+    BitSetIterator expected = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    expected.advance(42);
+
+    stream.forEach(
+        doc -> {
+          assertEquals(expected.docID(), doc);
+          assertTrue(doc < 80);
+          expected.nextDoc();
+        });
+    assertTrue(expected.docID() >= 80);
+  }
+
+  public void testCount() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    iterator.advance(42);
+    DocIdStream stream = new DISIDocIdStream(iterator, 80, new FixedBitSet(40));
+    assertEquals(bitSet.cardinality(42, 80), stream.count());
+  }
+
+  public void testForEachUpTo() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    iterator.advance(42);
+    DocIdStream stream = new DISIDocIdStream(iterator, 80, new FixedBitSet(40));
+
+    BitSetIterator expected = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    expected.advance(42);
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(40, _ -> fail());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        60,
+        doc -> {
+          assertEquals(expected.docID(), doc);
+          assertTrue(doc < 60);
+          expected.nextDoc();
+        });
+    assertTrue(expected.docID() >= 60);
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        120,
+        doc -> {
+          assertEquals(expected.docID(), doc);
+          assertTrue(doc < 80);
+          expected.nextDoc();
+        });
+    assertTrue(expected.docID() >= 80);
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+
+  public void testCountUpTo() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    iterator.advance(42);
+    DocIdStream stream = new DISIDocIdStream(iterator, 80, new FixedBitSet(40));
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(0, stream.count(40));
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(42, 60), stream.count(60));
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(60, 80), stream.count(120));
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+
+  public void testMixForEachCountUpTo() throws IOException {
+    FixedBitSet bitSet = randomBitSet();
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    iterator.advance(42);
+    DocIdStream stream = new DISIDocIdStream(iterator, 80, new FixedBitSet(40));
+
+    BitSetIterator expected = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(40, _ -> fail());
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(42, 50), stream.count(50));
+
+    assertTrue(stream.mayHaveRemaining());
+    expected.advance(50);
+    stream.forEach(
+        60,
+        doc -> {
+          assertEquals(expected.docID(), doc);
+          assertTrue(doc < 60);
+          expected.nextDoc();
+        });
+    assertTrue(expected.docID() >= 60);
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(bitSet.cardinality(60, 70), stream.count(70));
+
+    assertTrue(stream.mayHaveRemaining());
+    expected.advance(70);
+    stream.forEach(
+        120,
+        doc -> {
+          assertEquals(expected.docID(), doc);
+          assertTrue(doc < 80);
+          expected.nextDoc();
+        });
+    assertTrue(expected.docID() >= 80);
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java b/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java
new file mode 100644
index 000000000000..0018f75fb25b
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/search/TestRangeDocIdStream.java
@@ -0,0 +1,104 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.tests.util.LuceneTestCase;
+
+public class TestRangeDocIdStream extends LuceneTestCase {
+
+  public void testForEach() throws IOException {
+    RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+    int[] expected = new int[] {42};
+    stream.forEach(
+        doc -> {
+          assertEquals(expected[0]++, doc);
+        });
+    assertEquals(100, expected[0]);
+  }
+
+  public void testCount() throws IOException {
+    RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+    assertEquals(100 - 42, stream.count());
+  }
+
+  public void testForEachUpTo() throws IOException {
+    RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+    int[] expected = new int[] {42};
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(20, _ -> fail());
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        65,
+        doc -> {
+          assertEquals(expected[0]++, doc);
+        });
+    assertEquals(65, expected[0]);
+
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        120,
+        doc -> {
+          assertEquals(expected[0]++, doc);
+        });
+    assertEquals(100, expected[0]);
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+
+  public void testCountUpTo() throws IOException {
+    RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(0, stream.count(20));
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(65 - 42, stream.count(65));
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(100 - 65, stream.count(120));
+    assertFalse(stream.mayHaveRemaining());
+  }
+
+  public void testMixForEachCountUpTo() throws IOException {
+    RangeDocIdStream stream = new RangeDocIdStream(42, 100);
+    int[] expected = new int[] {42};
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        65,
+        doc -> {
+          assertEquals(expected[0]++, doc);
+        });
+    assertEquals(65, expected[0]);
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(80 - 65, stream.count(80));
+
+    expected[0] = 80;
+    assertTrue(stream.mayHaveRemaining());
+    stream.forEach(
+        90,
+        doc -> {
+          assertEquals(expected[0]++, doc);
+        });
+    assertEquals(90, expected[0]);
+
+    assertTrue(stream.mayHaveRemaining());
+    assertEquals(100 - 90, stream.count(120));
+
+    assertFalse(stream.mayHaveRemaining());
+  }
+}
diff --git a/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java b/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java
index 39acd5ea209e..efdebee71585 100644
--- a/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java
+++ b/lucene/core/src/test/org/apache/lucene/util/TestFixedBitSet.java
@@ -17,7 +17,9 @@
 package org.apache.lucene.util;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Random;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.tests.util.BaseBitSetTestCase;
@@ -714,4 +716,63 @@ public void testAndRange() {
       }
     }
   }
+
+  private static int slowCardinality(FixedBitSet bitSet, int from, int to) {
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.cardinality());
+    int cardinality = 0;
+    for (int i = iterator.advance(from); i < to; i = iterator.nextDoc()) {
+      cardinality += 1;
+    }
+    return cardinality;
+  }
+
+  public void testRangeCardinality() {
+    FixedBitSet bitSet = new FixedBitSet(0);
+    assertEquals(0, bitSet.cardinality(0, 0));
+
+    bitSet = new FixedBitSet(TestUtil.nextInt(random(), 64, 4096));
+    for (int i = 0; i < bitSet.length(); i += TestUtil.nextInt(random(), 1, 16)) {
+      bitSet.set(i);
+    }
+    assertEquals(0, bitSet.cardinality(0, 0));
+    assertEquals(0, bitSet.cardinality(bitSet.length(), bitSet.length()));
+    assertEquals(slowCardinality(bitSet, 0, 1), bitSet.cardinality(0, 1));
+    assertEquals(
+        slowCardinality(bitSet, bitSet.length() - 1, bitSet.length()),
+        bitSet.cardinality(bitSet.length() - 1, bitSet.length()));
+    assertEquals(slowCardinality(bitSet, 0, 63), bitSet.cardinality(0, 63));
+    assertEquals(
+        slowCardinality(bitSet, bitSet.length() - 63, bitSet.length()),
+        bitSet.cardinality(bitSet.length() - 63, bitSet.length()));
+    assertEquals(
+        slowCardinality(bitSet, 1, bitSet.length() - 1),
+        bitSet.cardinality(1, bitSet.length() - 1));
+  }
+
+  public void testForEach() throws IOException {
+    FixedBitSet bitSet = new FixedBitSet(TestUtil.nextInt(random(), 64, 4096));
+    for (int i = 0; i < bitSet.length(); i += TestUtil.nextInt(random(), 1, 16)) {
+      bitSet.set(i);
+    }
+    for (int iter = 0; iter < 100; ++iter) {
+      int from = TestUtil.nextInt(random(), 0, bitSet.length());
+      int to = TestUtil.nextInt(random(), from, bitSet.length());
+      doTestForEach(bitSet, from, to);
+    }
+  }
+
+  private void doTestForEach(FixedBitSet bitSet, int from, int to) throws IOException {
+    int base = TestUtil.nextInt(random(), 0, 100);
+
+    List<Integer> expected = new ArrayList<>();
+    BitSetIterator iterator = new BitSetIterator(bitSet, bitSet.approximateCardinality());
+    for (int doc = iterator.advance(from); doc < to; doc = iterator.nextDoc()) {
+      expected.add(base + doc);
+    }
+
+    List<Integer> actual = new ArrayList<>();
+    bitSet.forEach(from, to, base, actual::add);
+
+    assertEquals(expected, actual);
+  }
 }
diff --git a/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java b/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java
index 722e03ab08f7..fcb95699a6bb 100644
--- a/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java
+++ b/lucene/sandbox/src/java/org/apache/lucene/sandbox/facet/plain/histograms/HistogramCollector.java
@@ -27,6 +27,7 @@
 import org.apache.lucene.internal.hppc.LongIntHashMap;
 import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.DocIdStream;
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Scorable;
 import org.apache.lucene.search.ScoreMode;
@@ -252,6 +253,28 @@ public void collect(int doc) throws IOException {
       }
     }
 
+    @Override
+    public void collect(DocIdStream stream) throws IOException {
+      for (; ; ) {
+        int upToExclusive = upToInclusive + 1;
+        if (upToExclusive < 0) { // overflow
+          upToExclusive = Integer.MAX_VALUE;
+        }
+
+        if (upToSameBucket) {
+          counts[upToBucketIndex] += stream.count(upToExclusive);
+        } else {
+          stream.forEach(upToExclusive, this::collect);
+        }
+
+        if (stream.mayHaveRemaining()) {
+          advanceSkipper(upToExclusive);
+        } else {
+          break;
+        }
+      }
+    }
+
     @Override
     public void finish() throws IOException {
       // Put counts that we computed in the int[] back into the hash map.
diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
index c0a6eb9a34db..ca05ef76f878 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
@@ -120,7 +120,7 @@ public void finish() throws IOException {
   private class AssertingDocIdStream extends DocIdStream {
 
     private final DocIdStream stream;
-    private boolean consumed;
+    private int lastUpTo = -1;
 
     AssertingDocIdStream(DocIdStream stream) {
       this.stream = stream;
@@ -128,7 +128,7 @@ private class AssertingDocIdStream extends DocIdStream {
 
     @Override
     public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
-      assert consumed == false : "A terminal operation has already been called";
+      assert lastUpTo != DocIdSetIterator.NO_MORE_DOCS : "exhausted";
       stream.forEach(
           doc -> {
             assert doc > lastCollected : "Out of order : " + lastCollected + " " + doc;
@@ -137,15 +137,54 @@ public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException
             consumer.accept(doc);
             lastCollected = doc;
           });
-      consumed = true;
+      lastUpTo = DocIdSetIterator.NO_MORE_DOCS;
+      assert stream.mayHaveRemaining() == false;
+    }
+
+    @Override
+    public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+      assert lastUpTo < upTo : "upTo=" + upTo + " but previous upTo=" + lastUpTo;
+      stream.forEach(
+          doc -> {
+            assert doc > lastCollected : "Out of order : " + lastCollected + " " + doc;
+            assert doc >= min : "Out of range: " + doc + " < " + min;
+            assert doc < max : "Out of range: " + doc + " >= " + max;
+            consumer.accept(doc);
+            lastCollected = doc;
+          });
+      lastUpTo = upTo;
+      if (upTo == DocIdSetIterator.NO_MORE_DOCS) {
+        assert stream.mayHaveRemaining() == false;
+      }
     }
 
     @Override
     public int count() throws IOException {
-      assert consumed == false : "A terminal operation has already been called";
+      assert lastUpTo != DocIdSetIterator.NO_MORE_DOCS : "exhausted";
       int count = stream.count();
-      consumed = true;
+      lastUpTo = DocIdSetIterator.NO_MORE_DOCS;
+      assert stream.mayHaveRemaining() == false;
+      return count;
+    }
+
+    @Override
+    public int count(int upTo) throws IOException {
+      assert lastUpTo < upTo : "upTo=" + upTo + " but previous upTo=" + lastUpTo;
+      int count = stream.count(upTo);
+      lastUpTo = upTo;
+      if (upTo == DocIdSetIterator.NO_MORE_DOCS) {
+        assert stream.mayHaveRemaining() == false;
+      }
       return count;
     }
+
+    @Override
+    public boolean mayHaveRemaining() {
+      boolean mayHaveRemaining = stream.mayHaveRemaining();
+      if (lastUpTo == DocIdSetIterator.NO_MORE_DOCS) {
+        assert mayHaveRemaining == false;
+      }
+      return mayHaveRemaining;
+    }
   }
 }


Reviews:
user: login: epotyom
body: I like the idea! Looks like we can do similar trick for range facets and long values facets?
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-02-24T14:07:20Z
user: login: gsmiller
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T14:19:19Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-26T16:24:16Z
user: login: gsmiller
body: Left some minor feedback. Looks good though!
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:12:35Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:17:07Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:18:53Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:19:27Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:24:13Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T20:27:01Z
user: login: gsmiller
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-27T23:35:09Z
user: login: gsmiller
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T00:52:38Z
user: login: jpountz
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T13:25:46Z
user: login: gsmiller
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T14:17:19Z
user: login: gsmiller
body: Looks good. Thanks for incorporating the minor feedback!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-28T14:20:10Z

Review comments:
diff_hunk: @@ -119,15 +119,54 @@ public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException
             consumer.accept(doc);
             lastCollected = doc;
           });
-      consumed = true;
+      fullyConsumed = true;
+      assert stream.mayHaveRemaining() == false;
+    }
+
+    @Override
+    public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+      assert fullyConsumed == false : "A terminal operation has already been called";
+      stream.forEach(
+          doc -> {
+            assert doc > lastCollected : "Out of order : " + lastCollected + " " + doc;
+            assert doc >= min : "Out of range: " + doc + " < " + min;
+            assert doc < max : "Out of range: " + doc + " >= " + max;
+            consumer.accept(doc);
+            lastCollected = doc;
+          });
+      fullyConsumed = upTo == DocIdSetIterator.NO_MORE_DOCS;
+      if (fullyConsumed) {
+        assert stream.mayHaveRemaining() == false;
+      }
     }
 
     @Override
     public int count() throws IOException {
-      assert consumed == false : "A terminal operation has already been called";
+      assert fullyConsumed == false : "A terminal operation has already been called";
       int count = stream.count();
-      consumed = true;
+      fullyConsumed = true;
+      assert stream.mayHaveRemaining() == false;
       return count;
     }
+
+    @Override
+    public int count(int upTo) throws IOException {
+      assert fullyConsumed == false : "A terminal operation has already been called";
+      int count = stream.count();
path: lucene/test-framework/src/java/org/apache/lucene/tests/search/AssertingLeafCollector.java
user: login: epotyom
body: Should it be `count(upTo)`?
created_at: 2025-02-23T11:33:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,12 +33,35 @@ protected DocIdStream() {}
    * Iterate over doc IDs contained in this stream in order, calling the given {@link
    * CheckedIntConsumer} on them. This is a terminal operation.
    */
-  public abstract void forEach(CheckedIntConsumer<IOException> consumer) throws IOException;
+  public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
+    forEach(DocIdSetIterator.NO_MORE_DOCS, consumer);
+  }
+
+  /**
+   * Iterate over doc IDs contained in this doc ID stream up to the given {@code upTo} exclusive,
+   * calling the given {@link CheckedIntConsumer} on them. It is not possible to iterate these doc
+   * IDs again later on.
+   */
+  public abstract void forEach(int upTo, CheckedIntConsumer<IOException> consumer)
+      throws IOException;
 
   /** Count the number of entries in this stream. This is a terminal operation. */
   public int count() throws IOException {
     int[] count = new int[1];
     forEach(doc -> count[0]++);
     return count[0];
   }
+
+  /**
+   * Count the number of doc IDs in this stream that are below the given {@code upTo}. These doc IDs
+   * may not be consumed again later.
+   */
+  public int count(int upTo) throws IOException {
+    int[] count = new int[1];
+    forEach(upTo, doc -> count[0]++);
+    return count[0];
+  }
+
+  /** Return {@code true} if this stream may have remaining doc IDs. */
path: lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
user: login: epotyom
body: Maybe I'm nitpicking, but is it worth mentioning that it must eventually returning `false`, otherwise `return true` may sound like a correct implementation?
created_at: 2025-02-24T13:59:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,12 +33,35 @@ protected DocIdStream() {}
    * Iterate over doc IDs contained in this stream in order, calling the given {@link
    * CheckedIntConsumer} on them. This is a terminal operation.
    */
-  public abstract void forEach(CheckedIntConsumer<IOException> consumer) throws IOException;
+  public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
+    forEach(DocIdSetIterator.NO_MORE_DOCS, consumer);
+  }
+
+  /**
+   * Iterate over doc IDs contained in this doc ID stream up to the given {@code upTo} exclusive,
+   * calling the given {@link CheckedIntConsumer} on them. It is not possible to iterate these doc
+   * IDs again later on.
+   */
+  public abstract void forEach(int upTo, CheckedIntConsumer<IOException> consumer)
+      throws IOException;
 
   /** Count the number of entries in this stream. This is a terminal operation. */
   public int count() throws IOException {
     int[] count = new int[1];
     forEach(doc -> count[0]++);
     return count[0];
   }
+
+  /**
+   * Count the number of doc IDs in this stream that are below the given {@code upTo}. These doc IDs
+   * may not be consumed again later.
+   */
+  public int count(int upTo) throws IOException {
path: lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
user: login: gsmiller
body: This only becomes an optimization if we specialize this method right? The specializations I'm aware of rely on `FixedBitSet#cardinality`. Are you thinking of peeking into these bit sets to provide cardinality up to the specific doc? (Or maybe I'm missing something?)
created_at: 2025-03-26T14:19:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,12 +33,35 @@ protected DocIdStream() {}
    * Iterate over doc IDs contained in this stream in order, calling the given {@link
    * CheckedIntConsumer} on them. This is a terminal operation.
    */
-  public abstract void forEach(CheckedIntConsumer<IOException> consumer) throws IOException;
+  public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
+    forEach(DocIdSetIterator.NO_MORE_DOCS, consumer);
+  }
+
+  /**
+   * Iterate over doc IDs contained in this doc ID stream up to the given {@code upTo} exclusive,
+   * calling the given {@link CheckedIntConsumer} on them. It is not possible to iterate these doc
+   * IDs again later on.
+   */
+  public abstract void forEach(int upTo, CheckedIntConsumer<IOException> consumer)
+      throws IOException;
 
   /** Count the number of entries in this stream. This is a terminal operation. */
   public int count() throws IOException {
     int[] count = new int[1];
     forEach(doc -> count[0]++);
     return count[0];
   }
+
+  /**
+   * Count the number of doc IDs in this stream that are below the given {@code upTo}. These doc IDs
+   * may not be consumed again later.
+   */
+  public int count(int upTo) throws IOException {
path: lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
user: login: jpountz
body: > Are you thinking of peeking into these bit sets to provide cardinality up to the specific doc? (Or maybe I'm missing something?)

Yes exactly. I have something locally already, I need to beef up testing a bit.

The bitset-based `DocIdStream` is one interesting implementation, the other interesting implementation is the one that is backed by a range of doc IDs that all match. It is internally used by queries that fully match a segment (e.g. `PointRangeQuery` when all the segment's values are contained in the query range, or `MatchAllDocsQuery`) or queries on fields that are part of (or correlate with) the index sort fields. See #14312 for reference.
created_at: 2025-03-26T16:24:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,12 +33,34 @@ protected DocIdStream() {}
    * Iterate over doc IDs contained in this stream in order, calling the given {@link
    * CheckedIntConsumer} on them. This is a terminal operation.
    */
-  public abstract void forEach(CheckedIntConsumer<IOException> consumer) throws IOException;
+  public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
+    forEach(DocIdSetIterator.NO_MORE_DOCS, consumer);
+  }
+
+  /**
+   * Iterate over doc IDs contained in this doc ID stream up to the given {@code upTo} exclusive,
+   * calling the given {@link CheckedIntConsumer} on them. It is not possible to iterate these doc
+   * IDs again later on.
+   */
+  public abstract void forEach(int upTo, CheckedIntConsumer<IOException> consumer)
+      throws IOException;
 
   /** Count the number of entries in this stream. This is a terminal operation. */
   public int count() throws IOException {
-    int[] count = new int[1];
-    forEach(_ -> count[0]++);
-    return count[0];
+    return count(DocIdSetIterator.NO_MORE_DOCS);
   }
+
+  /**
+   * Count the number of doc IDs in this stream that are below the given {@code upTo}. These doc IDs
+   * may not be consumed again later.
+   */
+  // Note: it's abstract rather than having a default impl that delegates to #forEach because doing
path: lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
user: login: gsmiller
body: Thanks for adding this comment! +1 to the rationale as well.
created_at: 2025-03-27T17:17:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class DISIDocIdStream extends DocIdStream {
+
+  private final DocIdSetIterator iterator;
+  private final int to;
path: lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
user: login: gsmiller
body: minor: maybe `max` to be consistent with the other implementations?
created_at: 2025-03-27T17:28:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class DISIDocIdStream extends DocIdStream {
+
+  private final DocIdSetIterator iterator;
+  private final int to;
+  private final FixedBitSet spare;
+
+  DISIDocIdStream(DocIdSetIterator iterator, int to, FixedBitSet spare) {
+    if (to - iterator.docID() > spare.length()) {
+      throw new IllegalArgumentException("Bit set is too small to hold all potential matches");
+    }
+    this.iterator = iterator;
+    this.to = to;
+    this.spare = spare;
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return iterator.docID() < to;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    // If there are no live docs to apply, loading matching docs into a bit set and then iterating
+    // bits is unlikely to beat iterating the iterator directly.
+    upTo = Math.min(upTo, to);
+    for (int doc = iterator.docID(); doc < upTo; doc = iterator.nextDoc()) {
+      consumer.accept(doc);
+    }
+  }
+
+  @Override
+  public int count(int upTo) throws IOException {
+    if (iterator.docID() >= upTo) {
+      return 0;
+    }
+    // If the collector is just interested in the count, loading in a bit set and counting bits is
+    // often faster than incrementing a counter on every call to nextDoc().
+    assert spare.scanIsEmpty();
path: lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
user: login: gsmiller
body: Nice. I appreciate this assert in here!
created_at: 2025-03-27T17:29:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -204,6 +205,39 @@ public int cardinality() {
     return Math.toIntExact(tot);
   }
 
+  /**
+   * Return the number of set bits between indexes {@code from} inclusive and {@code to} exclusive.
+   */
+  public int cardinality(int from, int to) {
+    Objects.checkFromToIndex(from, to, length());
+
+    int cardinality = 0;
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        bits &= (1L << (to - from)) - 1L;
+        return Long.bitCount(bits);
+      }
+      cardinality += Long.bitCount(bits);
+      from += numBitsTilNextWord;
+    }
+
path: lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
user: login: gsmiller
body: minor: what about `assert (from & 0x3F) == 0;` right here?
created_at: 2025-03-27T19:37:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -790,4 +824,47 @@ public void applyMask(FixedBitSet bitSet, int offset) {
       throw new IllegalArgumentException("Some bits are set beyond the end of live docs");
     }
   }
+
+  /**
+   * For each set bit from {@code from} inclusive to {@code to} exclusive, add {@code base} to the
+   * bit index and call {@code consumer} on it. This is internally used by queries that use bit sets
+   * as intermediate representations of their matches.
+   */
+  public void forEach(int from, int to, int base, CheckedIntConsumer<IOException> consumer)
+      throws IOException {
+    Objects.checkFromToIndex(from, to, length());
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        // All bits are in a single word
+        bits &= (1L << (to - from)) - 1L;
+        forEach(bits, from + base, consumer);
+        return;
+      }
+      forEach(bits, from + base, consumer);
+      from += numBitsTilNextWord;
+    }
+
path: lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
user: login: gsmiller
body: minor: same suggestion of `assert (from & 0x3F) == 0;`
created_at: 2025-03-27T19:51:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class DISIDocIdStream extends DocIdStream {
+
+  private final DocIdSetIterator iterator;
+  private final int to;
+  private final FixedBitSet spare;
+
+  DISIDocIdStream(DocIdSetIterator iterator, int to, FixedBitSet spare) {
path: lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
user: login: gsmiller
body: Would it make sense to offer another ctor that takes care of allocating `spare`? I suppose there's not actually a use-case for it right now since the one calling code point of this currently has a bit set that can be reused, so maybe not worth adding until there's an actual reason.
created_at: 2025-03-27T19:59:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -207,8 +164,32 @@ private void scoreWindowIntoBitSetAndReplay(
       acceptDocs.applyMask(matching, base);
     }
 
-    docIdStreamView.base = base;
-    collector.collect(docIdStreamView);
+    if (buckets == null) {
+      if (acceptDocs != null) {
+        // In this case, live docs have not been applied yet.
+        acceptDocs.applyMask(matching, base);
+      }
path: lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
user: login: gsmiller
body: If I'm looking at this diff properly, I don't think you should need this block of code at all since you're still applying `acceptDocs` immediately prior?
created_at: 2025-03-27T20:06:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class BitSetDocIdStream extends DocIdStream {
+
+  private final FixedBitSet bitSet;
+  private final int offset, max;
+  private int upTo;
+
+  BitSetDocIdStream(FixedBitSet bitSet, int offset) {
+    this.bitSet = bitSet;
+    this.offset = offset;
+    upTo = offset;
+    max = (int) Math.min(Integer.MAX_VALUE, (long) offset + bitSet.length());
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo >= this.upTo) {
path: lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
user: login: gsmiller
body: I think you want `>` here instead? (`>=` still functionally works since `forEach` is tolerant of the `==` case, but I think you want to short circuit if `upTo == this.upTo`?)

(Also possible I'm getting this wrong... but I did make the changes locally and all the testing still seems to pass)
created_at: 2025-03-27T20:11:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class BitSetDocIdStream extends DocIdStream {
+
+  private final FixedBitSet bitSet;
+  private final int offset, max;
+  private int upTo;
+
+  BitSetDocIdStream(FixedBitSet bitSet, int offset) {
+    this.bitSet = bitSet;
+    this.offset = offset;
+    upTo = offset;
+    max = (int) Math.min(Integer.MAX_VALUE, (long) offset + bitSet.length());
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo >= this.upTo) {
+      upTo = Math.min(upTo, max);
+      bitSet.forEach(this.upTo - offset, upTo - offset, offset, consumer);
+      this.upTo = upTo;
+    }
+  }
+
+  @Override
+  public int count(int upTo) throws IOException {
+    if (upTo >= this.upTo) {
path: lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
user: login: gsmiller
body: Another place where I think you want `>`
created_at: 2025-03-27T20:11:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+
+final class RangeDocIdStream extends DocIdStream {
+
+  private int upTo;
+  private final int max;
+
+  RangeDocIdStream(int min, int max) {
+    if (max < min) {
+      throw new IllegalArgumentException("min = " + min + " > max = " + max);
+    }
+    this.upTo = min;
+    this.max = max;
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo >= this.upTo) {
path: lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
user: login: gsmiller
body: Another place where I think you want `>`
created_at: 2025-03-27T20:12:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+
+final class RangeDocIdStream extends DocIdStream {
+
+  private int upTo;
+  private final int max;
+
+  RangeDocIdStream(int min, int max) {
+    if (max < min) {
+      throw new IllegalArgumentException("min = " + min + " > max = " + max);
+    }
+    this.upTo = min;
+    this.max = max;
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo >= this.upTo) {
+      upTo = Math.min(upTo, max);
+      for (int doc = this.upTo; doc < upTo; ++doc) {
+        consumer.accept(doc);
+      }
+      this.upTo = upTo;
+    }
+  }
+
+  @Override
+  public int count(int upTo) throws IOException {
+    if (upTo >= this.upTo) {
path: lucene/core/src/java/org/apache/lucene/search/RangeDocIdStream.java
user: login: gsmiller
body: And one more spot where I think you want `>`?
created_at: 2025-03-27T20:12:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,12 +33,34 @@ protected DocIdStream() {}
    * Iterate over doc IDs contained in this stream in order, calling the given {@link
    * CheckedIntConsumer} on them. This is a terminal operation.
    */
-  public abstract void forEach(CheckedIntConsumer<IOException> consumer) throws IOException;
+  public void forEach(CheckedIntConsumer<IOException> consumer) throws IOException {
+    forEach(DocIdSetIterator.NO_MORE_DOCS, consumer);
+  }
+
+  /**
+   * Iterate over doc IDs contained in this doc ID stream up to the given {@code upTo} exclusive,
+   * calling the given {@link CheckedIntConsumer} on them. It is not possible to iterate these doc
+   * IDs again later on.
+   */
+  public abstract void forEach(int upTo, CheckedIntConsumer<IOException> consumer)
+      throws IOException;
 
   /** Count the number of entries in this stream. This is a terminal operation. */
   public int count() throws IOException {
-    int[] count = new int[1];
-    forEach(_ -> count[0]++);
-    return count[0];
+    return count(DocIdSetIterator.NO_MORE_DOCS);
   }
+
+  /**
+   * Count the number of doc IDs in this stream that are below the given {@code upTo}. These doc IDs
+   * may not be consumed again later.
+   */
+  // Note: it's abstract rather than having a default impl that delegates to #forEach because doing
path: lucene/core/src/java/org/apache/lucene/search/DocIdStream.java
user: login: jpountz
body: Ah right, it was based on your previous feedback. :)
created_at: 2025-03-27T20:17:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class DISIDocIdStream extends DocIdStream {
+
+  private final DocIdSetIterator iterator;
+  private final int to;
+  private final FixedBitSet spare;
+
+  DISIDocIdStream(DocIdSetIterator iterator, int to, FixedBitSet spare) {
path: lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
user: login: jpountz
body: I'm leaning towards waiting until there's a need for it before adding it.
created_at: 2025-03-27T20:18:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class BitSetDocIdStream extends DocIdStream {
+
+  private final FixedBitSet bitSet;
+  private final int offset, max;
+  private int upTo;
+
+  BitSetDocIdStream(FixedBitSet bitSet, int offset) {
+    this.bitSet = bitSet;
+    this.offset = offset;
+    upTo = offset;
+    max = (int) Math.min(Integer.MAX_VALUE, (long) offset + bitSet.length());
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo >= this.upTo) {
path: lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
user: login: jpountz
body: You're right, I'll change this.
created_at: 2025-03-27T20:19:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -207,8 +164,32 @@ private void scoreWindowIntoBitSetAndReplay(
       acceptDocs.applyMask(matching, base);
     }
 
-    docIdStreamView.base = base;
-    collector.collect(docIdStreamView);
+    if (buckets == null) {
+      if (acceptDocs != null) {
+        // In this case, live docs have not been applied yet.
+        acceptDocs.applyMask(matching, base);
+      }
path: lucene/core/src/java/org/apache/lucene/search/BooleanScorer.java
user: login: jpountz
body: Thanks for catching, this bit of code was mistakenly duplicated.
created_at: 2025-03-27T20:24:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class BitSetDocIdStream extends DocIdStream {
+
+  private final FixedBitSet bitSet;
+  private final int offset, max;
+  private int upTo;
+
+  BitSetDocIdStream(FixedBitSet bitSet, int offset) {
+    this.bitSet = bitSet;
+    this.offset = offset;
+    upTo = offset;
+    max = (int) Math.min(Integer.MAX_VALUE, (long) offset + bitSet.length());
+  }
+
+  @Override
+  public boolean mayHaveRemaining() {
+    return upTo < max;
+  }
+
+  @Override
+  public void forEach(int upTo, CheckedIntConsumer<IOException> consumer) throws IOException {
+    if (upTo >= this.upTo) {
path: lucene/core/src/java/org/apache/lucene/search/BitSetDocIdStream.java
user: login: jpountz
body: (For reference, it's not as much for short-circuiting as for the fact that logic under the if block would fail if trying to move backwards. But since we're doing a check anyway, I agree it's also worth excluding the trivial case when the range of docs to collect is empty.)
created_at: 2025-03-27T20:27:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -204,6 +205,40 @@ public int cardinality() {
     return Math.toIntExact(tot);
   }
 
+  /**
+   * Return the number of set bits between indexes {@code from} inclusive and {@code to} exclusive.
+   */
+  public int cardinality(int from, int to) {
+    Objects.checkFromToIndex(from, to, length());
+
+    int cardinality = 0;
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        bits &= (1L << (to - from)) - 1L;
+        return Long.bitCount(bits);
+      }
+      cardinality += Long.bitCount(bits);
+      from += numBitsTilNextWord;
+      assert (from & 0x3F) == 0;
+    }
+
+    for (int i = from >> 6, end = to >> 6; i < end; ++i) {
+      cardinality += Long.bitCount(bits[i]);
+    }
+
+    // Now handle bits between the last complete word and to
+    if ((to & 0x3F) != 0) {
+      long bits = this.bits[to >> 6] << -to;
path: lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
user: login: gsmiller
body: minor: one other small comment. I noticed in the new `forEach` method you added you use `long bits = this.bits[to >> 6] & ((1L << to) - 1);`. Is the rationale here that you need to persist the correct number of trailing zeros in the `forEach` implementation but not in this case since you're doing a bit count? Is this approach (shifting by `-to`) more performant (I ask since you could use the same approach as `forEach` here too for consistency, so I assume you had a reason ;))
created_at: 2025-03-27T23:35:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search;
+
+import java.io.IOException;
+import org.apache.lucene.util.FixedBitSet;
+
+final class DISIDocIdStream extends DocIdStream {
+
+  private final DocIdSetIterator iterator;
+  private final int to;
+  private final FixedBitSet spare;
+
+  DISIDocIdStream(DocIdSetIterator iterator, int to, FixedBitSet spare) {
path: lucene/core/src/java/org/apache/lucene/search/DISIDocIdStream.java
user: login: gsmiller
body: I think that's fair. Thanks!
created_at: 2025-03-28T00:52:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -204,6 +205,40 @@ public int cardinality() {
     return Math.toIntExact(tot);
   }
 
+  /**
+   * Return the number of set bits between indexes {@code from} inclusive and {@code to} exclusive.
+   */
+  public int cardinality(int from, int to) {
+    Objects.checkFromToIndex(from, to, length());
+
+    int cardinality = 0;
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        bits &= (1L << (to - from)) - 1L;
+        return Long.bitCount(bits);
+      }
+      cardinality += Long.bitCount(bits);
+      from += numBitsTilNextWord;
+      assert (from & 0x3F) == 0;
+    }
+
+    for (int i = from >> 6, end = to >> 6; i < end; ++i) {
+      cardinality += Long.bitCount(bits[i]);
+    }
+
+    // Now handle bits between the last complete word and to
+    if ((to & 0x3F) != 0) {
+      long bits = this.bits[to >> 6] << -to;
path: lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
user: login: jpountz
body: You are right. We need the low "to % 64" bits of `bits[to >> 6]`. `x << -to` is nice because it requires fewer instructions but it also moves bits to a different index. This makes iterating over set bits slightly more cumbersome, but doesn't matter for counting bits. Hence why `forEach` applies a mask instead. I haven't checked actual performance, I suspect that it doesn't matter in practice.
created_at: 2025-03-28T13:25:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -204,6 +205,40 @@ public int cardinality() {
     return Math.toIntExact(tot);
   }
 
+  /**
+   * Return the number of set bits between indexes {@code from} inclusive and {@code to} exclusive.
+   */
+  public int cardinality(int from, int to) {
+    Objects.checkFromToIndex(from, to, length());
+
+    int cardinality = 0;
+
+    // First, align `from` with a word start, ie. a multiple of Long.SIZE (64)
+    if ((from & 0x3F) != 0) {
+      long bits = this.bits[from >> 6] >>> from;
+      int numBitsTilNextWord = -from & 0x3F;
+      if (to - from < numBitsTilNextWord) {
+        bits &= (1L << (to - from)) - 1L;
+        return Long.bitCount(bits);
+      }
+      cardinality += Long.bitCount(bits);
+      from += numBitsTilNextWord;
+      assert (from & 0x3F) == 0;
+    }
+
+    for (int i = from >> 6, end = to >> 6; i < end; ++i) {
+      cardinality += Long.bitCount(bits[i]);
+    }
+
+    // Now handle bits between the last complete word and to
+    if ((to & 0x3F) != 0) {
+      long bits = this.bits[to >> 6] << -to;
path: lucene/core/src/java/org/apache/lucene/util/FixedBitSet.java
user: login: gsmiller
body: Thanks for the detailed explanation. This makes sense to me, it just left me scratching my head for a little bit initially to figure out the intention behind the different approaches :)
created_at: 2025-03-28T14:17:19Z
author_association: CONTRIBUTOR

PR comments:
user: login: jpountz
created_at: 2025-02-21T15:42:35Z
body: @epotyom You may be interested in taking a look.
author_association: CONTRIBUTOR
user: login: jpountz
created_at: 2025-02-24T15:59:34Z
body: > Looks like we can do similar trick for range facets and long values facets?

This is right.
author_association: CONTRIBUTOR
user: login: gsmiller
created_at: 2025-03-25T15:26:22Z
body: +1 to this optimization. Love the idea!
author_association: CONTRIBUTOR
user: login: jpountz
created_at: 2025-03-25T21:38:58Z
body: Quick update: we now have more queries that collect hits using `collect(DocIdStream)`, which makes this optimization more appealing.
author_association: CONTRIBUTOR
user: login: gsmiller
created_at: 2025-03-26T14:15:41Z
body: > I like the idea! Looks like we can do similar trick for range facets and long values facets?

I _think_ we could optimize these use-cases even further by potentially skipping over docs that don't fall into any of the ranges/values as well. With the histogram collection use-case, we case about the entire value range of the field we're interested in, but that's not necessarily true of these other use-cases. If we have a skipper, I think we ought to also be able to use competitive iterators to jump over blocks of docs we know we won't collect based on their values?

Maybe we need a spin-off issue :). I created #14406 
author_association: CONTRIBUTOR
user: login: jpountz
created_at: 2025-03-26T16:20:17Z
body: > If we have a skipper, I think we ought to also be able to use competitive iterators to jump over blocks of docs we know we won't collect based on their values?

This is correct. I plan on doing something similar when sorting: it is safe to skip blocks whose values all compare worse than the current k-th value. It's similar to what block-max WAND/MAXSCORE do: when a block's best possible score is less than the k-th best score so far, it can safely be skipped.
author_association: CONTRIBUTOR
user: login: jpountz
created_at: 2025-03-26T23:20:43Z
body: It should be ready for review now. Now that `DocIdStream` has become more sophisticated, I extracted impls to proper classes that could be better tested. This causes some diffs in our boolean scorers, hence the high number of lines changed.
author_association: CONTRIBUTOR
user: login: jpountz
created_at: 2025-03-26T23:25:50Z
body: I'll try to run some simple benchmarks next.
author_association: CONTRIBUTOR
user: login: jpountz
created_at: 2025-03-27T14:33:00Z
body: I played with the geonames dataset, by filtering out docs that don't have a value for the `elevation` field (2.3M docs left), enabling index sorting on the `elevation` field and computing histograms on the `elevation` field with a bucket width of 100.

| Query |  | Latency on main (ms) | Latency on branch (ms) |
| - | - | - | - |
| `MatchAllDocsQuery` | Uses `RangeDocIdStream` under the hood | 6.9 | 4.3 |
| `featureClass:(S P)` | Matches spots or cities, 1.2 matching docs, uses `BitSetDocIdStream` under the hood | 4.8 | 2.4 |

I also checked wikibigall, no slowdowns were detected.
author_association: CONTRIBUTOR