Don't always rewrite the Lucene query in search phases

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

@@ -191,7 +191,7 @@ protected ShardValidateQueryResponse shardOperation(ShardValidateQueryRequest re
         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 @@ protected ShardValidateQueryResponse shardOperation(ShardValidateQueryRequest re
     }
 
     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 @@ protected ExplainResponse shardOperation(ExplainRequest request, ShardId shardId
                 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.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 @@ public void preProcess(boolean rewrite) {
             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 @@ public ParsedQuery parsedQuery() {
         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 @@ protected SearchContext createContext(ReaderContext readerContext,
             }
             context.setTask(task);
 
-            // pre process
-            queryPhase.preProcess(context);
+            context.preProcess();
         } catch (Exception e) {
             context.close();
             throw e;
@@ -1095,7 +1094,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
                  * 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 CollectionStatistics collectionStatistics(String field) throws IOExceptio
                 }
             };
 
-            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 SearchLookup searchLookup() {
     }
 
     /**
-     * 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 void setNextReader(LeafReaderContext readerContext) {
             @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 FetchSubPhaseProcessor getProcessor(FetchContext context) throws IOExcept
             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 SearchContext storedFieldsContext(StoredFieldsContext storedFieldsContext
     }
 
     @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.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.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 final void close() {
 
     /**
      * 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 final void assignRescoreDocIds(RescoreDocIds rescoreDocIds) {
     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 SubSearchContext(SearchContext context) {
     }
 
     @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.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.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 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 void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep
         }
 
         // 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 @@ static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExe
         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 @@ static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExe
                 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 @@ static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOExcep
     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 @@ protected Engine.Searcher acquireSearcherInternal(String source) {
             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 @@ protected Engine.Searcher acquireSearcherInternal(String source) {
             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 @@ protected Engine.Searcher acquireSearcherInternal(String source) {
             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 ScrollContext scrollContext() {
             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 ScrollContext scrollContext() {
             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 ScrollContext scrollContext() {
                 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);