From faeb0cb8b4cbb2ec20ce27540e4a7a98c0b075d0 Mon Sep 17 00:00:00 2001
From: Luca Cavanna <javanna@apache.org>
Date: Tue, 24 Oct 2023 21:53:19 +0200
Subject: [PATCH] Don't fork slices creation in ContextIndexSearcher

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.
---
 .../search/internal/ContextIndexSearcher.java | 31 ++++++++-----------
 .../internal/ContextIndexSearcherTests.java   |  4 +--
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java
index 1733029af9024..3c69db98c7588 100644
--- a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java
+++ b/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 ContextIndexSearcher(
         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 Weight createWeight(Query query, ScoreMode scoreMode, float boost) throws
         }
     }
 
-    /**
-     * 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 @@ private <C extends Collector, T> T search(Weight weight, CollectorManager<C, T>
         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));
diff --git a/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java b/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java
index 3f7d67d292761..120982bea36de 100644
--- a/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java
+++ b/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java
@@ -564,7 +564,7 @@ public void testCancelSliceTasksOnException() throws Exception {
                         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 void testCancelSliceTasksOnTimeout() throws Exception {
                         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)];