Skip to content

Commit be93df9

Browse files
jimczijimczi
committed
Compute the took time of the query after the expand phase (#24902)
The took time computed for search requests does not take in account the expand search phase. This change delays the computation to after the expand phase finishes. Relates #24900
1 parent 3f02234 commit be93df9

File tree

4 files changed

+37
-41
lines changed

4 files changed

+37
-41
lines changed

‎core/src/main/java/org/elasticsearch/action/search/ExpandSearchPhase.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.elasticsearch.search.SearchHits;
2929
import org.elasticsearch.search.builder.SearchSourceBuilder;
3030
import org.elasticsearch.search.collapse.CollapseBuilder;
31+
import org.elasticsearch.search.internal.InternalSearchResponse;
3132

3233
import java.io.IOException;
3334
import java.util.HashMap;
@@ -42,11 +43,11 @@
4243
*/
4344
final class ExpandSearchPhase extends SearchPhase {
4445
private final SearchPhaseContext context;
45-
private final SearchResponse searchResponse;
46-
private final Function<SearchResponse, SearchPhase> nextPhaseFactory;
46+
private final InternalSearchResponse searchResponse;
47+
private final Function<InternalSearchResponse, SearchPhase> nextPhaseFactory;
4748

48-
ExpandSearchPhase(SearchPhaseContext context, SearchResponse searchResponse,
49-
Function<SearchResponse, SearchPhase> nextPhaseFactory) {
49+
ExpandSearchPhase(SearchPhaseContext context, InternalSearchResponse searchResponse,
50+
Function<InternalSearchResponse, SearchPhase> nextPhaseFactory) {
5051
super("expand");
5152
this.context = context;
5253
this.searchResponse = searchResponse;
@@ -65,15 +66,15 @@ private boolean isCollapseRequest() {
6566

6667
@Override
6768
public void run() throws IOException {
68-
if (isCollapseRequest() && searchResponse.getHits().getHits().length > 0) {
69+
if (isCollapseRequest() && searchResponse.hits().getHits().length > 0) {
6970
SearchRequest searchRequest = context.getRequest();
7071
CollapseBuilder collapseBuilder = searchRequest.source().collapse();
7172
final List<InnerHitBuilder> innerHitBuilders = collapseBuilder.getInnerHits();
7273
MultiSearchRequest multiRequest = new MultiSearchRequest();
7374
if (collapseBuilder.getMaxConcurrentGroupRequests() > 0) {
7475
multiRequest.maxConcurrentSearchRequests(collapseBuilder.getMaxConcurrentGroupRequests());
7576
}
76-
for (SearchHit hit : searchResponse.getHits()) {
77+
for (SearchHit hit : searchResponse.hits().getHits()) {
7778
BoolQueryBuilder groupQuery = new BoolQueryBuilder();
7879
Object collapseValue = hit.field(collapseBuilder.getField()).getValue();
7980
if (collapseValue != null) {
@@ -97,7 +98,7 @@ public void run() throws IOException {
9798
context.getSearchTransport().sendExecuteMultiSearch(multiRequest, context.getTask(),
9899
ActionListener.wrap(response -> {
99100
Iterator<MultiSearchResponse.Item> it = response.iterator();
100-
for (SearchHit hit : searchResponse.getHits()) {
101+
for (SearchHit hit : searchResponse.hits.getHits()) {
101102
for (InnerHitBuilder innerHitBuilder : innerHitBuilders) {
102103
MultiSearchResponse.Item item = it.next();
103104
if (item.isFailure()) {

‎core/src/main/java/org/elasticsearch/action/search/FetchSearchPhase.java‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636

3737
import java.io.IOException;
3838
import java.util.List;
39-
import java.util.function.Function;
39+
import java.util.function.BiFunction;
4040

4141
/**
4242
* This search phase merges the query results from the previous phase together and calculates the topN hits for this search.
@@ -46,7 +46,7 @@ final class FetchSearchPhase extends SearchPhase {
4646
private final AtomicArray<FetchSearchResult> fetchResults;
4747
private final SearchPhaseController searchPhaseController;
4848
private final AtomicArray<SearchPhaseResult> queryResults;
49-
private final Function<SearchResponse, SearchPhase> nextPhaseFactory;
49+
private final BiFunction<InternalSearchResponse, String, SearchPhase> nextPhaseFactory;
5050
private final SearchPhaseContext context;
5151
private final Logger logger;
5252
private final InitialSearchPhase.SearchPhaseResults<SearchPhaseResult> resultConsumer;
@@ -55,13 +55,13 @@ final class FetchSearchPhase extends SearchPhase {
5555
SearchPhaseController searchPhaseController,
5656
SearchPhaseContext context) {
5757
this(resultConsumer, searchPhaseController, context,
58-
(response) -> new ExpandSearchPhase(context, response, // collapse only happens if the request has inner hits
59-
(finalResponse) -> sendResponsePhase(finalResponse, context)));
58+
(response, scrollId) -> new ExpandSearchPhase(context, response, // collapse only happens if the request has inner hits
59+
(finalResponse) -> sendResponsePhase(finalResponse, scrollId, context)));
6060
}
6161

6262
FetchSearchPhase(InitialSearchPhase.SearchPhaseResults<SearchPhaseResult> resultConsumer,
6363
SearchPhaseController searchPhaseController,
64-
SearchPhaseContext context, Function<SearchResponse, SearchPhase> nextPhaseFactory) {
64+
SearchPhaseContext context, BiFunction<InternalSearchResponse, String, SearchPhase> nextPhaseFactory) {
6565
super("fetch");
6666
if (context.getNumShards() != resultConsumer.getNumShards()) {
6767
throw new IllegalStateException("number of shards must match the length of the query results but doesn't:"
@@ -205,14 +205,14 @@ private void moveToNextPhase(SearchPhaseController searchPhaseController,
205205
AtomicArray<? extends SearchPhaseResult> fetchResultsArr) {
206206
final InternalSearchResponse internalResponse = searchPhaseController.merge(context.getRequest().scroll() != null,
207207
reducedQueryPhase, fetchResultsArr.asList(), fetchResultsArr::get);
208-
context.executeNextPhase(this, nextPhaseFactory.apply(context.buildSearchResponse(internalResponse, scrollId)));
208+
context.executeNextPhase(this, nextPhaseFactory.apply(internalResponse, scrollId));
209209
}
210210

211-
private static SearchPhase sendResponsePhase(SearchResponse response, SearchPhaseContext context) {
211+
private static SearchPhase sendResponsePhase(InternalSearchResponse response, String scrollId, SearchPhaseContext context) {
212212
return new SearchPhase("response") {
213213
@Override
214214
public void run() throws IOException {
215-
context.onResponse(response);
215+
context.onResponse(context.buildSearchResponse(response, scrollId));
216216
}
217217
};
218218
}

‎core/src/test/java/org/elasticsearch/action/search/ExpandSearchPhaseTests.java‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -108,13 +108,12 @@ void sendExecuteMultiSearch(MultiSearchRequest request, SearchTask task, ActionL
108108
Collections.singletonMap("someField", new SearchHitField("someField", Collections.singletonList(collapseValue))))},
109109
1, 1.0F);
110110
InternalSearchResponse internalSearchResponse = new InternalSearchResponse(hits, null, null, null, false, null, 1);
111-
SearchResponse response = mockSearchPhaseContext.buildSearchResponse(internalSearchResponse, null);
112111
AtomicReference<SearchResponse> reference = new AtomicReference<>();
113-
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, response, r ->
112+
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, internalSearchResponse, (r) ->
114113
new SearchPhase("test") {
115114
@Override
116115
public void run() throws IOException {
117-
reference.set(r);
116+
reference.set(mockSearchPhaseContext.buildSearchResponse(r, null));
118117
}
119118
}
120119
);
@@ -123,7 +122,6 @@ public void run() throws IOException {
123122
mockSearchPhaseContext.assertNoFailure();
124123
assertNotNull(reference.get());
125124
SearchResponse theResponse = reference.get();
126-
assertSame(theResponse, response);
127125
assertEquals(numInnerHits, theResponse.getHits().getHits()[0].getInnerHits().size());
128126

129127
for (int innerHitNum = 0; innerHitNum < numInnerHits; innerHitNum++) {
@@ -167,13 +165,12 @@ void sendExecuteMultiSearch(MultiSearchRequest request, SearchTask task, ActionL
167165
Collections.singletonMap("someField", new SearchHitField("someField", Collections.singletonList(collapseValue))))}, 1,
168166
1.0F);
169167
InternalSearchResponse internalSearchResponse = new InternalSearchResponse(hits, null, null, null, false, null, 1);
170-
SearchResponse response = mockSearchPhaseContext.buildSearchResponse(internalSearchResponse, null);
171168
AtomicReference<SearchResponse> reference = new AtomicReference<>();
172-
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, response, r ->
169+
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, internalSearchResponse, r ->
173170
new SearchPhase("test") {
174171
@Override
175172
public void run() throws IOException {
176-
reference.set(r);
173+
reference.set(mockSearchPhaseContext.buildSearchResponse(r, null));
177174
}
178175
}
179176
);
@@ -201,13 +198,12 @@ void sendExecuteMultiSearch(MultiSearchRequest request, SearchTask task, ActionL
201198
new SearchHit(2, "ID2", new Text("type"),
202199
Collections.singletonMap("someField", new SearchHitField("someField", Collections.singletonList(null))))}, 1, 1.0F);
203200
InternalSearchResponse internalSearchResponse = new InternalSearchResponse(hits, null, null, null, false, null, 1);
204-
SearchResponse response = mockSearchPhaseContext.buildSearchResponse(internalSearchResponse, null);
205201
AtomicReference<SearchResponse> reference = new AtomicReference<>();
206-
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, response, r ->
202+
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, internalSearchResponse, r ->
207203
new SearchPhase("test") {
208204
@Override
209205
public void run() throws IOException {
210-
reference.set(r);
206+
reference.set(mockSearchPhaseContext.buildSearchResponse(r, null));
211207
}
212208
}
213209
);
@@ -232,13 +228,12 @@ void sendExecuteMultiSearch(MultiSearchRequest request, SearchTask task, ActionL
232228

233229
SearchHits hits = new SearchHits(new SearchHit[0], 1, 1.0f);
234230
InternalSearchResponse internalSearchResponse = new InternalSearchResponse(hits, null, null, null, false, null, 1);
235-
SearchResponse response = mockSearchPhaseContext.buildSearchResponse(internalSearchResponse, null);
236231
AtomicReference<SearchResponse> reference = new AtomicReference<>();
237-
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, response, r ->
232+
ExpandSearchPhase phase = new ExpandSearchPhase(mockSearchPhaseContext, internalSearchResponse, r ->
238233
new SearchPhase("test") {
239234
@Override
240235
public void run() throws IOException {
241-
reference.set(r);
236+
reference.set(mockSearchPhaseContext.buildSearchResponse(r, null));
242237
}
243238
}
244239
);

‎core/src/test/java/org/elasticsearch/action/search/FetchSearchPhaseTests.java‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.elasticsearch.search.fetch.FetchSearchResult;
3333
import org.elasticsearch.search.fetch.QueryFetchSearchResult;
3434
import org.elasticsearch.search.fetch.ShardFetchSearchRequest;
35+
import org.elasticsearch.search.internal.InternalSearchResponse;
3536
import org.elasticsearch.search.query.QuerySearchResult;
3637
import org.elasticsearch.test.ESTestCase;
3738
import org.elasticsearch.transport.Transport;
@@ -66,10 +67,10 @@ public void testShortcutQueryAndFetchOptimization() throws IOException {
6667
}
6768

6869
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
69-
(searchResponse) -> new SearchPhase("test") {
70+
(searchResponse, scrollId) -> new SearchPhase("test") {
7071
@Override
7172
public void run() throws IOException {
72-
responseRef.set(searchResponse);
73+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
7374
}
7475
});
7576
assertEquals("fetch", phase.getName());
@@ -119,10 +120,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
119120
};
120121
mockSearchPhaseContext.searchTransport = searchTransportService;
121122
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
122-
(searchResponse) -> new SearchPhase("test") {
123+
(searchResponse, scrollId) -> new SearchPhase("test") {
123124
@Override
124125
public void run() throws IOException {
125-
responseRef.set(searchResponse);
126+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
126127
}
127128
});
128129
assertEquals("fetch", phase.getName());
@@ -173,10 +174,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
173174
};
174175
mockSearchPhaseContext.searchTransport = searchTransportService;
175176
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
176-
(searchResponse) -> new SearchPhase("test") {
177+
(searchResponse, scrollId) -> new SearchPhase("test") {
177178
@Override
178179
public void run() throws IOException {
179-
responseRef.set(searchResponse);
180+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
180181
}
181182
});
182183
assertEquals("fetch", phase.getName());
@@ -225,10 +226,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
225226
mockSearchPhaseContext.searchTransport = searchTransportService;
226227
CountDownLatch latch = new CountDownLatch(1);
227228
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
228-
(searchResponse) -> new SearchPhase("test") {
229+
(searchResponse, scrollId) -> new SearchPhase("test") {
229230
@Override
230231
public void run() throws IOException {
231-
responseRef.set(searchResponse);
232+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
232233
latch.countDown();
233234
}
234235
});
@@ -291,10 +292,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
291292
};
292293
mockSearchPhaseContext.searchTransport = searchTransportService;
293294
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
294-
(searchResponse) -> new SearchPhase("test") {
295+
(searchResponse, scrollId) -> new SearchPhase("test") {
295296
@Override
296297
public void run() throws IOException {
297-
responseRef.set(searchResponse);
298+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
298299
}
299300
});
300301
assertEquals("fetch", phase.getName());
@@ -340,10 +341,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
340341
};
341342
mockSearchPhaseContext.searchTransport = searchTransportService;
342343
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
343-
(searchResponse) -> new SearchPhase("test") {
344+
(searchResponse, scrollId) -> new SearchPhase("test") {
344345
@Override
345346
public void run() throws IOException {
346-
responseRef.set(searchResponse);
347+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
347348
}
348349
});
349350
assertEquals("fetch", phase.getName());
@@ -358,5 +359,4 @@ public void run() throws IOException {
358359
assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
359360
assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(123L));
360361
}
361-
362362
}

0 commit comments

Comments
 (0)