Support DocIdSetBuilder with partition bounds

lucene/core/src/java/org/apache/lucene/search/FilterWeight.java

@@ -67,4 +67,10 @@ public Matches matches(LeafReaderContext context, int doc) throws IOException {
   public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
     return in.scorerSupplier(context);
   }
+
+  @Override
+  public ScorerSupplier scorerSupplier(IndexSearcher.LeafReaderContextPartition partition)
+      throws IOException {
+    return in.scorerSupplier(partition);
+  }
 }

lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java

@@ -321,7 +321,7 @@ public QueryCachingPolicy getQueryCachingPolicy() {
    * href="https://github.com/apache/lucene/issues/13745">the corresponding github issue</a>.
    */
   protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
-    return slices(leaves, MAX_DOCS_PER_SLICE, MAX_SEGMENTS_PER_SLICE, false);
+    return slices(leaves, MAX_DOCS_PER_SLICE, MAX_SEGMENTS_PER_SLICE, true);
   }
 
   /**
@@ -828,7 +828,16 @@ protected void searchLeaf(
       // continue with the following leaf
       return;
     }
-    ScorerSupplier scorerSupplier = weight.scorerSupplier(ctx);
+
+    // Create partition object to pass to weight
+    final LeafReaderContextPartition partition;
+    if (minDocId == 0 && maxDocId == DocIdSetIterator.NO_MORE_DOCS) {
+      partition = LeafReaderContextPartition.createForEntireSegment(ctx);
+    } else {
+      partition = LeafReaderContextPartition.createFromAndTo(ctx, minDocId, maxDocId);
+    }
+
+    ScorerSupplier scorerSupplier = weight.scorerSupplier(partition);
     if (scorerSupplier != null) {
       scorerSupplier.setTopLevelScoringClause();
       BulkScorer scorer = scorerSupplier.bulkScorer();

lucene/core/src/java/org/apache/lucene/search/PointRangeQuery.java

@@ -249,8 +249,9 @@ public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
       }
 
       @Override
-      public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
-        LeafReader reader = context.reader();
+      public ScorerSupplier scorerSupplier(IndexSearcher.LeafReaderContextPartition partition)
+          throws IOException {
+        LeafReader reader = partition.ctx.reader();
 
         PointValues values = reader.getPointValues(field);
         if (checkValidPointValues(values) == false) {
@@ -298,7 +299,11 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
         } else {
           return new ConstantScoreScorerSupplier(score(), scoreMode, reader.maxDoc()) {
 
-            final DocIdSetBuilder result = new DocIdSetBuilder(reader.maxDoc(), values);
+            // Create partition-aware DocIdSetBuilder that filters docs and uses partition-sized
+            // threshold
+            final DocIdSetBuilder result =
+                new DocIdSetBuilder(
+                    reader.maxDoc(), values, partition.minDocId, partition.maxDocId);
             final IntersectVisitor visitor = getIntersectVisitor(result);
             long cost = -1;
 
@@ -336,6 +341,13 @@ public long cost() {
         }
       }
 
+      @Override
+      public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException {
+        // Delegate to partition-aware version for entire segment
+        return scorerSupplier(
+            IndexSearcher.LeafReaderContextPartition.createForEntireSegment(context));
+      }
+
       @Override
       public int count(LeafReaderContext context) throws IOException {
         LeafReader reader = context.reader();

lucene/core/src/java/org/apache/lucene/search/Weight.java

@@ -21,6 +21,7 @@
 import org.apache.lucene.index.IndexReaderContext;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.IndexSearcher.LeafReaderContextPartition;
 import org.apache.lucene.util.Bits;
 
 /**
@@ -149,6 +150,30 @@ public final Scorer scorer(LeafReaderContext context) throws IOException {
    */
   public abstract ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOException;
 
+  /**
+   * Returns a {@link ScorerSupplier}, which can then be used to get a {@link Scorer} for a
+   * partition of a leaf reader context.
+   *
+   * <p>This method allows queries to optimize for intra-segment concurrency by knowing the specific
+   * doc ID range being searched within the segment. The default implementation delegates to {@link
+   * #scorerSupplier(LeafReaderContext)} ignoring the partition bounds. Queries that can benefit
+   * from partition awareness (e.g., by creating smaller data structures scoped to the partition)
+   * should override this method.
+   *
+   * <p>A scorer supplier for the same {@link LeafReaderContext} instance may be requested multiple
+   * times as part of a single search call, potentially from different threads searching different
+   * doc ID ranges concurrently.
+   *
+   * @param partition the leaf reader context partition containing the context and doc ID range
+   * @return a {@link ScorerSupplier} providing the scorer, or null if scorer is null
+   * @throws IOException if an IOException occurs
+   * @see LeafReaderContextPartition
+   * @since 10.1
+   */
+  public ScorerSupplier scorerSupplier(LeafReaderContextPartition partition) throws IOException {
+    return scorerSupplier(partition.ctx);
+  }
+
   /**
    * Helper method that delegates to {@link #scorerSupplier(LeafReaderContext)}. It is implemented
    * as

lucene/core/src/java/org/apache/lucene/util/DocIdSetBuilder.java

@@ -41,7 +41,8 @@ public final class DocIdSetBuilder {
    *
    * @see DocIdSetBuilder#grow
    */
-  public sealed interface BulkAdder permits FixedBitSetAdder, BufferAdder {
+  public sealed interface BulkAdder
+      permits PartitionAwareFixedBitSetAdder, PartitionAwareBufferAdder {
     void add(int doc);
 
     void add(IntsRef docs);
@@ -51,32 +52,54 @@ public sealed interface BulkAdder permits FixedBitSetAdder, BufferAdder {
     void add(IntsRef docs, int docLowerBoundInclusive);
   }
 
-  private record FixedBitSetAdder(FixedBitSet bitSet) implements BulkAdder {
+  /**
+   * Partition-aware FixedBitSetAdder that filters docs to only include those within the specified
+   * range. Stores docs using partition-relative indices (doc - offset) to save memory.
+   *
+   * @param bitSet the partition-sized bitset to store relative doc indices
+   * @param minDocId minimum doc ID (inclusive) to accept
+   * @param maxDocId maximum doc ID (exclusive) to accept
+   * @param offset the value to subtract from absolute doc IDs (typically minDocId)
+   */
+  private record PartitionAwareFixedBitSetAdder(
+      FixedBitSet bitSet, int minDocId, int maxDocId, int offset) implements BulkAdder {
 
     @Override
     public void add(int doc) {
-      bitSet.set(doc);
+      if (doc >= minDocId && doc < maxDocId) {
+        bitSet.set(doc - offset);
+      }
     }
 
     @Override
     public void add(IntsRef docs) {
       for (int i = docs.offset, to = docs.offset + docs.length; i < to; i++) {
-        bitSet.set(docs.ints[i]);
+        int doc = docs.ints[i];
+        if (doc >= minDocId && doc < maxDocId) {
+          bitSet.set(doc - offset);
+        }
       }
     }
 
     @Override
     public void add(DocIdSetIterator iterator) throws IOException {
-      iterator.nextDoc();
-      iterator.intoBitSet(DocIdSetIterator.NO_MORE_DOCS, bitSet, 0);
+      // Advance iterator to minDocId first
+      int doc = iterator.nextDoc();
+      if (doc < minDocId) {
+        doc = iterator.advance(minDocId);
+      }
+      // Use optimized intoBitSet with partition boundaries and offset
+      if (doc < maxDocId) {
+        iterator.intoBitSet(maxDocId, bitSet, offset);
+      }
     }
 
     @Override
     public void add(IntsRef docs, int docLowerBoundInclusive) {
       for (int i = docs.offset, to = docs.offset + docs.length; i < to; i++) {
         int doc = docs.ints[i];
-        if (doc >= docLowerBoundInclusive) {
-          bitSet.set(doc);
+        if (doc >= Math.max(docLowerBoundInclusive, minDocId) && doc < maxDocId) {
+          bitSet.set(doc - offset);
         }
       }
     }
@@ -97,17 +120,29 @@ private static class Buffer {
     }
   }
 
-  private record BufferAdder(Buffer buffer) implements BulkAdder {
+  /**
+   * Partition-aware BufferAdder that filters docs to only include those within the specified range.
+   */
+  private record PartitionAwareBufferAdder(Buffer buffer, int minDocId, int maxDocId)
+      implements BulkAdder {
 
     @Override
     public void add(int doc) {
-      buffer.array[buffer.length++] = doc;
+      if (doc >= minDocId && doc < maxDocId) {
+        buffer.array[buffer.length++] = doc;
+      }
     }
 
     @Override
     public void add(IntsRef docs) {
-      System.arraycopy(docs.ints, docs.offset, buffer.array, buffer.length, docs.length);
-      buffer.length += docs.length;
+      int index = buffer.length;
+      for (int i = docs.offset, to = docs.offset + docs.length; i < to; i++) {
+        int doc = docs.ints[i];
+        if (doc >= minDocId && doc < maxDocId) {
+          buffer.array[index++] = doc;
+        }
+      }
+      buffer.length = index;
     }
 
     @Override
@@ -123,7 +158,7 @@ public void add(IntsRef docs, int docLowerBoundInclusive) {
       int index = buffer.length;
       for (int i = docs.offset, to = docs.offset + docs.length; i < to; i++) {
         int doc = docs.ints[i];
-        if (doc >= docLowerBoundInclusive) {
+        if (doc >= Math.max(docLowerBoundInclusive, minDocId) && doc < maxDocId) {
           buffer.array[index++] = doc;
         }
       }
@@ -137,6 +172,10 @@ public void add(IntsRef docs, int docLowerBoundInclusive) {
   final boolean multivalued;
   final double numValuesPerDoc;
 
+  // Partition filtering support - filters docs to only include those within [minDocId, maxDocId)
+  private final int minDocId; // inclusive
+  private final int maxDocId; // exclusive
+
   private List<Buffer> buffers = new ArrayList<>();
   private int totalAllocated; // accumulated size of the allocated buffers
 
@@ -166,8 +205,57 @@ public DocIdSetBuilder(int maxDoc, PointValues values) throws IOException {
     this(maxDoc, values.getDocCount(), values.size());
   }
 
+  /**
+   * Create a partition-aware {@link DocIdSetBuilder} instance that only accepts doc IDs within the
+   * specified range. This is useful for intra-segment concurrency where each partition only needs
+   * to collect docs within its assigned range.
+   *
+   * @param maxDoc the maximum doc ID in the segment
+   * @param minDocId the minimum doc ID (inclusive) to accept
+   * @param maxDocId the maximum doc ID (exclusive) to accept
+   * @param docCount estimated document count
+   * @param valueCount estimated value count
+   */
+  public DocIdSetBuilder(int maxDoc, int minDocId, int maxDocId, int docCount, long valueCount) {
+    this(maxDoc, docCount, valueCount, minDocId, maxDocId);
+  }
+
+  /**
+   * Create a partition-aware {@link DocIdSetBuilder} for {@link PointValues} that only accepts doc
+   * IDs within the specified range.
+   *
+   * @param maxDoc the maximum doc ID in the segment
+   * @param values the point values
+   * @param minDocId the minimum doc ID (inclusive) to accept
+   * @param maxDocId the maximum doc ID (exclusive) to accept
+   */
+  public DocIdSetBuilder(int maxDoc, PointValues values, int minDocId, int maxDocId)
+      throws IOException {
+    this(maxDoc, values.getDocCount(), values.size(), minDocId, maxDocId);
+  }
+
+  /**
+   * Create a partition-aware {@link DocIdSetBuilder} for {@link Terms} that only accepts doc IDs
+   * within the specified range.
+   *
+   * @param maxDoc the maximum doc ID in the segment
+   * @param terms the terms
+   * @param minDocId the minimum doc ID (inclusive) to accept
+   * @param maxDocId the maximum doc ID (exclusive) to accept
+   */
+  public DocIdSetBuilder(int maxDoc, Terms terms, int minDocId, int maxDocId) throws IOException {
+    this(maxDoc, terms.getDocCount(), terms.getSumDocFreq(), minDocId, maxDocId);
+  }
+
   DocIdSetBuilder(int maxDoc, int docCount, long valueCount) {
+    this(maxDoc, docCount, valueCount, 0, maxDoc);
+  }
+
+  private DocIdSetBuilder(int maxDoc, int docCount, long valueCount, int minDocId, int maxDocId) {
     this.maxDoc = maxDoc;
+    this.minDocId = minDocId;
+    this.maxDocId = maxDocId;
+
     this.multivalued = docCount < 0 || docCount != valueCount;
     if (docCount <= 0 || valueCount < 0) {
       // assume one value per doc, this means the cost will be overestimated
@@ -184,7 +272,12 @@ public DocIdSetBuilder(int maxDoc, PointValues values) throws IOException {
     // maxDoc >>> 7 is a good value if you want to save memory, lower values
     // such as maxDoc >>> 11 should provide faster building but at the expense
     // of using a full bitset even for quite sparse data
-    this.threshold = maxDoc >>> 7;
+    //
+    // When filtering to a partition (minDocId > 0 or maxDocId < maxDoc), use the partition size
+    // for threshold calculation to ensure the threshold scales correctly with the partition size
+    boolean isPartition = (minDocId > 0 || maxDocId < maxDoc);
+    int effectiveMaxDoc = isPartition ? (maxDocId - minDocId) : maxDoc;
+    this.threshold = effectiveMaxDoc >>> 7;
 
     this.bitSet = null;
   }
@@ -267,7 +360,10 @@ private int additionalCapacity(int numDocs) {
   private Buffer addBuffer(int len) {
     Buffer buffer = new Buffer(len);
     buffers.add(buffer);
-    adder = new BufferAdder(buffer);
+    // Always use partition-aware adder to avoid megamorphic call sites
+    // For non-partitioned case (minDocId=0, maxDocId=maxDoc), the bounds check
+    // becomes a predictable branch that the JIT optimizes away
+    adder = new PartitionAwareBufferAdder(buffer, minDocId, maxDocId);
     totalAllocated += buffer.array.length;
     return buffer;
   }
@@ -279,20 +375,33 @@ private void growBuffer(Buffer buffer, int additionalCapacity) {
 
   private void upgradeToBitSet() {
     assert bitSet == null;
-    FixedBitSet bitSet = new FixedBitSet(maxDoc);
+
+    // For partitions, create a smaller bitset sized to the partition range only
+    // This saves memory by not allocating bits outside [minDocId, maxDocId)
+    boolean isPartition = (minDocId > 0 || maxDocId < maxDoc);
+    int bitSetSize = isPartition ? (maxDocId - minDocId) : maxDoc;
+
+    FixedBitSet bitSet = new FixedBitSet(bitSetSize);
     long counter = 0;
     for (Buffer buffer : buffers) {
       int[] array = buffer.array;
       int length = buffer.length;
       counter += length;
       for (int i = 0; i < length; ++i) {
-        bitSet.set(array[i]);
+        // For partitions, convert absolute doc ID to partition-relative index
+        int docId = array[i];
+        int bitIndex = isPartition ? (docId - minDocId) : docId;
+        bitSet.set(bitIndex);
       }
     }
     this.bitSet = bitSet;
     this.counter = counter;
     this.buffers = null;
-    this.adder = new FixedBitSetAdder(bitSet);
+    // Always use partition-aware adder to avoid megamorphic call sites
+    // For non-partitioned case, use offset=0 and bounds=[0, maxDoc)
+    // The JIT will optimize away the redundant checks for non-partitioned case
+    int offset = isPartition ? minDocId : 0;
+    this.adder = new PartitionAwareFixedBitSetAdder(bitSet, minDocId, maxDocId, offset);
   }
 
   /** Build a {@link DocIdSet} from the accumulated doc IDs. */
@@ -301,7 +410,14 @@ public DocIdSet build() {
       if (bitSet != null) {
         assert counter >= 0;
         final long cost = Math.round(counter / numValuesPerDoc);
-        return new BitDocIdSet(bitSet, cost);
+
+        // For partition-relative bitsets, wrap with offset to return absolute doc IDs
+        boolean isPartition = (minDocId > 0 || maxDocId < maxDoc);
+        if (isPartition) {
+          return new OffsetBitDocIdSet(bitSet, cost, minDocId);
+        } else {
+          return new BitDocIdSet(bitSet, cost);
+        }
       } else {
         Buffer concatenated = concat(buffers);
         LSBRadixSorter sorter = new LSBRadixSorter();