Better Segment To Thread Mapping Algorithm [LUCENE-8757]

[asfimport]

The current segments to threads allocation algorithm always allocates one thread per segment. This is detrimental to performance in case of skew in segment sizes since small segments also get their dedicated thread. This can lead to performance degradation due to context switching overheads.

 

A better algorithm which is cognizant of size skew would have better performance for realistic scenarios

---

Migrated from LUCENE-8757 by Atri Sharma (@atris), resolved May 21 2019
Attachments: image-2020-12-04-14-09-22-924.png, LUCENE-8757.patch (versions: 10)
Linked issues:

• org.apache.lucene.search.TestTopDocsMerge.testSort_1 failure [LUCENE-8819] #9863
• TopDocs#Merge is Tightly Coupled To Number Of Collectors Involved [LUCENE-8829] #9873

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

Attached is a first cut for the patch. The main idea there is that smaller segments are coalesced into single threads with a hard cap on the number of segments per thread to avoid overwhelming a single thread.

 

This can be enhanced with taking number of cores on the machine and current CPU utilization, but that might lead to a higher IndexSearcher bootstrap time.

 

Comments and thoughts are welcome.

[asfimport]

Michael McCandless (@mikemccand) (migrated from JIRA)

Thanks @atris – I agree it's important to have better defaults for how we coalesce segments into per-query-per-thread work units.  A few small comments:

• Can you insert _ in the big number constants (e.g. 25000000)?  Makes it easier to read, and open-source code is written for reading :)
• I think something is wrong with docSum – you only set it, and never add to it?  I think the intention is to sum up docs in multiple adjacent (sorted by maxDoc) segments until that count exceeds 25000000?
• How did you pick 25000000 and 100 as good constants?  We are using much smaller values in our production infrastructure – 250_000 and 5, admittedly after only a little experimentation. 
• Can you add some tests?  You can maybe make the slice method a package private static method and then create test cases with "interesting" LeafReaderContext combinations?  In particular, a test case exposing the docSum bug would be great, then fix that bug, then see the test case pass.

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

Hi @mikemccand

 

Thanks for taking a look at the patch. I have attached an updated patch which fixes your comments.

For the constants, I think I was being too conservative in not using too many threads :)

 

Please let me know if the current patch seems sane

[asfimport]

Simon Willnauer (@s1monw) (migrated from JIRA)

Hey Atri,

thanks for putting up this patch, here is some additional feedback:

• can we stick with an protected non-static method on IndexSearcher subclasses should be able to override your impl. I think it's ok to have a static method like this:

   public static LeafSlice[] slices (List<LeafReaderContext> leaves, int maxDocsPerSlice, int maxSegPerSlice)

  that you can call from the protected method with your defaults?

• you might want to change your sort to something like this: 

  Collections.sort(leaves, Collections.reverseOrder(Comparator.comparingInt(l -> l.reader().maxDoc())));

• I think the Leaves class is unnecessary we can just use List<LeafReaderContext> instead?

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

LUCENE-8757.patchHi @s1monw

 

Thanks for reviewing the patch and the comments. Attached is an updated patch.

Regards,

Atri

LUCENE-8757.patch

From a181fa0e6a9630e7f38760cb16827c4a80ce0d86 Mon Sep 17 00:00:00 2001
From: Atri Sharma <atri@apache.org>
Date: Mon, 22 Apr 2019 19:44:52 +0530
Subject: [PATCH] LUCENE-8757: Improving Default Segments To Thread Mapping
 Algorithm

The current slicing algorithm assigns a thread per segment, which
can be detrimental to performance in case the distribution has
a large number of small segments. The patch introduces a slicing
algorithm which coalesces smaller segments to a single thread,
thus reducing the impact of context switching by limiting the
number of threads
---
 .../apache/lucene/search/IndexSearcher.java   |  68 ++++-
 .../index/TestSegmentToThreadMapping.java     | 268 ++++++++++++++++++
 .../lucene/search/AssertingCollector.java     |   5 +
 .../apache/lucene/util/LuceneTestCase.java    |   9 +
 4 files changed, 342 insertions(+), 8 deletions(-)
 create mode 100644 lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java

diff --git a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
index 689409f5cc..1748dd4821 100644
--- a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
+++ b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
@@ -22,6 +22,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -100,6 +101,13 @@ public class IndexSearcher {
    */
   private static final int TOTAL_HITS_THRESHOLD = 1000;
 
+  /**
+   * Thresholds for index slice allocation logic. To change the default, extend
+   * <code> IndexSearcher</code> and use custom values
+   */
+  private static final int MAX_DOCS_PER_SLICE = 250_000;
+  private static final int MAX_SEGMENTS_PER_SLICE = 5;
+
   final IndexReader reader; // package private for testing!
   
   // NOTE: these members might change in incompatible ways
@@ -268,17 +276,60 @@ public class IndexSearcher {
 
   /**
    * Expert: Creates an array of leaf slices each holding a subset of the given leaves.
-   * Each {@link LeafSlice} is executed in a single thread. By default there
-   * will be one {@link LeafSlice} per leaf ({@link org.apache.lucene.index.LeafReaderContext}).
+   * Each {@link LeafSlice} is executed in a single thread. By default, segments with more than
+   * MAX_DOCS_PER_SLICE will get their own thread
    */
   protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
-    LeafSlice[] slices = new LeafSlice[leaves.size()];
-    for (int i = 0; i < slices.length; i++) {
-      slices[i] = new LeafSlice(leaves.get(i));
+    return slices(leaves, MAX_DOCS_PER_SLICE, MAX_SEGMENTS_PER_SLICE);
+  }
+
+  /**
+   * Static method to segregate LeafReaderContexts amongst multiple slices
+   */
+  public static LeafSlice[] slices (List<LeafReaderContext> leaves, int maxDocsPerSlice,
+                                    int maxSegmentsPerSlice) {
+    // 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())));
+
+    final List<List<LeafReaderContext>> groupedLeaves = new ArrayList<>();
+    long docSum = 0;
+    List<LeafReaderContext> group = null;
+    for (LeafReaderContext ctx : sortedLeaves) {
+      if (ctx.reader().maxDoc() > maxDocsPerSlice) {
+        assert group == null;
+        groupedLeaves.add(Collections.singletonList(ctx));
+      } else {
+        if (group == null) {
+          group = new ArrayList<>();
+          group.add(ctx);
+
+          groupedLeaves.add(group);
+        } else {
+          group.add(ctx);
+        }
+
+        docSum += ctx.reader().maxDoc();
+        if (group.size() >= maxSegmentsPerSlice || docSum > maxDocsPerSlice) {
+          group = null;
+          docSum = 0;
+        }
+      }
+    }
+
+    LeafSlice[] slices = new LeafSlice[groupedLeaves.size()];
+    int upto = 0;
+    for (List<LeafReaderContext> currentLeaf : groupedLeaves) {
+      slices[upto] = new LeafSlice(currentLeaf);
+      ++upto;
     }
+
     return slices;
   }
-  
+
   /** Return the {@link IndexReader} this searches. */
   public IndexReader getIndexReader() {
     return reader;
@@ -743,8 +794,9 @@ public class IndexSearcher {
      *  @lucene.experimental */
     public final LeafReaderContext[] leaves;
     
-    public LeafSlice(LeafReaderContext... leaves) {
-      this.leaves = leaves;
+    public LeafSlice(List<LeafReaderContext> leavesList) {
+      Collections.sort(leavesList, Comparator.comparingInt(l -> l.docBase));
+      this.leaves = leavesList.toArray(new LeafReaderContext[0]);
     }
   }
 
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
new file mode 100644
index 0000000000..62299dfd91
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
@@ -0,0 +1,268 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.LuceneTestCase;
+import org.apache.lucene.util.NamedThreadFactory;
+import org.apache.lucene.util.Version;
+
+public class TestSegmentToThreadMapping extends LuceneTestCase {
+
+  public LeafReader dummyIndexReader(final int maxDoc) {
+    return new LeafReader() {
+      @Override
+      public int maxDoc() {
+        return maxDoc;
+      }
+
+      @Override
+      public int numDocs() {
+        return maxDoc;
+      }
+
+      @Override
+      public FieldInfos getFieldInfos() {
+        return FieldInfos.EMPTY;
+      }
+
+      @Override
+      public Bits getLiveDocs() {
+        return null;
+      }
+
+      @Override
+      public Terms terms(String field) throws IOException {
+        return null;
+      }
+
+      @Override
+      public Fields getTermVectors(int doc) {
+        return null;
+      }
+
+      @Override
+      public NumericDocValues getNumericDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public BinaryDocValues getBinaryDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedDocValues getSortedDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedNumericDocValues getSortedNumericDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedSetDocValues getSortedSetDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public NumericDocValues getNormValues(String field) {
+        return null;
+      }
+
+      @Override
+      public PointValues getPointValues(String field) {
+        return null;
+      }
+
+      @Override
+      protected void doClose() {
+      }
+
+      @Override
+      public void document(int doc, StoredFieldVisitor visitor) {
+      }
+
+      @Override
+      public void checkIntegrity() throws IOException {
+      }
+
+      @Override
+      public LeafMetaData getMetaData() {
+        return new LeafMetaData(Version.LATEST.major, Version.LATEST, null);
+      }
+
+      @Override
+      public CacheHelper getCoreCacheHelper() {
+        return null;
+      }
+
+      @Override
+      public CacheHelper getReaderCacheHelper() {
+        return null;
+      }
+    };
+  }
+
+  public void testSingleSlice() {
+    LeafReader largeSegmentReader = dummyIndexReader(50_000);
+    LeafReader firstMediumSegmentReader = dummyIndexReader(30_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(30__000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(30_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(largeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 1);
+
+    final LeafReaderContext[] leaves = resultSlices[0].leaves;
+
+    assertTrue(leaves.length == 4);
+  }
+
+  public void testSmallSegments() {
+    LeafReader firstMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader fourthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader fifthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader sixthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader seventhLargeSegmentReader = dummyIndexReader(130_000);
+    LeafReader eigthLargeSegmentReader = dummyIndexReader(130_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(fourthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(fifthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(sixthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(seventhLargeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(eigthLargeSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 3);
+
+    final LeafReaderContext[] firstSliceleaves = resultSlices[0].leaves;
+    final LeafReaderContext[] secondSliceleaves = resultSlices[1].leaves;
+    final LeafReaderContext[] thirdSliceleaves = resultSlices[2].leaves;
+
+    assertTrue(firstSliceleaves.length == 2);
+    assertTrue(secondSliceleaves.length == 5);
+    assertTrue(thirdSliceleaves.length == 1);
+  }
+
+  public void testLargeSlices() {
+    LeafReader largeSegmentReader = dummyIndexReader(290_900);
+    LeafReader firstMediumSegmentReader = dummyIndexReader(170_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(170_000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(170_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(largeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 3);
+
+    final LeafReaderContext[] firstSliceleaves = resultSlices[0].leaves;
+    final LeafReaderContext[] secondSliceleaves = resultSlices[1].leaves;
+    final LeafReaderContext[] thirdSliceleaves = resultSlices[2].leaves;
+
+    assertTrue(firstSliceleaves.length == 1);
+    assertTrue(secondSliceleaves.length == 2);
+    assertTrue(thirdSliceleaves.length == 1);
+  }
+
+  public void testIntraSliceDocIDOrder() throws Exception {
+    Directory dir = newDirectory();
+    RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+    w.addDocument(new Document());
+    w.addDocument(new Document());
+    w.commit();
+    w.addDocument(new Document());
+    w.addDocument(new Document());
+    w.commit();
+    IndexReader r = w.getReader();
+    w.close();
+
+    ExecutorService service = new ThreadPoolExecutor(4, 4, 0L, TimeUnit.MILLISECONDS,
+        new LinkedBlockingQueue<Runnable>(),
+        new NamedThreadFactory("TestSegmentToThreadMapping"));
+    IndexSearcher s = new IndexSearcher(r, service);
+    Query query = new MatchAllDocsQuery();
+
+    s.search(query, Integer.MAX_VALUE);
+
+    IndexSearcher.LeafSlice[] slices = s.getSlices();
+    assertNotNull(slices);
+
+    for (IndexSearcher.LeafSlice leafSlice : slices) {
+      LeafReaderContext[] leafReaderContexts = leafSlice.leaves;
+      int previousDocBase = leafReaderContexts[0].docBase;
+
+      for (LeafReaderContext leafReaderContext : leafReaderContexts) {
+        assertTrue(previousDocBase <= leafReaderContext.docBase);
+        previousDocBase = leafReaderContext.docBase;
+      }
+    }
+
+    service.shutdown();
+    IOUtils.close(r, dir);
+  }
+
+  public void testRandom() {
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+    int max = 500_000;
+    int min = 10_000;
+    int numSegments = random().nextInt(50);
+
+    for (int i = 0; i < numSegments; i++) {
+      leafReaderContexts.add(new LeafReaderContext(dummyIndexReader(random().nextInt((max - min) + 1) + min)));
+    }
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length > 0);
+  }
+}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
index 125d681298..fb7991d254 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
@@ -26,6 +26,7 @@ import org.apache.lucene.index.LeafReaderContext;
 class AssertingCollector extends FilterCollector {
 
   private int maxDoc = -1;
+  private int previousLeafMaxDoc = -1;
 
   /** Wrap the given collector in order to add assertions. */
   public static Collector wrap(Collector in) {
@@ -50,6 +51,10 @@ class AssertingCollector extends FilterCollector {
         // not only per segment
         assert docBase + doc >= maxDoc : "collection is not in order: current doc="
             + (docBase + doc) + " while " + maxDoc + " has already been collected";
+
+        assert context.docBase >= previousLeafMaxDoc;
+        previousLeafMaxDoc = context.docBase + context.reader().maxDoc();
+
         super.collect(doc);
         maxDoc = docBase + doc;
       }
diff --git a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
index 6aa6f03753..92d364d04b 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
@@ -1932,6 +1932,15 @@ public abstract class LuceneTestCase extends Assert {
         ret = random.nextBoolean()
             ? new AssertingIndexSearcher(random, r, ex)
             : new AssertingIndexSearcher(random, r.getContext(), ex);
+      } else if (random.nextBoolean()) {
+        int maxDocPerSlice = 1 + random.nextInt(100000);
+        int maxSegmentsPerSlice = 1 + random.nextInt(20);
+        ret = new IndexSearcher(r, ex) {
+          @Override
+          protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+            return slices(leaves, maxDocPerSlice, maxSegmentsPerSlice);
+          }
+        };
       } else {
         ret = random.nextBoolean()
             ? new IndexSearcher(r, ex)
-- 
2.17.2 (Apple Git-113)

[asfimport]

Simon Willnauer (@s1monw) (migrated from JIRA)

Thanks @atris, can you bring back the javadocs for 

protected LeafSlice[] slices(List<LeafReaderContext> leaves)

please don't reassign an argument like here:

leaves = new ArrayList<>(leaves);

The rest of the patch looks OK to me yet I am not so sure about the defaults. I do wonder if we should look at this from a different perspective. Rather than using hard numbers can we try to evenly balance the total number of documents across N threads and make N the variable? @mikemccand WDYT?

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

@s1monw Attached is an updated patch.

My two cents are that segregating segments to keep the document count fair is a more complex operation that what the slices API does today (and in this patch). Fair segmentation is a known hard problem (integer partitioning, for eg).

 

We should also consider how much of a bootstrap time latency would a more complex algorithm add. Given that a user has the option of overriding IndexSearcher to add their own ways of splicing, I feel our default algorithm should do well on the common usecase, but not more than that.

 

Happy to discuss the alternatives.

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

Hi @s1monw,

 

Spending a bit more time thinking about your suggestions, I agree that it is a great idea, albeit requiring more thought and effort than what this Jira proposes to achieve.

I have opened #9838 - Cost Based Slice Allocation Algorithm for discussing the same. Please share your thoughts.

 

Would it make sense to push this patch, and then let users consume it and provide feedback while we iterate on the more sophisticated version? We could even have both of the methods available as options to users, potentially

 

Thoughts?

 

[asfimport]

Simon Willnauer (@s1monw) (migrated from JIRA)

> Would it make sense to push this patch, and then let users consume it and provide feedback while we iterate on the more sophisticated version? We could even have both of the methods available as options to users, potentially

I don't think we should push this if we already know we wanna do something different. That said, I am not convinced the numbers are good defaults. At the same time I don't have any numbers here do you have anything to back these defaults up?

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

I don't think we should push this if we already know we wanna do something different. That said, I am not convinced the numbers are good defaults. At the same time I don't have any numbers here do you have anything to back these defaults up?

 

Sure. The reason I was suggesting pushing this patch per se is because the other approach we are advancing would require a couple of new semantics to be introduced, so we could potentially want users to have an option to opt-in for either of the two. That said, I believe the cost based algorithm would also require some hard defaults to be present – to ensure that small segments do not get independent threads even if system had the capacity.

 

RE: The default constant values, these numbers are derived from empirical testing across different datasets in ESRally (nyc_taxis, logging) and looking at the default segment size distribution of wikipedia10M dataset in luceneutil. However, this might not be a good default size to split on.

 

One thing we could do (albeit expensive) is to take the mean number of documents in the corresponding LeafReaderContexts for a query as the split point. Would that be a better dynamic way?

[asfimport]

Simon Willnauer (@s1monw) (migrated from JIRA)

@atris actually I thought about these defaults again and I am starting to think it's an ok default. The reason for this is that we try to prevent having dedicated threads for smallish segments so we group them together. I still do wonder if we need to have 2 parameters? Wouldn't it be enough to just say that we group things together until we have at least 250k docs per thread to be searched? is it really necessary to have another parameter that limits the number of segmetns per slice? I think a single parameter would be great and simpler. WDYT?

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

@s1monw The reasoning behind adding the second parameter was to ensure that we do not bias against the case where there are a large number of small segments. For eg, if there are 100 segments and all of them are small, then we should still allow parallel searches to get some performance gains. Although this should be a rare case since merging will coalesce them.

 

However, I agree with you that this might be contradicting the whole idea of adding the 250K docs split point. If all segments together in an index do not add up to 250K, then the index is small enough to not need parallelism.

 

Attached is an updated patch

 

LUCENE-8757.patch

LUCENE-8757.patch

From a181fa0e6a9630e7f38760cb16827c4a80ce0d86 Mon Sep 17 00:00:00 2001
From: Atri Sharma <atri@apache.org>
Date: Mon, 22 Apr 2019 19:44:52 +0530
Subject: [PATCH] LUCENE-8757: Improving Default Segments To Thread Mapping
 Algorithm

The current slicing algorithm assigns a thread per segment, which
can be detrimental to performance in case the distribution has
a large number of small segments. The patch introduces a slicing
algorithm which coalesces smaller segments to a single thread,
thus reducing the impact of context switching by limiting the
number of threads
---
 .../apache/lucene/search/IndexSearcher.java   |  68 ++++-
 .../index/TestSegmentToThreadMapping.java     | 268 ++++++++++++++++++
 .../lucene/search/AssertingCollector.java     |   5 +
 .../apache/lucene/util/LuceneTestCase.java    |   9 +
 4 files changed, 342 insertions(+), 8 deletions(-)
 create mode 100644 lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java

diff --git a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
index 689409f5cc..1748dd4821 100644
--- a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
+++ b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
@@ -22,6 +22,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -100,6 +101,13 @@ public class IndexSearcher {
    */
   private static final int TOTAL_HITS_THRESHOLD = 1000;
 
+  /**
+   * Thresholds for index slice allocation logic. To change the default, extend
+   * <code> IndexSearcher</code> and use custom values
+   */
+  private static final int MAX_DOCS_PER_SLICE = 250_000;
+  private static final int MAX_SEGMENTS_PER_SLICE = 5;
+
   final IndexReader reader; // package private for testing!
   
   // NOTE: these members might change in incompatible ways
@@ -268,17 +276,60 @@ public class IndexSearcher {
 
   /**
    * Expert: Creates an array of leaf slices each holding a subset of the given leaves.
-   * Each {@link LeafSlice} is executed in a single thread. By default there
-   * will be one {@link LeafSlice} per leaf ({@link org.apache.lucene.index.LeafReaderContext}).
+   * Each {@link LeafSlice} is executed in a single thread. By default, segments with more than
+   * MAX_DOCS_PER_SLICE will get their own thread
    */
   protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
-    LeafSlice[] slices = new LeafSlice[leaves.size()];
-    for (int i = 0; i < slices.length; i++) {
-      slices[i] = new LeafSlice(leaves.get(i));
+    return slices(leaves, MAX_DOCS_PER_SLICE, MAX_SEGMENTS_PER_SLICE);
+  }
+
+  /**
+   * Static method to segregate LeafReaderContexts amongst multiple slices
+   */
+  public static LeafSlice[] slices (List<LeafReaderContext> leaves, int maxDocsPerSlice,
+                                    int maxSegmentsPerSlice) {
+    // 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())));
+
+    final List<List<LeafReaderContext>> groupedLeaves = new ArrayList<>();
+    long docSum = 0;
+    List<LeafReaderContext> group = null;
+    for (LeafReaderContext ctx : sortedLeaves) {
+      if (ctx.reader().maxDoc() > maxDocsPerSlice) {
+        assert group == null;
+        groupedLeaves.add(Collections.singletonList(ctx));
+      } else {
+        if (group == null) {
+          group = new ArrayList<>();
+          group.add(ctx);
+
+          groupedLeaves.add(group);
+        } else {
+          group.add(ctx);
+        }
+
+        docSum += ctx.reader().maxDoc();
+        if (group.size() >= maxSegmentsPerSlice || docSum > maxDocsPerSlice) {
+          group = null;
+          docSum = 0;
+        }
+      }
+    }
+
+    LeafSlice[] slices = new LeafSlice[groupedLeaves.size()];
+    int upto = 0;
+    for (List<LeafReaderContext> currentLeaf : groupedLeaves) {
+      slices[upto] = new LeafSlice(currentLeaf);
+      ++upto;
     }
+
     return slices;
   }
-  
+
   /** Return the {@link IndexReader} this searches. */
   public IndexReader getIndexReader() {
     return reader;
@@ -743,8 +794,9 @@ public class IndexSearcher {
      *  @lucene.experimental */
     public final LeafReaderContext[] leaves;
     
-    public LeafSlice(LeafReaderContext... leaves) {
-      this.leaves = leaves;
+    public LeafSlice(List<LeafReaderContext> leavesList) {
+      Collections.sort(leavesList, Comparator.comparingInt(l -> l.docBase));
+      this.leaves = leavesList.toArray(new LeafReaderContext[0]);
     }
   }
 
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
new file mode 100644
index 0000000000..62299dfd91
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
@@ -0,0 +1,268 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.LuceneTestCase;
+import org.apache.lucene.util.NamedThreadFactory;
+import org.apache.lucene.util.Version;
+
+public class TestSegmentToThreadMapping extends LuceneTestCase {
+
+  public LeafReader dummyIndexReader(final int maxDoc) {
+    return new LeafReader() {
+      @Override
+      public int maxDoc() {
+        return maxDoc;
+      }
+
+      @Override
+      public int numDocs() {
+        return maxDoc;
+      }
+
+      @Override
+      public FieldInfos getFieldInfos() {
+        return FieldInfos.EMPTY;
+      }
+
+      @Override
+      public Bits getLiveDocs() {
+        return null;
+      }
+
+      @Override
+      public Terms terms(String field) throws IOException {
+        return null;
+      }
+
+      @Override
+      public Fields getTermVectors(int doc) {
+        return null;
+      }
+
+      @Override
+      public NumericDocValues getNumericDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public BinaryDocValues getBinaryDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedDocValues getSortedDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedNumericDocValues getSortedNumericDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedSetDocValues getSortedSetDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public NumericDocValues getNormValues(String field) {
+        return null;
+      }
+
+      @Override
+      public PointValues getPointValues(String field) {
+        return null;
+      }
+
+      @Override
+      protected void doClose() {
+      }
+
+      @Override
+      public void document(int doc, StoredFieldVisitor visitor) {
+      }
+
+      @Override
+      public void checkIntegrity() throws IOException {
+      }
+
+      @Override
+      public LeafMetaData getMetaData() {
+        return new LeafMetaData(Version.LATEST.major, Version.LATEST, null);
+      }
+
+      @Override
+      public CacheHelper getCoreCacheHelper() {
+        return null;
+      }
+
+      @Override
+      public CacheHelper getReaderCacheHelper() {
+        return null;
+      }
+    };
+  }
+
+  public void testSingleSlice() {
+    LeafReader largeSegmentReader = dummyIndexReader(50_000);
+    LeafReader firstMediumSegmentReader = dummyIndexReader(30_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(30__000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(30_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(largeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 1);
+
+    final LeafReaderContext[] leaves = resultSlices[0].leaves;
+
+    assertTrue(leaves.length == 4);
+  }
+
+  public void testSmallSegments() {
+    LeafReader firstMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader fourthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader fifthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader sixthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader seventhLargeSegmentReader = dummyIndexReader(130_000);
+    LeafReader eigthLargeSegmentReader = dummyIndexReader(130_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(fourthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(fifthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(sixthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(seventhLargeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(eigthLargeSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 3);
+
+    final LeafReaderContext[] firstSliceleaves = resultSlices[0].leaves;
+    final LeafReaderContext[] secondSliceleaves = resultSlices[1].leaves;
+    final LeafReaderContext[] thirdSliceleaves = resultSlices[2].leaves;
+
+    assertTrue(firstSliceleaves.length == 2);
+    assertTrue(secondSliceleaves.length == 5);
+    assertTrue(thirdSliceleaves.length == 1);
+  }
+
+  public void testLargeSlices() {
+    LeafReader largeSegmentReader = dummyIndexReader(290_900);
+    LeafReader firstMediumSegmentReader = dummyIndexReader(170_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(170_000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(170_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(largeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 3);
+
+    final LeafReaderContext[] firstSliceleaves = resultSlices[0].leaves;
+    final LeafReaderContext[] secondSliceleaves = resultSlices[1].leaves;
+    final LeafReaderContext[] thirdSliceleaves = resultSlices[2].leaves;
+
+    assertTrue(firstSliceleaves.length == 1);
+    assertTrue(secondSliceleaves.length == 2);
+    assertTrue(thirdSliceleaves.length == 1);
+  }
+
+  public void testIntraSliceDocIDOrder() throws Exception {
+    Directory dir = newDirectory();
+    RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+    w.addDocument(new Document());
+    w.addDocument(new Document());
+    w.commit();
+    w.addDocument(new Document());
+    w.addDocument(new Document());
+    w.commit();
+    IndexReader r = w.getReader();
+    w.close();
+
+    ExecutorService service = new ThreadPoolExecutor(4, 4, 0L, TimeUnit.MILLISECONDS,
+        new LinkedBlockingQueue<Runnable>(),
+        new NamedThreadFactory("TestSegmentToThreadMapping"));
+    IndexSearcher s = new IndexSearcher(r, service);
+    Query query = new MatchAllDocsQuery();
+
+    s.search(query, Integer.MAX_VALUE);
+
+    IndexSearcher.LeafSlice[] slices = s.getSlices();
+    assertNotNull(slices);
+
+    for (IndexSearcher.LeafSlice leafSlice : slices) {
+      LeafReaderContext[] leafReaderContexts = leafSlice.leaves;
+      int previousDocBase = leafReaderContexts[0].docBase;
+
+      for (LeafReaderContext leafReaderContext : leafReaderContexts) {
+        assertTrue(previousDocBase <= leafReaderContext.docBase);
+        previousDocBase = leafReaderContext.docBase;
+      }
+    }
+
+    service.shutdown();
+    IOUtils.close(r, dir);
+  }
+
+  public void testRandom() {
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+    int max = 500_000;
+    int min = 10_000;
+    int numSegments = random().nextInt(50);
+
+    for (int i = 0; i < numSegments; i++) {
+      leafReaderContexts.add(new LeafReaderContext(dummyIndexReader(random().nextInt((max - min) + 1) + min)));
+    }
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length > 0);
+  }
+}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
index 125d681298..fb7991d254 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
@@ -26,6 +26,7 @@ import org.apache.lucene.index.LeafReaderContext;
 class AssertingCollector extends FilterCollector {
 
   private int maxDoc = -1;
+  private int previousLeafMaxDoc = -1;
 
   /** Wrap the given collector in order to add assertions. */
   public static Collector wrap(Collector in) {
@@ -50,6 +51,10 @@ class AssertingCollector extends FilterCollector {
         // not only per segment
         assert docBase + doc >= maxDoc : "collection is not in order: current doc="
             + (docBase + doc) + " while " + maxDoc + " has already been collected";
+
+        assert context.docBase >= previousLeafMaxDoc;
+        previousLeafMaxDoc = context.docBase + context.reader().maxDoc();
+
         super.collect(doc);
         maxDoc = docBase + doc;
       }
diff --git a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
index 6aa6f03753..92d364d04b 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
@@ -1932,6 +1932,15 @@ public abstract class LuceneTestCase extends Assert {
         ret = random.nextBoolean()
             ? new AssertingIndexSearcher(random, r, ex)
             : new AssertingIndexSearcher(random, r.getContext(), ex);
+      } else if (random.nextBoolean()) {
+        int maxDocPerSlice = 1 + random.nextInt(100000);
+        int maxSegmentsPerSlice = 1 + random.nextInt(20);
+        ret = new IndexSearcher(r, ex) {
+          @Override
+          protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+            return slices(leaves, maxDocPerSlice, maxSegmentsPerSlice);
+          }
+        };
       } else {
         ret = random.nextBoolean()
             ? new IndexSearcher(r, ex)
-- 
2.17.2 (Apple Git-113)

[asfimport]

Simon Willnauer (@s1monw) (migrated from JIRA)

@atris I think the assertion in this part doesn't hold:

+    for (LeafReaderContext ctx : sortedLeaves) {
+      if (ctx.reader().maxDoc() > maxDocsPerSlice) {
+        assert group == null;
+        List singleSegmentSlice = new ArrayList();

if the previous segment was smallish then group is non-null? I think you should test these cases, maybe add a random test and randomize the order or the segments?

This:

+        List singleSegmentSlice = new ArrayList();
+
+        singleSegmentSlice.add(ctx);
+        groupedLeaves.add(singleSegmentSlice);

can and should be replaced by:

groupedLeaves.add(Collections.singletonList(ctx));

otherwise it looks good.

[asfimport]

Atri Sharma (@atris) (migrated from JIRA)

Hi @s1monw

if the previous segment was smallish then group is non-null? I think you should test these cases, maybe add a random test and randomize the order or the segments?

I dont think that case is possible, since we sort LeafReaderContexts based on the number of documents per segment in descending order. Hence, no LeafReaderContext can be succeded by one which has more documents than its predecessor. I agree with your thought of having a random test with variety of configurations for segment size distributions.

bq.can and should be replaced by:

Fixed, thanks.

LUCENE-8757.patch

 

LUCENE-8757.patch

From a181fa0e6a9630e7f38760cb16827c4a80ce0d86 Mon Sep 17 00:00:00 2001
From: Atri Sharma <atri@apache.org>
Date: Mon, 22 Apr 2019 19:44:52 +0530
Subject: [PATCH] LUCENE-8757: Improving Default Segments To Thread Mapping
 Algorithm

The current slicing algorithm assigns a thread per segment, which
can be detrimental to performance in case the distribution has
a large number of small segments. The patch introduces a slicing
algorithm which coalesces smaller segments to a single thread,
thus reducing the impact of context switching by limiting the
number of threads
---
 .../apache/lucene/search/IndexSearcher.java   |  68 ++++-
 .../index/TestSegmentToThreadMapping.java     | 268 ++++++++++++++++++
 .../lucene/search/AssertingCollector.java     |   5 +
 .../apache/lucene/util/LuceneTestCase.java    |   9 +
 4 files changed, 342 insertions(+), 8 deletions(-)
 create mode 100644 lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java

diff --git a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
index 689409f5cc..1748dd4821 100644
--- a/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
+++ b/lucene/core/src/java/org/apache/lucene/search/IndexSearcher.java
@@ -22,6 +22,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -100,6 +101,13 @@ public class IndexSearcher {
    */
   private static final int TOTAL_HITS_THRESHOLD = 1000;
 
+  /**
+   * Thresholds for index slice allocation logic. To change the default, extend
+   * <code> IndexSearcher</code> and use custom values
+   */
+  private static final int MAX_DOCS_PER_SLICE = 250_000;
+  private static final int MAX_SEGMENTS_PER_SLICE = 5;
+
   final IndexReader reader; // package private for testing!
   
   // NOTE: these members might change in incompatible ways
@@ -268,17 +276,60 @@ public class IndexSearcher {
 
   /**
    * Expert: Creates an array of leaf slices each holding a subset of the given leaves.
-   * Each {@link LeafSlice} is executed in a single thread. By default there
-   * will be one {@link LeafSlice} per leaf ({@link org.apache.lucene.index.LeafReaderContext}).
+   * Each {@link LeafSlice} is executed in a single thread. By default, segments with more than
+   * MAX_DOCS_PER_SLICE will get their own thread
    */
   protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
-    LeafSlice[] slices = new LeafSlice[leaves.size()];
-    for (int i = 0; i < slices.length; i++) {
-      slices[i] = new LeafSlice(leaves.get(i));
+    return slices(leaves, MAX_DOCS_PER_SLICE, MAX_SEGMENTS_PER_SLICE);
+  }
+
+  /**
+   * Static method to segregate LeafReaderContexts amongst multiple slices
+   */
+  public static LeafSlice[] slices (List<LeafReaderContext> leaves, int maxDocsPerSlice,
+                                    int maxSegmentsPerSlice) {
+    // 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())));
+
+    final List<List<LeafReaderContext>> groupedLeaves = new ArrayList<>();
+    long docSum = 0;
+    List<LeafReaderContext> group = null;
+    for (LeafReaderContext ctx : sortedLeaves) {
+      if (ctx.reader().maxDoc() > maxDocsPerSlice) {
+        assert group == null;
+        groupedLeaves.add(Collections.singletonList(ctx));
+      } else {
+        if (group == null) {
+          group = new ArrayList<>();
+          group.add(ctx);
+
+          groupedLeaves.add(group);
+        } else {
+          group.add(ctx);
+        }
+
+        docSum += ctx.reader().maxDoc();
+        if (group.size() >= maxSegmentsPerSlice || docSum > maxDocsPerSlice) {
+          group = null;
+          docSum = 0;
+        }
+      }
+    }
+
+    LeafSlice[] slices = new LeafSlice[groupedLeaves.size()];
+    int upto = 0;
+    for (List<LeafReaderContext> currentLeaf : groupedLeaves) {
+      slices[upto] = new LeafSlice(currentLeaf);
+      ++upto;
     }
+
     return slices;
   }
-  
+
   /** Return the {@link IndexReader} this searches. */
   public IndexReader getIndexReader() {
     return reader;
@@ -743,8 +794,9 @@ public class IndexSearcher {
      *  @lucene.experimental */
     public final LeafReaderContext[] leaves;
     
-    public LeafSlice(LeafReaderContext... leaves) {
-      this.leaves = leaves;
+    public LeafSlice(List<LeafReaderContext> leavesList) {
+      Collections.sort(leavesList, Comparator.comparingInt(l -> l.docBase));
+      this.leaves = leavesList.toArray(new LeafReaderContext[0]);
     }
   }
 
diff --git a/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
new file mode 100644
index 0000000000..62299dfd91
--- /dev/null
+++ b/lucene/core/src/test/org/apache/lucene/index/TestSegmentToThreadMapping.java
@@ -0,0 +1,268 @@
+/*
+ * 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.index;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+
+
+import org.apache.lucene.document.Document;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.IOUtils;
+import org.apache.lucene.util.LuceneTestCase;
+import org.apache.lucene.util.NamedThreadFactory;
+import org.apache.lucene.util.Version;
+
+public class TestSegmentToThreadMapping extends LuceneTestCase {
+
+  public LeafReader dummyIndexReader(final int maxDoc) {
+    return new LeafReader() {
+      @Override
+      public int maxDoc() {
+        return maxDoc;
+      }
+
+      @Override
+      public int numDocs() {
+        return maxDoc;
+      }
+
+      @Override
+      public FieldInfos getFieldInfos() {
+        return FieldInfos.EMPTY;
+      }
+
+      @Override
+      public Bits getLiveDocs() {
+        return null;
+      }
+
+      @Override
+      public Terms terms(String field) throws IOException {
+        return null;
+      }
+
+      @Override
+      public Fields getTermVectors(int doc) {
+        return null;
+      }
+
+      @Override
+      public NumericDocValues getNumericDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public BinaryDocValues getBinaryDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedDocValues getSortedDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedNumericDocValues getSortedNumericDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public SortedSetDocValues getSortedSetDocValues(String field) {
+        return null;
+      }
+
+      @Override
+      public NumericDocValues getNormValues(String field) {
+        return null;
+      }
+
+      @Override
+      public PointValues getPointValues(String field) {
+        return null;
+      }
+
+      @Override
+      protected void doClose() {
+      }
+
+      @Override
+      public void document(int doc, StoredFieldVisitor visitor) {
+      }
+
+      @Override
+      public void checkIntegrity() throws IOException {
+      }
+
+      @Override
+      public LeafMetaData getMetaData() {
+        return new LeafMetaData(Version.LATEST.major, Version.LATEST, null);
+      }
+
+      @Override
+      public CacheHelper getCoreCacheHelper() {
+        return null;
+      }
+
+      @Override
+      public CacheHelper getReaderCacheHelper() {
+        return null;
+      }
+    };
+  }
+
+  public void testSingleSlice() {
+    LeafReader largeSegmentReader = dummyIndexReader(50_000);
+    LeafReader firstMediumSegmentReader = dummyIndexReader(30_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(30__000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(30_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(largeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 1);
+
+    final LeafReaderContext[] leaves = resultSlices[0].leaves;
+
+    assertTrue(leaves.length == 4);
+  }
+
+  public void testSmallSegments() {
+    LeafReader firstMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader fourthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader fifthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader sixthMediumSegmentReader = dummyIndexReader(10_000);
+    LeafReader seventhLargeSegmentReader = dummyIndexReader(130_000);
+    LeafReader eigthLargeSegmentReader = dummyIndexReader(130_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(fourthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(fifthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(sixthMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(seventhLargeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(eigthLargeSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 3);
+
+    final LeafReaderContext[] firstSliceleaves = resultSlices[0].leaves;
+    final LeafReaderContext[] secondSliceleaves = resultSlices[1].leaves;
+    final LeafReaderContext[] thirdSliceleaves = resultSlices[2].leaves;
+
+    assertTrue(firstSliceleaves.length == 2);
+    assertTrue(secondSliceleaves.length == 5);
+    assertTrue(thirdSliceleaves.length == 1);
+  }
+
+  public void testLargeSlices() {
+    LeafReader largeSegmentReader = dummyIndexReader(290_900);
+    LeafReader firstMediumSegmentReader = dummyIndexReader(170_000);
+    LeafReader secondMediumSegmentReader = dummyIndexReader(170_000);
+    LeafReader thirdMediumSegmentReader = dummyIndexReader(170_000);
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+
+    leafReaderContexts.add(new LeafReaderContext(largeSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(firstMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(secondMediumSegmentReader));
+    leafReaderContexts.add(new LeafReaderContext(thirdMediumSegmentReader));
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length == 3);
+
+    final LeafReaderContext[] firstSliceleaves = resultSlices[0].leaves;
+    final LeafReaderContext[] secondSliceleaves = resultSlices[1].leaves;
+    final LeafReaderContext[] thirdSliceleaves = resultSlices[2].leaves;
+
+    assertTrue(firstSliceleaves.length == 1);
+    assertTrue(secondSliceleaves.length == 2);
+    assertTrue(thirdSliceleaves.length == 1);
+  }
+
+  public void testIntraSliceDocIDOrder() throws Exception {
+    Directory dir = newDirectory();
+    RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+    w.addDocument(new Document());
+    w.addDocument(new Document());
+    w.commit();
+    w.addDocument(new Document());
+    w.addDocument(new Document());
+    w.commit();
+    IndexReader r = w.getReader();
+    w.close();
+
+    ExecutorService service = new ThreadPoolExecutor(4, 4, 0L, TimeUnit.MILLISECONDS,
+        new LinkedBlockingQueue<Runnable>(),
+        new NamedThreadFactory("TestSegmentToThreadMapping"));
+    IndexSearcher s = new IndexSearcher(r, service);
+    Query query = new MatchAllDocsQuery();
+
+    s.search(query, Integer.MAX_VALUE);
+
+    IndexSearcher.LeafSlice[] slices = s.getSlices();
+    assertNotNull(slices);
+
+    for (IndexSearcher.LeafSlice leafSlice : slices) {
+      LeafReaderContext[] leafReaderContexts = leafSlice.leaves;
+      int previousDocBase = leafReaderContexts[0].docBase;
+
+      for (LeafReaderContext leafReaderContext : leafReaderContexts) {
+        assertTrue(previousDocBase <= leafReaderContext.docBase);
+        previousDocBase = leafReaderContext.docBase;
+      }
+    }
+
+    service.shutdown();
+    IOUtils.close(r, dir);
+  }
+
+  public void testRandom() {
+    List<LeafReaderContext> leafReaderContexts = new ArrayList<>();
+    int max = 500_000;
+    int min = 10_000;
+    int numSegments = random().nextInt(50);
+
+    for (int i = 0; i < numSegments; i++) {
+      leafReaderContexts.add(new LeafReaderContext(dummyIndexReader(random().nextInt((max - min) + 1) + min)));
+    }
+
+    IndexSearcher.LeafSlice[] resultSlices = IndexSearcher.slices(leafReaderContexts, 250_000, 5);
+
+    assertTrue(resultSlices.length > 0);
+  }
+}
diff --git a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
index 125d681298..fb7991d254 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/search/AssertingCollector.java
@@ -26,6 +26,7 @@ import org.apache.lucene.index.LeafReaderContext;
 class AssertingCollector extends FilterCollector {
 
   private int maxDoc = -1;
+  private int previousLeafMaxDoc = -1;
 
   /** Wrap the given collector in order to add assertions. */
   public static Collector wrap(Collector in) {
@@ -50,6 +51,10 @@ class AssertingCollector extends FilterCollector {
         // not only per segment
         assert docBase + doc >= maxDoc : "collection is not in order: current doc="
             + (docBase + doc) + " while " + maxDoc + " has already been collected";
+
+        assert context.docBase >= previousLeafMaxDoc;
+        previousLeafMaxDoc = context.docBase + context.reader().maxDoc();
+
         super.collect(doc);
         maxDoc = docBase + doc;
       }
diff --git a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
index 6aa6f03753..92d364d04b 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
@@ -1932,6 +1932,15 @@ public abstract class LuceneTestCase extends Assert {
         ret = random.nextBoolean()
             ? new AssertingIndexSearcher(random, r, ex)
             : new AssertingIndexSearcher(random, r.getContext(), ex);
+      } else if (random.nextBoolean()) {
+        int maxDocPerSlice = 1 + random.nextInt(100000);
+        int maxSegmentsPerSlice = 1 + random.nextInt(20);
+        ret = new IndexSearcher(r, ex) {
+          @Override
+          protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+            return slices(leaves, maxDocPerSlice, maxSegmentsPerSlice);
+          }
+        };
       } else {
         ret = random.nextBoolean()
             ? new IndexSearcher(r, ex)
-- 
2.17.2 (Apple Git-113)

[asfimport]

Simon Willnauer (@s1monw) (migrated from JIRA)

Thanks for the additional iteration, now that we simplified this can we remove the sorting? I don't necessearily see how the sort makes things simpler. If we see a segment > threshold we can just add it as a group? I though you did that already and hence my comment about the assertion. WDYT?

I also want to suggest to beef up testing a bit with a randomized version of this like this:

diff --git a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
index 7c63a817adb..76ccca64ee7 100644
--- a/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
+++ b/lucene/test-framework/src/java/org/apache/lucene/util/LuceneTestCase.java
`@@` -1933,6 +1933,14 `@@` public abstract class LuceneTestCase extends Assert {
         ret = random.nextBoolean()
             ? new AssertingIndexSearcher(random, r, ex)
             : new AssertingIndexSearcher(random, r.getContext(), ex);
+      } else if (random.nextBoolean()) {
+        int maxDocPerSlice = 1 + random.nextInt(100000);
+        ret = new IndexSearcher(r, ex) {
+          `@Override`
+          protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+            return slices(leaves, maxDocPerSlice);
+          }
+        };
       } else {
         ret = random.nextBoolean()
             ? new IndexSearcher(r, ex)