Skip to content

Commit 86be0f0

Browse files
not-napoleonmartijnvgelasticmachine
authored
Limit the resutls objects SearchContext creates (#94405)
This restricts the results objects created by DefaultSearchContext to only those that will be used in that phase. In most paths, this means creating only one results object, but we also have an optimization where we return the query and fetch results at the same time, via the QueryFetchSearchResult object, and in that case we need to create both a QuerySearchResult and a FetchSearchResult. --------- Co-authored-by: Martijn van Groningen <[email protected]> Co-authored-by: Elastic Machine <[email protected]>
1 parent ce8714a commit 86be0f0

File tree

9 files changed

+141
-8
lines changed

9 files changed

+141
-8
lines changed

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

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.apache.lucene.search.FieldDoc;
1616
import org.apache.lucene.search.MatchNoDocsQuery;
1717
import org.apache.lucene.search.Query;
18+
import org.apache.lucene.search.TotalHits;
1819
import org.elasticsearch.action.search.SearchShardTask;
1920
import org.elasticsearch.action.search.SearchType;
2021
import org.elasticsearch.common.lucene.search.Queries;
@@ -74,9 +75,9 @@ final class DefaultSearchContext extends SearchContext {
7475
private final IndexShard indexShard;
7576
private final IndexService indexService;
7677
private final ContextIndexSearcher searcher;
77-
private final DfsSearchResult dfsResult;
78-
private final QuerySearchResult queryResult;
79-
private final FetchSearchResult fetchResult;
78+
private DfsSearchResult dfsResult;
79+
private QuerySearchResult queryResult;
80+
private FetchSearchResult fetchResult;
8081
private final float queryBoost;
8182
private final boolean lowLevelCancellation;
8283
private TimeValue timeout;
@@ -142,9 +143,6 @@ final class DefaultSearchContext extends SearchContext {
142143
this.fetchPhase = fetchPhase;
143144
this.searchType = request.searchType();
144145
this.shardTarget = shardTarget;
145-
this.dfsResult = new DfsSearchResult(readerContext.id(), shardTarget, request);
146-
this.queryResult = new QuerySearchResult(readerContext.id(), shardTarget, request);
147-
this.fetchResult = new FetchSearchResult(readerContext.id(), shardTarget);
148146
this.indexService = readerContext.indexService();
149147
this.indexShard = readerContext.indexShard();
150148

@@ -172,6 +170,21 @@ final class DefaultSearchContext extends SearchContext {
172170
this.lowLevelCancellation = lowLevelCancellation;
173171
}
174172

173+
@Override
174+
public void addFetchResult() {
175+
this.fetchResult = new FetchSearchResult(this.readerContext.id(), this.shardTarget);
176+
}
177+
178+
@Override
179+
public void addQueryResult() {
180+
this.queryResult = new QuerySearchResult(this.readerContext.id(), this.shardTarget, this.request);
181+
}
182+
183+
@Override
184+
public void addDfsResult() {
185+
this.dfsResult = new DfsSearchResult(this.readerContext.id(), this.shardTarget, this.request);
186+
}
187+
175188
/**
176189
* Should be called before executing the main query and after all other parameters have been set.
177190
*/
@@ -700,6 +713,22 @@ public QuerySearchResult queryResult() {
700713
return queryResult;
701714
}
702715

716+
@Override
717+
public TotalHits getTotalHits() {
718+
if (queryResult != null) {
719+
return queryResult.getTotalHits();
720+
}
721+
return null;
722+
}
723+
724+
@Override
725+
public float getMaxScore() {
726+
if (queryResult != null) {
727+
return queryResult.getMaxScore();
728+
}
729+
return Float.NaN;
730+
}
731+
703732
@Override
704733
public FetchPhase fetchPhase() {
705734
return fetchPhase;

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,7 @@ private DfsSearchResult executeDfsPhase(ShardSearchRequest request, SearchShardT
449449
Releasable ignored = readerContext.markAsUsed(getKeepAlive(request));
450450
SearchContext context = createContext(readerContext, request, task, true)
451451
) {
452+
context.addDfsResult();
452453
dfsPhase.execute(context);
453454
return context.dfsResult();
454455
} catch (Exception e) {
@@ -630,6 +631,7 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, SearchSh
630631
SearchContext context = createContext(readerContext, request, task, true)
631632
) {
632633
tracer.startTrace("executeQueryPhase", Map.of());
634+
context.addQueryResult();
633635
final long afterQueryTime;
634636
try (SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(context)) {
635637
loadOrExecuteQueryPhase(request, context);
@@ -641,6 +643,7 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, SearchSh
641643
tracer.stopTrace();
642644
}
643645
if (request.numberOfShards() == 1) {
646+
context.addFetchResult();
644647
return executeFetchPhase(readerContext, context, afterQueryTime);
645648
} else {
646649
// Pass the rescoreDocIds to the queryResult to send them the coordinating node and receive them back in the fetch phase.
@@ -698,6 +701,7 @@ public void executeQueryPhase(
698701
SearchContext searchContext = createContext(readerContext, shardSearchRequest, task, false);
699702
SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(searchContext)
700703
) {
704+
searchContext.addQueryResult();
701705
searchContext.searcher().setAggregatedDfs(readerContext.getAggregatedDfs(null));
702706
processScroll(request, readerContext, searchContext);
703707
QueryPhase.execute(searchContext);
@@ -722,6 +726,7 @@ public void executeQueryPhase(QuerySearchRequest request, SearchShardTask task,
722726
SearchContext searchContext = createContext(readerContext, shardSearchRequest, task, true);
723727
SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(searchContext)
724728
) {
729+
searchContext.addQueryResult();
725730
searchContext.searcher().setAggregatedDfs(request.dfs());
726731
QueryPhase.execute(searchContext);
727732
if (searchContext.queryResult().hasSearchContext() == false && readerContext.singleSession()) {
@@ -777,9 +782,11 @@ public void executeFetchPhase(
777782
SearchContext searchContext = createContext(readerContext, shardSearchRequest, task, false);
778783
SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(searchContext)
779784
) {
785+
searchContext.addFetchResult();
780786
searchContext.assignRescoreDocIds(readerContext.getRescoreDocIds(null));
781787
searchContext.searcher().setAggregatedDfs(readerContext.getAggregatedDfs(null));
782788
processScroll(request, readerContext, searchContext);
789+
searchContext.addQueryResult();
783790
QueryPhase.execute(searchContext);
784791
final long afterQueryTime = executor.success();
785792
QueryFetchSearchResult fetchSearchResult = executeFetchPhase(readerContext, searchContext, afterQueryTime);
@@ -799,6 +806,7 @@ public void executeFetchPhase(ShardFetchRequest request, SearchShardTask task, A
799806
final Releasable markAsUsed = readerContext.markAsUsed(getKeepAlive(shardSearchRequest));
800807
runAsync(getExecutor(readerContext.indexShard()), () -> {
801808
try (SearchContext searchContext = createContext(readerContext, shardSearchRequest, task, false)) {
809+
searchContext.addFetchResult();
802810
if (request.lastEmittedDoc() != null) {
803811
searchContext.scrollContext().lastEmittedDoc = request.lastEmittedDoc();
804812
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -164,8 +164,8 @@ protected SearchHit nextDoc(int doc) throws IOException {
164164
throw new TaskCancelledException("cancelled");
165165
}
166166

167-
TotalHits totalHits = context.queryResult().getTotalHits();
168-
return new SearchHits(hits, totalHits, context.queryResult().getMaxScore());
167+
TotalHits totalHits = context.getTotalHits();
168+
return new SearchHits(hits, totalHits, context.getMaxScore());
169169
}
170170

171171
private static StoredFieldLoader buildStoredFieldsLoader(StoredFieldsSpec spec) {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ public final class QueryFetchSearchResult extends SearchPhaseResult {
2424

2525
public QueryFetchSearchResult(StreamInput in) throws IOException {
2626
super(in);
27+
// TODO: Delegate refcounting to QuerySearchResult (see https://github.com/elastic/elasticsearch/pull/94023)
2728
queryResult = new QuerySearchResult(in);
2829
fetchResult = new FetchSearchResult(in);
2930
}

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.apache.lucene.search.Collector;
1212
import org.apache.lucene.search.FieldDoc;
1313
import org.apache.lucene.search.Query;
14+
import org.apache.lucene.search.TotalHits;
1415
import org.elasticsearch.action.search.SearchShardTask;
1516
import org.elasticsearch.action.search.SearchType;
1617
import org.elasticsearch.core.TimeValue;
@@ -368,16 +369,41 @@ public DfsSearchResult dfsResult() {
368369
return in.dfsResult();
369370
}
370371

372+
@Override
373+
public void addDfsResult() {
374+
in.addDfsResult();
375+
}
376+
371377
@Override
372378
public QuerySearchResult queryResult() {
373379
return in.queryResult();
374380
}
375381

382+
@Override
383+
public void addQueryResult() {
384+
in.addQueryResult();
385+
}
386+
387+
@Override
388+
public TotalHits getTotalHits() {
389+
return in.getTotalHits();
390+
}
391+
392+
@Override
393+
public float getMaxScore() {
394+
return in.getMaxScore();
395+
}
396+
376397
@Override
377398
public FetchSearchResult fetchResult() {
378399
return in.fetchResult();
379400
}
380401

402+
@Override
403+
public void addFetchResult() {
404+
in.addFetchResult();
405+
}
406+
381407
@Override
382408
public FetchPhase fetchPhase() {
383409
return in.fetchPhase();

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.apache.lucene.search.Collector;
1111
import org.apache.lucene.search.FieldDoc;
1212
import org.apache.lucene.search.Query;
13+
import org.apache.lucene.search.TotalHits;
1314
import org.elasticsearch.action.search.SearchShardTask;
1415
import org.elasticsearch.action.search.SearchType;
1516
import org.elasticsearch.core.Nullable;
@@ -315,12 +316,34 @@ public final Query rewrittenQuery() {
315316

316317
public abstract DfsSearchResult dfsResult();
317318

319+
/**
320+
* Indicates that the caller will be using, and thus owning, a {@link DfsSearchResult} object. It is the caller's responsibility
321+
* to correctly cleanup this result object.
322+
*/
323+
public abstract void addDfsResult();
324+
318325
public abstract QuerySearchResult queryResult();
319326

327+
/**
328+
* Indicates that the caller will be using, and thus owning, a {@link QuerySearchResult} object. It is the caller's responsibility
329+
* to correctly cleanup this result object.
330+
*/
331+
public abstract void addQueryResult();
332+
333+
public abstract TotalHits getTotalHits();
334+
335+
public abstract float getMaxScore();
336+
320337
public abstract FetchPhase fetchPhase();
321338

322339
public abstract FetchSearchResult fetchResult();
323340

341+
/**
342+
* Indicates that the caller will be using, and thus owning, a {@link FetchSearchResult} object. It is the caller's responsibility
343+
* to correctly cleanup this result object.
344+
*/
345+
public abstract void addFetchResult();
346+
324347
/**
325348
* Return a handle over the profilers for the current search request, or {@code null} if profiling is not enabled.
326349
*/

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
package org.elasticsearch.search.internal;
99

1010
import org.apache.lucene.search.Query;
11+
import org.apache.lucene.search.TotalHits;
1112
import org.elasticsearch.core.TimeValue;
1213
import org.elasticsearch.index.query.ParsedQuery;
1314
import org.elasticsearch.search.aggregations.SearchContextAggregations;
@@ -308,4 +309,14 @@ public FetchSearchResult fetchResult() {
308309
public long getRelativeTimeInMillis() {
309310
throw new UnsupportedOperationException("Not supported");
310311
}
312+
313+
@Override
314+
public TotalHits getTotalHits() {
315+
return querySearchResult.getTotalHits();
316+
}
317+
318+
@Override
319+
public float getMaxScore() {
320+
return querySearchResult.getMaxScore();
321+
}
311322
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1232,8 +1232,17 @@ public void testCreateSearchContext() throws IOException {
12321232
assertEquals(expectedIndexName, searchShardTarget.getFullyQualifiedIndexName());
12331233
assertEquals(clusterAlias, searchShardTarget.getClusterAlias());
12341234
assertEquals(shardId, searchShardTarget.getShardId());
1235+
1236+
assertNull(searchContext.dfsResult());
1237+
searchContext.addDfsResult();
12351238
assertSame(searchShardTarget, searchContext.dfsResult().getSearchShardTarget());
1239+
1240+
assertNull(searchContext.queryResult());
1241+
searchContext.addQueryResult();
12361242
assertSame(searchShardTarget, searchContext.queryResult().getSearchShardTarget());
1243+
1244+
assertNull(searchContext.fetchResult());
1245+
searchContext.addFetchResult();
12371246
assertSame(searchShardTarget, searchContext.fetchResult().getSearchShardTarget());
12381247
}
12391248
}

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.apache.lucene.search.Collector;
1111
import org.apache.lucene.search.FieldDoc;
1212
import org.apache.lucene.search.Query;
13+
import org.apache.lucene.search.TotalHits;
1314
import org.elasticsearch.action.search.SearchShardTask;
1415
import org.elasticsearch.action.search.SearchType;
1516
import org.elasticsearch.core.TimeValue;
@@ -473,16 +474,41 @@ public DfsSearchResult dfsResult() {
473474
return null;
474475
}
475476

477+
@Override
478+
public void addDfsResult() {
479+
// this space intentionally left blank
480+
}
481+
476482
@Override
477483
public QuerySearchResult queryResult() {
478484
return queryResult;
479485
}
480486

487+
@Override
488+
public void addQueryResult() {
489+
// this space intentionally left blank
490+
}
491+
492+
@Override
493+
public TotalHits getTotalHits() {
494+
return queryResult.getTotalHits();
495+
}
496+
497+
@Override
498+
public float getMaxScore() {
499+
return queryResult.getMaxScore();
500+
}
501+
481502
@Override
482503
public FetchSearchResult fetchResult() {
483504
return null;
484505
}
485506

507+
@Override
508+
public void addFetchResult() {
509+
// this space intentionally left blank
510+
}
511+
486512
@Override
487513
public FetchPhase fetchPhase() {
488514
return null;

0 commit comments

Comments
 (0)