Browse Source

Don't fork slices creation in ContextIndexSearcher (#101282)

We used to have to fork the slices creation with Lucene 9.7, as slices
were used for knn searches too. With Lucene 9.8, we can customize how
IndexSearcher creates slices, and rely on leafSlices retrieved via
IndexSearcher#getSlices.
Luca Cavanna 2 years ago
parent
commit
41ab9495ba

+ 13 - 18
server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java

@@ -87,7 +87,7 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
     private QueryProfiler profiler;
     private final MutableQueryTimeout cancellable;
 
-    private final LeafSlice[] leafSlices;
+    private final int maximumNumberOfSlices;
     // don't create slices with less than this number of docs
     private final int minimumDocsPerSlice;
 
@@ -150,13 +150,15 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
         setQueryCachingPolicy(queryCachingPolicy);
         this.cancellable = cancellable;
         this.minimumDocsPerSlice = minimumDocsPerSlice;
-        if (executor == null) {
-            this.leafSlices = null;
-        } else {
-            // we offload to the executor unconditionally, including requests that don't support concurrency
-            this.leafSlices = computeSlices(getLeafContexts(), maximumNumberOfSlices, minimumDocsPerSlice);
-            assert this.leafSlices.length <= maximumNumberOfSlices : "more slices created than the maximum allowed";
-        }
+        this.maximumNumberOfSlices = maximumNumberOfSlices;
+    }
+
+    @Override
+    protected LeafSlice[] slices(List<LeafReaderContext> leaves) {
+        // we offload to the executor unconditionally, including requests that don't support concurrency
+        LeafSlice[] leafSlices = computeSlices(getLeafContexts(), maximumNumberOfSlices, minimumDocsPerSlice);
+        assert leafSlices.length <= maximumNumberOfSlices : "more slices created than the maximum allowed";
+        return leafSlices;
     }
 
     // package private for testing
@@ -238,15 +240,6 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
         }
     }
 
-    /**
-     * Returns the slices created by this {@link ContextIndexSearcher}, different from those created by the base class and
-     * returned by {@link IndexSearcher#getSlices()}. The former are used for parallelizing the collection, while the latter are used
-     * for now to parallelize rewrite (e.g. knn query rewrite)
-     */
-    final LeafSlice[] getSlicesForCollection() {
-        return leafSlices;
-    }
-
     /**
      * Each computed slice contains at least 10% of the total data in the leaves with a
      * minimum given by the <code>minDocsPerSlice</code> parameter and the final number
@@ -346,7 +339,9 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
         if (getExecutor() == null) {
             search(leafContexts, weight, firstCollector);
             return collectorManager.reduce(Collections.singletonList(firstCollector));
-        } else if (leafSlices.length == 0) {
+        }
+        LeafSlice[] leafSlices = getSlices();
+        if (leafSlices.length == 0) {
             assert leafContexts.isEmpty();
             doAggregationPostCollection(firstCollector);
             return collectorManager.reduce(Collections.singletonList(firstCollector));

+ 2 - 2
server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java

@@ -564,7 +564,7 @@ public class ContextIndexSearcherTests extends ESTestCase {
                         1
                     )
                 ) {
-                    leafSlices = contextIndexSearcher.getSlicesForCollection();
+                    leafSlices = contextIndexSearcher.getSlices();
                     int numThrowingLeafSlices = randomIntBetween(1, 3);
                     for (int i = 0; i < numThrowingLeafSlices; i++) {
                         LeafSlice throwingLeafSlice = leafSlices[randomIntBetween(0, Math.min(leafSlices.length, numAvailableThreads) - 1)];
@@ -700,7 +700,7 @@ public class ContextIndexSearcherTests extends ESTestCase {
                         1
                     )
                 ) {
-                    leafSlices = contextIndexSearcher.getSlicesForCollection();
+                    leafSlices = contextIndexSearcher.getSlices();
                     int numThrowingLeafSlices = randomIntBetween(1, 3);
                     for (int i = 0; i < numThrowingLeafSlices; i++) {
                         LeafSlice throwingLeafSlice = leafSlices[randomIntBetween(0, Math.min(leafSlices.length, numAvailableThreads) - 1)];