Skip to content

Commit a09729c

Browse files
jimczijimczi
authored
Don't always rewrite the Lucene query in search phases (#79358) (#79367)
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.
1 parent ff6e589 commit a09729c

File tree

16 files changed

+70
-101
lines changed

16 files changed

+70
-101
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ protected ShardValidateQueryResponse shardOperation(ShardValidateQueryRequest re
191191
try {
192192
ParsedQuery parsedQuery = searchContext.getSearchExecutionContext().toQuery(request.query());
193193
searchContext.parsedQuery(parsedQuery);
194-
searchContext.preProcess(request.rewrite());
194+
searchContext.preProcess();
195195
valid = true;
196196
explanation = explain(searchContext, request.rewrite());
197197
} catch (QueryShardException|ParsingException e) {
@@ -208,7 +208,7 @@ protected ShardValidateQueryResponse shardOperation(ShardValidateQueryRequest re
208208
}
209209

210210
private String explain(SearchContext context, boolean rewritten) {
211-
Query query = context.query();
211+
Query query = rewritten ? context.rewrittenQuery() : context.query();
212212
if (rewritten && query instanceof MatchNoDocsQuery) {
213213
return context.parsedQuery().query().toString();
214214
} else {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,9 +117,9 @@ protected ExplainResponse shardOperation(ExplainRequest request, ShardId shardId
117117
return new ExplainResponse(shardId.getIndexName(), request.type(), request.id(), false);
118118
}
119119
context.parsedQuery(context.getSearchExecutionContext().toQuery(request.query()));
120-
context.preProcess(true);
120+
context.preProcess();
121121
int topLevelDocId = result.docIdAndVersion().docId + result.docIdAndVersion().docBase;
122-
Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
122+
Explanation explanation = context.searcher().explain(context.rewrittenQuery(), topLevelDocId);
123123
for (RescoreContext ctx : context.rescore()) {
124124
Rescorer rescorer = ctx.rescorer();
125125
explanation = rescorer.explain(topLevelDocId, context.searcher(), ctx, explanation);

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

Lines changed: 10 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,6 @@
5050
import org.elasticsearch.search.internal.ShardSearchContextId;
5151
import org.elasticsearch.search.internal.ShardSearchRequest;
5252
import org.elasticsearch.search.profile.Profilers;
53-
import org.elasticsearch.search.query.QueryPhaseExecutionException;
5453
import org.elasticsearch.search.query.QuerySearchResult;
5554
import org.elasticsearch.search.rescore.RescoreContext;
5655
import org.elasticsearch.search.slice.SliceBuilder;
@@ -176,7 +175,7 @@ final class DefaultSearchContext extends SearchContext {
176175
* Should be called before executing the main query and after all other parameters have been set.
177176
*/
178177
@Override
179-
public void preProcess(boolean rewrite) {
178+
public void preProcess() {
180179
if (hasOnlySuggest() ) {
181180
return;
182181
}
@@ -231,19 +230,20 @@ public void preProcess(boolean rewrite) {
231230
throw new UncheckedIOException(e);
232231
}
233232

234-
if (query() == null) {
233+
if (query == null) {
235234
parsedQuery(ParsedQuery.parsedMatchAllQuery());
236235
}
237236
if (queryBoost != AbstractQueryBuilder.DEFAULT_BOOST) {
238-
parsedQuery(new ParsedQuery(new BoostQuery(query(), queryBoost), parsedQuery()));
237+
parsedQuery(new ParsedQuery(new BoostQuery(query, queryBoost), parsedQuery()));
239238
}
240239
this.query = buildFilteredQuery(query);
241-
if (rewrite) {
242-
try {
243-
this.query = searcher.rewrite(query);
244-
} catch (IOException e) {
245-
throw new QueryPhaseExecutionException(shardTarget, "Failed to rewrite main query", e);
246-
}
240+
if (lowLevelCancellation) {
241+
searcher().addQueryCancellation(() -> {
242+
final SearchShardTask task = getTask();
243+
if (task != null) {
244+
task.ensureNotCancelled();
245+
}
246+
});
247247
}
248248
}
249249

@@ -586,9 +586,6 @@ public ParsedQuery parsedQuery() {
586586
return this.originalQuery;
587587
}
588588

589-
/**
590-
* The query to execute, in its rewritten form.
591-
*/
592589
@Override
593590
public Query query() {
594591
return this.query;

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -894,8 +894,7 @@ protected SearchContext createContext(ReaderContext readerContext,
894894
}
895895
context.setTask(task);
896896

897-
// pre process
898-
queryPhase.preProcess(context);
897+
context.preProcess();
899898
} catch (Exception e) {
900899
context.close();
901900
throw e;
@@ -1111,7 +1110,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
11111110
* the filter for nested documents or slicing so we have to
11121111
* delay reading it until the aggs ask for it.
11131112
*/
1114-
() -> context.query() == null ? new MatchAllDocsQuery() : context.query(),
1113+
() -> context.rewrittenQuery() == null ? new MatchAllDocsQuery() : context.rewrittenQuery(),
11151114
context.getProfilers() == null ? null : context.getProfilers().getAggregationProfiler(),
11161115
multiBucketConsumerService.create(),
11171116
() -> new SubSearchContext(context).parsedQuery(context.parsedQuery()).fetchFieldsContext(context.fetchFieldsContext()),

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public CollectionStatistics collectionStatistics(String field) throws IOExceptio
6060
}
6161
};
6262

63-
searcher.createWeight(context.searcher().rewrite(context.query()), ScoreMode.COMPLETE, 1);
63+
searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
6464
for (RescoreContext rescoreContext : context.rescore()) {
6565
for (Query query : rescoreContext.getQueries()) {
6666
searcher.createWeight(context.searcher().rewrite(query), ScoreMode.COMPLETE, 1);

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,12 +67,19 @@ public SearchLookup searchLookup() {
6767
}
6868

6969
/**
70-
* The original query
70+
* The original query, not rewritten.
7171
*/
7272
public Query query() {
7373
return searchContext.query();
7474
}
7575

76+
/**
77+
* The original query in its rewritten form.
78+
*/
79+
public Query rewrittenQuery() {
80+
return searchContext.rewrittenQuery();
81+
}
82+
7683
/**
7784
* The original query with additional filters and named queries
7885
*/

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public void setNextReader(LeafReaderContext readerContext) {
3434
@Override
3535
public void process(HitContext hitContext) throws IOException {
3636
final int topLevelDocId = hitContext.hit().docId();
37-
Explanation explanation = context.searcher().explain(context.query(), topLevelDocId);
37+
Explanation explanation = context.searcher().explain(context.rewrittenQuery(), topLevelDocId);
3838

3939
for (RescoreContext rescore : context.rescore()) {
4040
explanation = rescore.rescorer().explain(topLevelDocId, context.searcher(), rescore, explanation);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public FetchSubPhaseProcessor getProcessor(FetchContext context) throws IOExcept
2727
return null;
2828
}
2929
final IndexSearcher searcher = context.searcher();
30-
final Weight weight = searcher.createWeight(searcher.rewrite(context.query()), ScoreMode.COMPLETE, 1);
30+
final Weight weight = searcher.createWeight(context.rewrittenQuery(), ScoreMode.COMPLETE, 1);
3131
return new FetchSubPhaseProcessor() {
3232

3333
Scorer scorer;

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ public SearchContext storedFieldsContext(StoredFieldsContext storedFieldsContext
6363
}
6464

6565
@Override
66-
public void preProcess(boolean rewrite) {
67-
in.preProcess(rewrite);
66+
public void preProcess() {
67+
in.preProcess();
6868
}
6969

7070
@Override

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

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import org.elasticsearch.core.TimeValue;
2020
import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
2121
import org.elasticsearch.index.query.ParsedQuery;
22+
import org.elasticsearch.index.query.QueryShardException;
2223
import org.elasticsearch.index.query.SearchExecutionContext;
2324
import org.elasticsearch.index.shard.IndexShard;
2425
import org.elasticsearch.search.RescoreDocIds;
@@ -42,6 +43,7 @@
4243
import org.elasticsearch.search.sort.SortAndFormats;
4344
import org.elasticsearch.search.suggest.SuggestionSearchContext;
4445

46+
import java.io.IOException;
4547
import java.util.HashMap;
4648
import java.util.List;
4749
import java.util.Map;
@@ -65,6 +67,9 @@ public abstract class SearchContext implements Releasable {
6567
private final AtomicBoolean closed = new AtomicBoolean(false);
6668
private InnerHitsContext innerHitsContext;
6769

70+
private boolean rewritten;
71+
private Query rewriteQuery;
72+
6873
protected SearchContext() {}
6974

7075
public abstract void setTask(SearchShardTask task);
@@ -82,9 +87,8 @@ public final void close() {
8287

8388
/**
8489
* Should be called before executing the main query and after all other parameters have been set.
85-
* @param rewrite if the set query should be rewritten against the searcher returned from {@link #searcher()}
8690
*/
87-
public abstract void preProcess(boolean rewrite);
91+
public abstract void preProcess();
8892

8993
/** Automatically apply all required filters to the given query such as
9094
* alias filters, types filters, etc. */
@@ -252,10 +256,27 @@ public final void assignRescoreDocIds(RescoreDocIds rescoreDocIds) {
252256
public abstract ParsedQuery parsedQuery();
253257

254258
/**
255-
* The query to execute, might be rewritten.
259+
* The query to execute, not rewritten.
256260
*/
257261
public abstract Query query();
258262

263+
/**
264+
* The query to execute in its rewritten form.
265+
*/
266+
public final Query rewrittenQuery() {
267+
if (query() == null) {
268+
throw new IllegalStateException("preProcess must be called first");
269+
}
270+
if (rewriteQuery == null) {
271+
try {
272+
this.rewriteQuery = searcher().rewrite(query());
273+
} catch (IOException exc) {
274+
throw new QueryShardException(getSearchExecutionContext(), "rewrite failed", exc);
275+
}
276+
}
277+
return rewriteQuery;
278+
}
279+
259280
public abstract int from();
260281

261282
public abstract SearchContext from(int from);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public SubSearchContext(SearchContext context) {
6262
}
6363

6464
@Override
65-
public void preProcess(boolean rewrite) {
65+
public void preProcess() {
6666
}
6767

6868
@Override

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

Lines changed: 2 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import org.apache.lucene.search.SortField;
2727
import org.apache.lucene.search.TopDocs;
2828
import org.apache.lucene.search.TotalHits;
29-
import org.elasticsearch.action.search.SearchShardTask;
3029
import org.elasticsearch.common.lucene.Lucene;
3130
import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
3231
import org.elasticsearch.common.util.concurrent.QueueResizingEsThreadPoolExecutor;
@@ -58,7 +57,6 @@
5857
import static org.elasticsearch.search.query.QueryCollectorContext.createMultiCollectorContext;
5958
import static org.elasticsearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext;
6059

61-
6260
/**
6361
* Query phase of a search request, used to run the query and get back from each shard information about the matching documents
6462
* (document ids and score or sort criteria) so that matches can be reduced on the coordinating node
@@ -78,27 +76,6 @@ public QueryPhase() {
7876
this.rescorePhase = new RescorePhase();
7977
}
8078

81-
public void preProcess(SearchContext context) {
82-
final Runnable cancellation;
83-
if (context.lowLevelCancellation()) {
84-
cancellation = context.searcher().addQueryCancellation(() -> {
85-
SearchShardTask task = context.getTask();
86-
if (task != null) {
87-
task.ensureNotCancelled();
88-
}
89-
});
90-
} else {
91-
cancellation = null;
92-
}
93-
try {
94-
context.preProcess(true);
95-
} finally {
96-
if (cancellation != null) {
97-
context.searcher().removeQueryCancellation(cancellation);
98-
}
99-
}
100-
}
101-
10279
public void execute(SearchContext searchContext) throws QueryPhaseExecutionException {
10380
if (searchContext.hasOnlySuggest()) {
10481
suggestPhase.execute(searchContext);
@@ -113,7 +90,7 @@ public void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep
11390
}
11491

11592
// Pre-process aggregations as late as possible. In the case of a DFS_Q_T_F
116-
// request, preProcess is called on the DFS phase phase, this is why we pre-process them
93+
// request, preProcess is called on the DFS phase, this is why we pre-process them
11794
// here to make sure it happens during the QUERY phase
11895
aggregationPhase.preProcess(searchContext);
11996
boolean rescore = executeInternal(searchContext);
@@ -142,7 +119,7 @@ static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExe
142119
try {
143120
queryResult.from(searchContext.from());
144121
queryResult.size(searchContext.size());
145-
Query query = searchContext.query();
122+
Query query = searchContext.rewrittenQuery();
146123
assert query == searcher.rewrite(query); // already rewritten
147124

148125
final ScrollContext scrollContext = searchContext.scrollContext();
@@ -230,15 +207,6 @@ static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExe
230207
timeoutRunnable = null;
231208
}
232209

233-
if (searchContext.lowLevelCancellation()) {
234-
searcher.addQueryCancellation(() -> {
235-
SearchShardTask task = searchContext.getTask();
236-
if (task != null) {
237-
task.ensureNotCancelled();
238-
}
239-
});
240-
}
241-
242210
try {
243211
boolean shouldRescore = searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, timeoutSet);
244212
ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -417,7 +417,7 @@ static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOExcep
417417
static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searchContext,
418418
boolean hasFilterCollector) throws IOException {
419419
final IndexReader reader = searchContext.searcher().getIndexReader();
420-
final Query query = searchContext.query();
420+
final Query query = searchContext.rewrittenQuery();
421421
// top collectors don't like a size of 0
422422
final int totalNumDocs = Math.max(1, reader.numDocs());
423423
if (searchContext.size() == 0) {

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

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ protected Engine.Searcher acquireSearcherInternal(String source) {
138138
contextWithoutScroll.close();
139139

140140
// resultWindow greater than maxResultWindow and scrollContext is null
141-
IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> contextWithoutScroll.preProcess(false));
141+
IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> contextWithoutScroll.preProcess());
142142
assertThat(exception.getMessage(), equalTo("Result window is too large, from + size must be less than or equal to:"
143143
+ " [" + maxResultWindow + "] but was [310]. See the scroll api for a more efficient way to request large data sets. "
144144
+ "This limit can be set by changing the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey()
@@ -151,7 +151,7 @@ protected Engine.Searcher acquireSearcherInternal(String source) {
151151
DefaultSearchContext context1 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null,
152152
timeout, null, false, Version.CURRENT);
153153
context1.from(300);
154-
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
154+
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess());
155155
assertThat(exception.getMessage(), equalTo("Batch size is too large, size must be less than or equal to: ["
156156
+ maxResultWindow + "] but was [310]. Scroll batch sizes cost as much memory as result windows so they are "
157157
+ "controlled by the [" + IndexSettings.MAX_RESULT_WINDOW_SETTING.getKey() + "] index level setting."));
@@ -166,12 +166,12 @@ protected Engine.Searcher acquireSearcherInternal(String source) {
166166
when(rescoreContext.getWindowSize()).thenReturn(500);
167167
context1.addRescore(rescoreContext);
168168

169-
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
169+
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess());
170170
assertThat(exception.getMessage(), equalTo("Cannot use [sort] option in conjunction with [rescore]."));
171171

172172
// rescore is null but sort is not null and rescoreContext.getWindowSize() exceeds maxResultWindow
173173
context1.sort(null);
174-
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess(false));
174+
exception = expectThrows(IllegalArgumentException.class, () -> context1.preProcess());
175175

176176
assertThat(exception.getMessage(), equalTo("Rescore window [" + rescoreContext.getWindowSize() + "] is too large. "
177177
+ "It must be less than [" + maxRescoreWindow + "]. This prevents allocating massive heaps for storing the results "
@@ -197,7 +197,7 @@ public ScrollContext scrollContext() {
197197
when(sliceBuilder.getMax()).thenReturn(numSlices);
198198
context2.sliceBuilder(sliceBuilder);
199199

200-
exception = expectThrows(IllegalArgumentException.class, () -> context2.preProcess(false));
200+
exception = expectThrows(IllegalArgumentException.class, () -> context2.preProcess());
201201
assertThat(exception.getMessage(), equalTo("The number of slices [" + numSlices + "] is too large. It must "
202202
+ "be less than [" + maxSlicesPerScroll + "]. This limit can be set by changing the [" +
203203
IndexSettings.MAX_SLICES_PER_SCROLL.getKey() + "] index level setting."));
@@ -209,7 +209,7 @@ public ScrollContext scrollContext() {
209209
DefaultSearchContext context3 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null,
210210
timeout, null, false, Version.CURRENT);
211211
ParsedQuery parsedQuery = ParsedQuery.parsedMatchAllQuery();
212-
context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess(false);
212+
context3.sliceBuilder(null).parsedQuery(parsedQuery).preProcess();
213213
assertEquals(context3.query(), context3.buildFilteredQuery(parsedQuery.query()));
214214

215215
when(searchExecutionContext.getIndexSettings()).thenReturn(indexSettings);
@@ -220,9 +220,9 @@ public ScrollContext scrollContext() {
220220
searcherSupplier.get(), randomNonNegativeLong(), false);
221221
DefaultSearchContext context4 = new DefaultSearchContext(readerContext, shardSearchRequest, target, null,
222222
timeout, null, false, Version.CURRENT);
223-
context4.sliceBuilder(new SliceBuilder(1,2)).parsedQuery(parsedQuery).preProcess(false);
223+
context4.sliceBuilder(new SliceBuilder(1,2)).parsedQuery(parsedQuery).preProcess();
224224
Query query1 = context4.query();
225-
context4.sliceBuilder(new SliceBuilder(0,2)).parsedQuery(parsedQuery).preProcess(false);
225+
context4.sliceBuilder(new SliceBuilder(0,2)).parsedQuery(parsedQuery).preProcess();
226226
Query query2 = context4.query();
227227
assertTrue(query1 instanceof MatchNoDocsQuery || query2 instanceof MatchNoDocsQuery);
228228

0 commit comments

Comments
 (0)