Don't always rewrite the Lucene query in search phases (#79358)

Today the Lucene query is rewritten eagerly in the dfs, query and fetch search phases.
While this is needed in dfs and search, most of the fetch phases could avoid this cost.
This change adds an explicit accessor for SearchContext that rewrites the original query
once and returns its rewritten form. That allows to rewrite the Lucene query only when
it's required.

server/src/main/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryAction.java

@@ -191,7 +191,7 @@ public class TransportValidateQueryAction extends TransportBroadcastAction<
         try {
             ParsedQuery parsedQuery = searchContext.getSearchExecutionContext().toQuery(request.query());
             searchContext.parsedQuery(parsedQuery);
-            searchContext.preProcess(request.rewrite());
+            searchContext.preProcess();
             valid = true;
             explanation = explain(searchContext, request.rewrite());
         } catch (QueryShardException|ParsingException e) {
@@ -208,7 +208,7 @@ public class TransportValidateQueryAction extends TransportBroadcastAction<
     }
 
     private String explain(SearchContext context, boolean rewritten) {
-        Query query = context.query();
+        Query query = rewritten ? context.rewrittenQuery() : context.query();
         if (rewritten && query instanceof MatchNoDocsQuery) {
             return context.parsedQuery().query().toString();
         } else {

server/src/main/java/org/elasticsearch/action/explain/TransportExplainAction.java

@@ -104,9 +104,9 @@ public class TransportExplainAction extends TransportSingleShardAction<ExplainRe
                 return new ExplainResponse(shardId.getIndexName(), request.id(), false);
             }
             context.parsedQuery(context.getSearchExecutionContext().toQuery(request.query()));
-            context.preProcess(true);
+            context.preProcess();
             int topLevelDocId = result.docIdAndVersion().docId + result.docIdAndVersion().docBase;
-            Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
+            Explanation explanation = context.searcher().explain(context.rewrittenQuery(), topLevelDocId);
             for (RescoreContext ctx : context.rescore()) {
                 Rescorer rescorer = ctx.rescorer();
                 explanation = rescorer.explain(topLevelDocId, context.searcher(), ctx, explanation);

server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java

@@ -48,7 +48,6 @@ import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.profile.Profilers;
-import org.elasticsearch.search.query.QueryPhaseExecutionException;
 import org.elasticsearch.search.query.QuerySearchResult;
 import org.elasticsearch.search.rescore.RescoreContext;
 import org.elasticsearch.search.slice.SliceBuilder;
@@ -169,7 +168,7 @@ final class DefaultSearchContext extends SearchContext {
      * Should be called before executing the main query and after all other parameters have been set.
      */
     @Override
-    public void preProcess(boolean rewrite) {
+    public void preProcess() {
         if (hasOnlySuggest() ) {
             return;
         }
@@ -224,19 +223,20 @@ final class DefaultSearchContext extends SearchContext {
             throw new UncheckedIOException(e);
         }
 
-        if (query() == null) {
+        if (query == null) {
             parsedQuery(ParsedQuery.parsedMatchAllQuery());
         }
         if (queryBoost != AbstractQueryBuilder.DEFAULT_BOOST) {
-            parsedQuery(new ParsedQuery(new BoostQuery(query(), queryBoost), parsedQuery()));
+            parsedQuery(new ParsedQuery(new BoostQuery(query, queryBoost), parsedQuery()));
         }
         this.query = buildFilteredQuery(query);
-        if (rewrite) {
-            try {
-                this.query = searcher.rewrite(query);
-            } catch (IOException e) {
-                throw new QueryPhaseExecutionException(shardTarget, "Failed to rewrite main query", e);
-            }
+        if (lowLevelCancellation) {
+            searcher().addQueryCancellation(() -> {
+                final SearchShardTask task = getTask();
+                if (task != null) {
+                    task.ensureNotCancelled();
+                }
+            });
         }
     }
 
@@ -562,9 +562,6 @@ final class DefaultSearchContext extends SearchContext {
         return this.originalQuery;
     }
 
-    /**
-     * The query to execute, in its rewritten form.
-     */
     @Override
     public Query query() {
         return this.query;

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -879,8 +879,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
             }
             context.setTask(task);
 
-            // pre process
-            queryPhase.preProcess(context);
+            context.preProcess();
         } catch (Exception e) {
             context.close();
             throw e;
@@ -1095,7 +1094,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
                  * the filter for nested documents or slicing so we have to
                  * delay reading it until the aggs ask for it.
                  */
-                () -> context.query() == null ? new MatchAllDocsQuery() : context.query(),
+                () -> context.rewrittenQuery() == null ? new MatchAllDocsQuery() : context.rewrittenQuery(),
                 context.getProfilers() == null ? null : context.getProfilers().getAggregationProfiler(),
                 multiBucketConsumerService.create(),
                 () -> new SubSearchContext(context).parsedQuery(context.parsedQuery()).fetchFieldsContext(context.fetchFieldsContext()),

server/src/main/java/org/elasticsearch/search/dfs/DfsPhase.java

@@ -60,7 +60,7 @@ public class DfsPhase {
                 }
             };
 
-            searcher.createWeight(context.searcher().rewrite(context.query()), ScoreMode.COMPLETE, 1);
+            searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
             for (RescoreContext rescoreContext : context.rescore()) {
                 for (Query query : rescoreContext.getQueries()) {
                     searcher.createWeight(context.searcher().rewrite(query), ScoreMode.COMPLETE, 1);

server/src/main/java/org/elasticsearch/search/fetch/FetchContext.java

@@ -67,12 +67,19 @@ public class FetchContext {
     }
 
     /**
-     * The original query
+     * The original query, not rewritten.
      */
     public Query query() {
         return searchContext.query();
     }
 
+    /**
+     * The original query in its rewritten form.
+     */
+    public Query rewrittenQuery() {
+        return searchContext.rewrittenQuery();
+    }
+
     /**
      * The original query with additional filters and named queries
      */

server/src/main/java/org/elasticsearch/search/fetch/subphase/ExplainPhase.java

@@ -34,7 +34,7 @@ public final class ExplainPhase implements FetchSubPhase {
             @Override
             public void process(HitContext hitContext) throws IOException {
                 final int topLevelDocId = hitContext.hit().docId();
-                Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
+                Explanation explanation = context.searcher().explain(context.rewrittenQuery(), topLevelDocId);
 
                 for (RescoreContext rescore : context.rescore()) {
                     explanation = rescore.rescorer().explain(topLevelDocId, context.searcher(), rescore, explanation);

server/src/main/java/org/elasticsearch/search/fetch/subphase/FetchScorePhase.java

@@ -27,7 +27,7 @@ public class FetchScorePhase implements FetchSubPhase {
             return null;
         }
         final IndexSearcher searcher = context.searcher();
-        final Weight weight = searcher.createWeight(searcher.rewrite(context.query()), ScoreMode.COMPLETE, 1);
+        final Weight weight = searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
         return new FetchSubPhaseProcessor() {
 
             Scorer scorer;

server/src/main/java/org/elasticsearch/search/internal/FilteredSearchContext.java

@@ -63,8 +63,8 @@ public abstract class FilteredSearchContext extends SearchContext {
     }
 
     @Override
-    public void preProcess(boolean rewrite) {
-        in.preProcess(rewrite);
+    public void preProcess() {
+        in.preProcess();
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -19,6 +19,7 @@ import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.query.ParsedQuery;
+import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.search.RescoreDocIds;
@@ -42,6 +43,7 @@ import org.elasticsearch.search.rescore.RescoreContext;
 import org.elasticsearch.search.sort.SortAndFormats;
 import org.elasticsearch.search.suggest.SuggestionSearchContext;
 
+import java.io.IOException;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -65,6 +67,9 @@ public abstract class SearchContext implements Releasable {
     private final AtomicBoolean closed = new AtomicBoolean(false);
     private InnerHitsContext innerHitsContext;
 
+    private boolean rewritten;
+    private Query rewriteQuery;
+
     protected SearchContext() {}
 
     public abstract void setTask(SearchShardTask task);
@@ -82,9 +87,8 @@ public abstract class SearchContext implements Releasable {
 
     /**
      * Should be called before executing the main query and after all other parameters have been set.
-     * @param rewrite if the set query should be rewritten against the searcher returned from {@link #searcher()}
      */
-    public abstract void preProcess(boolean rewrite);
+    public abstract void preProcess();
 
     /** Automatically apply all required filters to the given query such as
      *  alias filters, types filters, etc. */
@@ -252,10 +256,27 @@ public abstract class SearchContext implements Releasable {
     public abstract ParsedQuery parsedQuery();
 
     /**
-     * The query to execute, might be rewritten.
+     * The query to execute, not rewritten.
      */
     public abstract Query query();
 
+    /**
+     * The query to execute in its rewritten form.
+     */
+    public final Query rewrittenQuery() {
+        if (query() == null) {
+            throw new IllegalStateException("preProcess must be called first");
+        }
+        if (rewriteQuery == null) {
+            try {
+                this.rewriteQuery = searcher().rewrite(query());
+            } catch (IOException exc) {
+                throw new QueryShardException(getSearchExecutionContext(), "rewrite failed", exc);
+            }
+        }
+        return rewriteQuery;
+    }
+
     public abstract int from();
 
     public abstract SearchContext from(int from);

server/src/main/java/org/elasticsearch/search/internal/SubSearchContext.java

@@ -62,7 +62,7 @@ public class SubSearchContext extends FilteredSearchContext {
     }
 
     @Override
-    public void preProcess(boolean rewrite) {
+    public void preProcess() {
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -21,7 +21,6 @@ import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TotalHits;
-import org.elasticsearch.action.search.SearchShardTask;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.common.util.concurrent.EWMATrackingEsThreadPoolExecutor;
@@ -50,7 +49,6 @@ import static org.elasticsearch.search.query.QueryCollectorContext.createMinScor
 import static org.elasticsearch.search.query.QueryCollectorContext.createMultiCollectorContext;
 import static org.elasticsearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext;
 
-
 /**
  * Query phase of a search request, used to run the query and get back from each shard information about the matching documents
  * (document ids and score or sort criteria) so that matches can be reduced on the coordinating node
@@ -68,27 +66,6 @@ public class QueryPhase {
         this.rescorePhase = new RescorePhase();
     }
 
-    public void preProcess(SearchContext context) {
-        final Runnable cancellation;
-        if (context.lowLevelCancellation()) {
-            cancellation = context.searcher().addQueryCancellation(() -> {
-                SearchShardTask task = context.getTask();
-                if (task != null) {
-                    task.ensureNotCancelled();
-                }
-            });
-        } else {
-            cancellation = null;
-        }
-        try {
-            context.preProcess(true);
-        } finally {
-            if (cancellation != null) {
-                context.searcher().removeQueryCancellation(cancellation);
-            }
-        }
-    }
-
     public void execute(SearchContext searchContext) throws QueryPhaseExecutionException {
         if (searchContext.hasOnlySuggest()) {
             suggestPhase.execute(searchContext);
@@ -103,7 +80,7 @@ public class QueryPhase {
         }
 
         // Pre-process aggregations as late as possible. In the case of a DFS_Q_T_F
-        // request, preProcess is called on the DFS phase phase, this is why we pre-process them
+        // request, preProcess is called on the DFS phase, this is why we pre-process them
         // here to make sure it happens during the QUERY phase
         aggregationPhase.preProcess(searchContext);
         boolean rescore = executeInternal(searchContext);
@@ -132,7 +109,7 @@ public class QueryPhase {
         try {
             queryResult.from(searchContext.from());
             queryResult.size(searchContext.size());
-            Query query = searchContext.query();
+            Query query = searchContext.rewrittenQuery();
             assert query == searcher.rewrite(query); // already rewritten
 
             final ScrollContext scrollContext = searchContext.scrollContext();
@@ -204,15 +181,6 @@ public class QueryPhase {
                 timeoutRunnable = null;
             }
 
-            if (searchContext.lowLevelCancellation()) {
-                searcher.addQueryCancellation(() -> {
-                    SearchShardTask task = searchContext.getTask();
-                    if (task != null) {
-                        task.ensureNotCancelled();
-                    }
-                });
-            }
-
             try {
                 boolean shouldRescore = searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, timeoutSet);
                 ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH);

server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java

@@ -417,7 +417,7 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
     static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searchContext,
                                                                  boolean hasFilterCollector) throws IOException {
         final IndexReader reader = searchContext.searcher().getIndexReader();
-        final Query query = searchContext.query();
+        final Query query = searchContext.rewrittenQuery();
         // top collectors don't like a size of 0
         final int totalNumDocs = Math.max(1, reader.numDocs());
         if (searchContext.size() == 0) {

server/src/test/java/org/elasticsearch/search/DefaultSearchContextTests.java

@@ -137,7 +137,7 @@ public class DefaultSearchContextTests extends ESTestCase {
             contextWithoutScroll.close();
 
             // resultWindow greater than maxResultWindow and scrollContext is null
-            IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> contextWithoutScroll.preProcess(false));
+            IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> contextWithoutScroll.preProcess());
             assertThat(exception.getMessage(), equalTo("Result window is too large, from + size must be less than or equal to:"
                 + " [" + maxResultWindow + "] but was [310]. See the scroll api for a more efficient way to request large data sets. "
                 + "This limit can be set by changing the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey()
@@ -150,7 +150,7 @@ public class DefaultSearchContextTests extends ESTestCase {
             DefaultSearchContext context1 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null,
                 timeout, null, false);
             context1.from(300);
-            exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
+            exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess());
             assertThat(exception.getMessage(), equalTo("Batch size is too large, size must be less than or equal to: ["
                 + maxResultWindow + "] but was [310]. Scroll batch sizes cost as much memory as result windows so they are "
                 + "controlled by the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey() + "] index level setting."));
@@ -165,12 +165,12 @@ public class DefaultSearchContextTests extends ESTestCase {
             when(rescoreContext.getWindowSize()).thenReturn(500);
             context1.addRescore(rescoreContext);
 
-            exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
+            exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess());
             assertThat(exception.getMessage(), equalTo("Cannot use [sort] option in conjunction with [rescore]."));
 
             // rescore is null but sort is not null and rescoreContext.getWindowSize() exceeds maxResultWindow
             context1.sort(null);
-            exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
+            exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess());
 
             assertThat(exception.getMessage(), equalTo("Rescore window [" + rescoreContext.getWindowSize() + "] is too large. "
                 + "It must be less than [" + maxRescoreWindow + "]. This prevents allocating massive heaps for storing the results "
@@ -196,7 +196,7 @@ public class DefaultSearchContextTests extends ESTestCase {
             when(sliceBuilder.getMax()).thenReturn(numSlices);
             context2.sliceBuilder(sliceBuilder);
 
-            exception = expectThrows(IllegalArgumentException.class, () -> context2.preProcess(false));
+            exception = expectThrows(IllegalArgumentException.class, () -> context2.preProcess());
             assertThat(exception.getMessage(), equalTo("The number of slices [" + numSlices + "] is too large. It must "
                 + "be less than [" + maxSlicesPerScroll + "]. This limit can be set by changing the [" +
                 IndexSettings.MAX_SLICES_PER_SCROLL.getKey() + "] index level setting."));
@@ -208,7 +208,7 @@ public class DefaultSearchContextTests extends ESTestCase {
             DefaultSearchContext context3 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null,
                 timeout, null, false);
             ParsedQuery parsedQuery = ParsedQuery.parsedMatchAllQuery();
-            context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess(false);
+            context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess();
             assertEquals(context3.query(), context3.buildFilteredQuery(parsedQuery.query()));
 
             when(searchExecutionContext.getIndexSettings()).thenReturn(indexSettings);
@@ -219,9 +219,9 @@ public class DefaultSearchContextTests extends ESTestCase {
                 searcherSupplier.get(), randomNonNegativeLong(), false);
             DefaultSearchContext context4 =
                 new DefaultSearchContext(readerContext, shardSearchRequest, target, null, timeout, null, false);
-            context4.sliceBuilder(new SliceBuilder(1,2)).parsedQuery(parsedQuery).preProcess(false);
+            context4.sliceBuilder(new SliceBuilder(1,2)).parsedQuery(parsedQuery).preProcess();
             Query query1 = context4.query();
-            context4.sliceBuilder(new SliceBuilder(0,2)).parsedQuery(parsedQuery).preProcess(false);
+            context4.sliceBuilder(new SliceBuilder(0,2)).parsedQuery(parsedQuery).preProcess();
             Query query2 = context4.query();
             assertTrue(query1 instanceof MatchNoDocsQuery || query2 instanceof MatchNoDocsQuery);
 

server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java

@@ -912,7 +912,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
         dir.close();
     }
 
-    public void testCancellationDuringPreprocess() throws IOException {
+    public void testCancellationDuringRewrite() throws IOException {
         try (Directory dir = newDirectory();
              RandomIndexWriter w = new RandomIndexWriter(random(), dir, newIndexWriterConfig())) {
 
@@ -925,42 +925,19 @@ public class QueryPhaseTests extends IndexShardTestCase {
             w.close();
 
             try (IndexReader reader = DirectoryReader.open(dir)) {
-                TestSearchContext context = new TestSearchContextWithRewriteAndCancellation(
-                    null, indexShard, newContextSearcher(reader));
+                TestSearchContext context = new TestSearchContext(null, indexShard, newContextSearcher(reader));
                 PrefixQuery prefixQuery = new PrefixQuery(new Term("foo", "a"));
                 prefixQuery.setRewriteMethod(MultiTermQuery.SCORING_BOOLEAN_REWRITE);
                 context.parsedQuery(new ParsedQuery(prefixQuery));
                 SearchShardTask task = new SearchShardTask(randomLong(), "transport", "", "", TaskId.EMPTY_TASK_ID, Collections.emptyMap());
                 TaskCancelHelper.cancel(task, "simulated");
                 context.setTask(task);
-                expectThrows(TaskCancelledException.class, () -> new QueryPhase().preProcess(context));
+                context.searcher().addQueryCancellation(task::ensureNotCancelled);
+                expectThrows(TaskCancelledException.class, () -> context.rewrittenQuery());
             }
         }
     }
 
-    private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext {
-
-        private TestSearchContextWithRewriteAndCancellation(SearchExecutionContext searchExecutionContext,
-                                                            IndexShard indexShard,
-                                                            ContextIndexSearcher searcher) {
-            super(searchExecutionContext, indexShard, searcher);
-        }
-
-        @Override
-        public void preProcess(boolean rewrite) {
-            try {
-                searcher().rewrite(query());
-            } catch (IOException e) {
-                fail("IOException shouldn't be thrown");
-            }
-        }
-
-        @Override
-        public boolean lowLevelCancellation() {
-            return true;
-        }
-    }
-
     private static ContextIndexSearcher newContextSearcher(IndexReader reader) throws IOException {
         return new ContextIndexSearcher(reader, IndexSearcher.getDefaultSimilarity(),
             IndexSearcher.getDefaultQueryCache(), IndexSearcher.getDefaultQueryCachingPolicy(), true);

test/framework/src/main/java/org/elasticsearch/test/TestSearchContext.java

@@ -115,7 +115,7 @@ public class TestSearchContext extends SearchContext {
     }
 
     @Override
-    public void preProcess(boolean rewrite) {
+    public void preProcess() {
     }
 
     @Override