Initial impl of intra segment concurrency

- Added LeafReaderContextPartition abstraction that holds LeafReaderContext + the range of doc ids it targets. A slice points now to specific subsets of segments, identified by their corresponding range of doc ids.
- Introduced additional IndexSearcher#search method that takes LeafReaderContextPartition[] instead of LeafReaderContext[], which calls `scorer.score(leafCollector, ctx.reader().getLiveDocs(), slice.minDocId, slice.maxDocId);` providing the range of doc ids in place of `scorer.score(leafCollector, ctx.reader().getLiveDocs());` that would score all documents
- Added override for new protected IndexSearcher#search to subclasses that require it
- Adjusted TotalHitCountCollectorManager to hold state used to track seen leaves and ensure consistent treatment for different partitions of the same leaf
- Updated IndexSearcher#getSlices to partition segments that go over maxDocsPerSlice into multiple partitions: this first iteration does not mix up partitions with entire segments in the same slice.

Problems found:
1) IndexSearcher#count / TotalHitCountCollector rely on `Weight#count(LeafReaderContext)`, which now gets called multiple times against the same leaf and leads to excessive counting of hits.
2) LRUQueryCache caches the return value of `Weight#count`. When we execute the same query against the same segment multiple times (as part of a single search call), the first time we do the actual counting for the docs that the first partition holds, and subsequent times we should do the same, count hits in each partition of the same segment instead of retrieving the count from the cache.
3) CheckHits verifies matches and may go outside of the bounds of the doc id range of the current slice

Both 1) and 2) have been addressed by tracking which seen leaves in TotalHitCountCollectorManager, and extending the TotalHitCountCollector that newCollector returns according to such internal state

Author: javanna

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

@@ -21,8 +21,10 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
 import java.util.function.Function;
@@ -233,7 +235,13 @@ public IndexSearcher(IndexReaderContext context, Executor executor) {
             ? leaves ->
                 leaves.isEmpty()
                     ? new LeafSlice[0]
-                    : new LeafSlice[] {new LeafSlice(new ArrayList<>(leaves))}
+                    : new LeafSlice[] {
+                      new LeafSlice(
+                          new ArrayList<>(
+                              leaves.stream()
+                                  .map(LeafReaderContextPartition::createForEntireSegment)
+                                  .toList()))
+                    }
             : this::slices;
     leafSlicesSupplier = new CachingLeafSlicesSupplier(slicesProvider, leafContexts);
   }
@@ -328,42 +336,65 @@ protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
   /** Static method to segregate LeafReaderContexts amongst multiple slices */
   public static LeafSlice[] slices(
       List<LeafReaderContext> leaves, int maxDocsPerSlice, int maxSegmentsPerSlice) {
+
+    // TODO this is a temporary hack to force testing against multiple leaf reader context slices.
+    // It must be reverted before merging.
+    // maxDocsPerSlice = 1;
+    // maxSegmentsPerSlice = 1;
+    // end hack
+
     // Make a copy so we can sort:
     List<LeafReaderContext> sortedLeaves = new ArrayList<>(leaves);
 
     // Sort by maxDoc, descending:
-    Collections.sort(
-        sortedLeaves, Collections.reverseOrder(Comparator.comparingInt(l -> l.reader().maxDoc())));
+    sortedLeaves.sort(Collections.reverseOrder(Comparator.comparingInt(l -> l.reader().maxDoc())));
 
-    final List<List<LeafReaderContext>> groupedLeaves = new ArrayList<>();
-    long docSum = 0;
-    List<LeafReaderContext> group = null;
+    final List<List<LeafReaderContextPartition>> groupedLeafPartitions = new ArrayList<>();
+    int currentSliceNumDocs = 0;
+    List<LeafReaderContextPartition> group = null;
     for (LeafReaderContext ctx : sortedLeaves) {
       if (ctx.reader().maxDoc() > maxDocsPerSlice) {
         assert group == null;
-        groupedLeaves.add(Collections.singletonList(ctx));
+        // if the segment does not fit in a single slice, we split it in multiple partitions of
+        // equal size
+        int numSlices = Math.ceilDiv(ctx.reader().maxDoc(), maxDocsPerSlice);
+        int numDocs = ctx.reader().maxDoc() / numSlices;
+        int maxDocId = numDocs;
+        int minDocId = 0;
+        for (int i = 0; i < numSlices - 1; i++) {
+          groupedLeafPartitions.add(
+              Collections.singletonList(
+                  LeafReaderContextPartition.createFromAndTo(ctx, minDocId, maxDocId)));
+          minDocId = maxDocId;
+          maxDocId += numDocs;
+        }
+        // the last slice gets all the remaining docs
+        groupedLeafPartitions.add(
+            Collections.singletonList(LeafReaderContextPartition.createFrom(ctx, minDocId)));
       } else {
         if (group == null) {
           group = new ArrayList<>();
-          group.add(ctx);
-
-          groupedLeaves.add(group);
-        } else {
-          group.add(ctx);
+          groupedLeafPartitions.add(group);
         }
-
-        docSum += ctx.reader().maxDoc();
-        if (group.size() >= maxSegmentsPerSlice || docSum > maxDocsPerSlice) {
+        group.add(LeafReaderContextPartition.createForEntireSegment(ctx));
+
+        currentSliceNumDocs += ctx.reader().maxDoc();
+        // We only split a segment when it does not fit entirely in a slice. We don't partition the
+        // segment that makes the current slice (which holds multiple segments) go over
+        // maxDocsPerSlice.
+        // This means that a slice either contains multiple entire segments, or a single partition
+        // of a segment.
+        if (group.size() >= maxSegmentsPerSlice || currentSliceNumDocs > maxDocsPerSlice) {
           group = null;
-          docSum = 0;
+          currentSliceNumDocs = 0;
         }
       }
     }
 
-    LeafSlice[] slices = new LeafSlice[groupedLeaves.size()];
+    LeafSlice[] slices = new LeafSlice[groupedLeafPartitions.size()];
     int upto = 0;
-    for (List<LeafReaderContext> currentLeaf : groupedLeaves) {
-      slices[upto] = new LeafSlice(currentLeaf);
+    for (List<LeafReaderContextPartition> currentGroup : groupedLeafPartitions) {
+      slices[upto] = new LeafSlice(currentGroup);
       ++upto;
     }
 
@@ -658,11 +689,11 @@ private <C extends Collector, T> T search(
       }
       final List<Callable<C>> listTasks = new ArrayList<>(leafSlices.length);
       for (int i = 0; i < leafSlices.length; ++i) {
-        final LeafReaderContext[] leaves = leafSlices[i].leaves;
+        final LeafReaderContextPartition[] leaves = leafSlices[i].leaves;
         final C collector = collectors.get(i);
         listTasks.add(
             () -> {
-              search(Arrays.asList(leaves), weight, collector);
+              search(leaves, weight, collector);
               return collector;
             });
       }
@@ -671,6 +702,52 @@ private <C extends Collector, T> T search(
     }
   }
 
+  // TODO this is a copy of the existing search protected method that takes a list of leaf reader
+  // contexts
+  // perhaps more methods need to be switched to use it in place of the other one?
+  protected void search(LeafReaderContextPartition[] leaves, Weight weight, Collector collector)
+      throws IOException {
+
+    collector.setWeight(weight);
+
+    for (LeafReaderContextPartition slice : leaves) { // search each subreader partition
+      LeafReaderContext ctx = slice.ctx;
+      final LeafCollector leafCollector;
+      try {
+        leafCollector = collector.getLeafCollector(ctx);
+      } catch (
+          @SuppressWarnings("unused")
+          CollectionTerminatedException e) {
+        // there is no doc of interest in this reader context
+        // continue with the following leaf
+        continue;
+      }
+
+      // TODO we are pulling bulk scorer twice for the same ctx. will that work?
+      BulkScorer scorer = weight.bulkScorer(ctx);
+      if (scorer != null) {
+        if (queryTimeout != null) {
+          scorer = new TimeLimitingBulkScorer(scorer, queryTimeout);
+        }
+        try {
+          scorer.score(leafCollector, ctx.reader().getLiveDocs(), slice.minDocId, slice.maxDocId);
+        } catch (
+            @SuppressWarnings("unused")
+            CollectionTerminatedException e) {
+          // collection was terminated prematurely
+          // continue with the following leaf
+        } catch (
+            @SuppressWarnings("unused")
+            TimeLimitingBulkScorer.TimeExceededException e) {
+          partialResult = true;
+        }
+      }
+      // Note: this is called if collection ran successfully, including the above special cases of
+      // CollectionTerminatedException and TimeExceededException, but no other exception.
+      leafCollector.finish();
+    }
+  }
+
   /**
    * Lower-level search API.
    *
@@ -685,6 +762,9 @@ private <C extends Collector, T> T search(
    * @throws TooManyClauses If a query would exceed {@link IndexSearcher#getMaxClauseCount()}
    *     clauses.
    */
+
+  // TODO legacy search method that does not support intra segment concurrency - what do we do about
+  // it?
   protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector)
       throws IOException {
 
@@ -873,11 +953,70 @@ public static class LeafSlice {
      *
      * @lucene.experimental
      */
-    public final LeafReaderContext[] leaves;
+    public final LeafReaderContextPartition[] leaves;
+
+    public LeafSlice(List<LeafReaderContextPartition> leafReaderContextPartitions) {
+      leafReaderContextPartitions.sort(Comparator.comparingInt(l -> l.ctx.docBase));
+      // TODO should we sort by minDocId too?
+      this.leaves = leafReaderContextPartitions.toArray(new LeafReaderContextPartition[0]);
+    }
 
-    public LeafSlice(List<LeafReaderContext> leavesList) {
-      Collections.sort(leavesList, Comparator.comparingInt(l -> l.docBase));
-      this.leaves = leavesList.toArray(new LeafReaderContext[0]);
+    /**
+     * Returns the total number of docs that a slice targets, by summing the number of docs that
+     * each of its leaf context partitions targets.
+     */
+    public int getNumDocs() {
+      return Arrays.stream(leaves)
+          .map(LeafReaderContextPartition::getNumDocs)
+          .reduce(Integer::sum)
+          .get();
+    }
+  }
+
+  /**
+   * Holds information about a specific leaf context and the corresponding range of doc ids to
+   * search within.
+   *
+   * @lucene.experimental
+   */
+  public static final class LeafReaderContextPartition {
+    private final int minDocId;
+    private final int maxDocId;
+    private final int numDocs;
+    public final LeafReaderContext ctx;
+
+    private LeafReaderContextPartition(
+        LeafReaderContext leafReaderContext, int minDocId, int maxDocId, int numDocs) {
+      this.ctx = leafReaderContext;
+      this.minDocId = minDocId;
+      this.maxDocId = maxDocId;
+      this.numDocs = numDocs;
+    }
+
+    /** Returns The number of docs that the doc id range of this partition targets */
+    public int getNumDocs() {
+      return numDocs;
+    }
+
+    /** Creates a partition of the provided leaf context that targets the entire segment */
+    public static LeafReaderContextPartition createForEntireSegment(LeafReaderContext ctx) {
+      return new LeafReaderContextPartition(
+          ctx, 0, DocIdSetIterator.NO_MORE_DOCS, ctx.reader().maxDoc());
+    }
+
+    /**
+     * Creates a partition of the provided leaf context that targets a subset of the entire segment,
+     * starting from the min doc id provided, until the end of the segment
+     */
+    public static LeafReaderContextPartition createFrom(LeafReaderContext ctx, int minDocId) {
+      return new LeafReaderContextPartition(
+          ctx, minDocId, DocIdSetIterator.NO_MORE_DOCS, ctx.reader().maxDoc() - minDocId);
+    }
+
+    public static LeafReaderContextPartition createFromAndTo(
+        LeafReaderContext ctx, int minDocId, int maxDocId) {
+      assert maxDocId != DocIdSetIterator.NO_MORE_DOCS;
+      return new LeafReaderContextPartition(ctx, minDocId, maxDocId, maxDocId - minDocId);
     }
   }
 
@@ -1007,10 +1146,33 @@ public LeafSlice[] get() {
             leafSlices =
                 Objects.requireNonNull(
                     sliceProvider.apply(leaves), "slices computed by the provider is null");
+            checkSlices(leafSlices);
           }
         }
       }
       return leafSlices;
     }
+
+    /**
+     * Enforce that there aren't multiple slices pointing to the same physical segment. It is a
+     * requirement that {@link Collector#getLeafCollector(LeafReaderContext)} gets called once per
+     * leaf context. Also, it does not make sense to partition a segment to then search those
+     * partitions as part of the same slice, because the goal of partitioning is parallel searching
+     * which happens at the slices level.
+     */
+    private static void checkSlices(LeafSlice[] leafSlices) {
+      for (LeafSlice leafSlice : leafSlices) {
+        Set<LeafReaderContext> distinctLeaves = new HashSet<>();
+        for (LeafReaderContextPartition leafPartition : leafSlice.leaves) {
+          distinctLeaves.add(leafPartition.ctx);
+        }
+        if (leafSlice.leaves.length != distinctLeaves.size()) {
+          throw new IllegalStateException(
+              "The same slice targets multiple partitions of the same leaf reader. "
+                  + "A segment should rather get partitioned to be searched concurrently from as many slices as the "
+                  + "number of partitions it is split into.");
+        }
+      }
+    }
   }
 }

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

@@ -362,6 +362,9 @@ public long cost() {
             final IntersectVisitor visitor = getIntersectVisitor(result);
             long cost = -1;
 
+            // maybe allowing calling get multiple times, or clone the scorer to avoid duplication
+            // across multiple
+            // slices that point to the same segment
             @Override
             public Scorer get(long leadCost) throws IOException {
               if (values.getDocCount() == reader.maxDoc()
@@ -373,6 +376,7 @@ && cost() > reader.maxDoc() / 2) {
                 final FixedBitSet result = new FixedBitSet(reader.maxDoc());
                 result.set(0, reader.maxDoc());
                 long[] cost = new long[] {reader.maxDoc()};
+                // bitset etc.
                 values.intersect(getInverseIntersectVisitor(result, cost));
                 final DocIdSetIterator iterator = new BitSetIterator(result, cost[0]);
                 return new ConstantScoreScorer(score(), scoreMode, iterator);

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

@@ -50,13 +50,17 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept
       totalHits += leafCount;
       throw new CollectionTerminatedException();
     }
+    return createLeafCollector();
+  }
+
+  protected final LeafCollector createLeafCollector() {
     return new LeafCollector() {
 
       @Override
-      public void setScorer(Scorable scorer) throws IOException {}
+      public void setScorer(Scorable scorer) {}
 
       @Override
-      public void collect(int doc) throws IOException {
+      public void collect(int doc) {
         totalHits++;
       }