Skip to content

Commit ec64c2c

Browse files
jimczijimczi
authored
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 a77b38c commit ec64c2c

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());
@@ -224,10 +225,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
224225
mockSearchPhaseContext.searchTransport = searchTransportService;
225226
CountDownLatch latch = new CountDownLatch(1);
226227
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
227-
(searchResponse) -> new SearchPhase("test") {
228+
(searchResponse, scrollId) -> new SearchPhase("test") {
228229
@Override
229230
public void run() throws IOException {
230-
responseRef.set(searchResponse);
231+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
231232
latch.countDown();
232233
}
233234
});
@@ -290,10 +291,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
290291
};
291292
mockSearchPhaseContext.searchTransport = searchTransportService;
292293
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
293-
(searchResponse) -> new SearchPhase("test") {
294+
(searchResponse, scrollId) -> new SearchPhase("test") {
294295
@Override
295296
public void run() throws IOException {
296-
responseRef.set(searchResponse);
297+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
297298
}
298299
});
299300
assertEquals("fetch", phase.getName());
@@ -339,10 +340,10 @@ public void sendExecuteFetch(Transport.Connection connection, ShardFetchSearchRe
339340
};
340341
mockSearchPhaseContext.searchTransport = searchTransportService;
341342
FetchSearchPhase phase = new FetchSearchPhase(results, controller, mockSearchPhaseContext,
342-
(searchResponse) -> new SearchPhase("test") {
343+
(searchResponse, scrollId) -> new SearchPhase("test") {
343344
@Override
344345
public void run() throws IOException {
345-
responseRef.set(searchResponse);
346+
responseRef.set(mockSearchPhaseContext.buildSearchResponse(searchResponse, null));
346347
}
347348
});
348349
assertEquals("fetch", phase.getName());
@@ -357,5 +358,4 @@ public void run() throws IOException {
357358
assertEquals(1, mockSearchPhaseContext.releasedSearchContexts.size());
358359
assertTrue(mockSearchPhaseContext.releasedSearchContexts.contains(123L));
359360
}
360-
361361
}

0 commit comments

Comments
 (0)