Skip to content

Commit 469eb25

Browse files
areekareek
authored
Merge pull request #19536 from areek/enhancement/completion_suggester_documents
Add support for returning documents with completion suggester
2 parents fbbb633 + fee013c commit 469eb25

20 files changed

+918
-206
lines changed

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.elasticsearch.search.internal.ShardSearchTransportRequest;
4747
import org.elasticsearch.search.query.QuerySearchResult;
4848
import org.elasticsearch.search.query.QuerySearchResultProvider;
49+
import org.elasticsearch.search.suggest.Suggest;
4950
import org.elasticsearch.threadpool.ThreadPool;
5051

5152
import java.util.List;
@@ -74,7 +75,7 @@ abstract class AbstractSearchAsyncAction<FirstResult extends SearchPhaseResult>
7475
protected final AtomicArray<FirstResult> firstResults;
7576
private volatile AtomicArray<ShardSearchFailure> shardFailures;
7677
private final Object shardFailuresMutex = new Object();
77-
protected volatile ScoreDoc[] sortedShardList;
78+
protected volatile ScoreDoc[] sortedShardDocs;
7879

7980
protected AbstractSearchAsyncAction(ESLogger logger, SearchTransportService searchTransportService, ClusterService clusterService,
8081
IndexNameExpressionResolver indexNameExpressionResolver,
@@ -321,8 +322,11 @@ protected void releaseIrrelevantSearchContexts(AtomicArray<? extends QuerySearch
321322
// we only release search context that we did not fetch from if we are not scrolling
322323
if (request.scroll() == null) {
323324
for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults.asList()) {
324-
final TopDocs topDocs = entry.value.queryResult().queryResult().topDocs();
325-
if (topDocs != null && topDocs.scoreDocs.length > 0 // the shard had matches
325+
QuerySearchResult queryResult = entry.value.queryResult().queryResult();
326+
final TopDocs topDocs = queryResult.topDocs();
327+
final Suggest suggest = queryResult.suggest();
328+
if (((topDocs != null && topDocs.scoreDocs.length > 0) // the shard had matches
329+
||suggest != null && suggest.hasScoreDocs()) // or had suggest docs
326330
&& docIdsToLoad.get(entry.index) == null) { // but none of them made it to the global top docs
327331
try {
328332
DiscoveryNode node = nodes.get(entry.value.queryResult().shardTarget().nodeId());
@@ -343,12 +347,8 @@ protected void sendReleaseSearchContext(long contextId, DiscoveryNode node) {
343347

344348
protected ShardFetchSearchRequest createFetchRequest(QuerySearchResult queryResult, AtomicArray.Entry<IntArrayList> entry,
345349
ScoreDoc[] lastEmittedDocPerShard) {
346-
if (lastEmittedDocPerShard != null) {
347-
ScoreDoc lastEmittedDoc = lastEmittedDocPerShard[entry.index];
348-
return new ShardFetchSearchRequest(request, queryResult.id(), entry.value, lastEmittedDoc);
349-
} else {
350-
return new ShardFetchSearchRequest(request, queryResult.id(), entry.value);
351-
}
350+
final ScoreDoc lastEmittedDoc = (lastEmittedDocPerShard != null) ? lastEmittedDocPerShard[entry.index] : null;
351+
return new ShardFetchSearchRequest(request, queryResult.id(), entry.value, lastEmittedDoc);
352352
}
353353

354354
protected abstract void sendExecuteFirstPhase(DiscoveryNode node, ShardSearchTransportRequest request,

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,8 @@ private void finishHim() {
118118
threadPool.executor(ThreadPool.Names.SEARCH).execute(new ActionRunnable<SearchResponse>(listener) {
119119
@Override
120120
public void doRun() throws IOException {
121-
sortedShardList = searchPhaseController.sortDocs(true, queryFetchResults);
122-
final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, queryFetchResults,
121+
sortedShardDocs = searchPhaseController.sortDocs(true, queryFetchResults);
122+
final InternalSearchResponse internalResponse = searchPhaseController.merge(true, sortedShardDocs, queryFetchResults,
123123
queryFetchResults);
124124
String scrollId = null;
125125
if (request.scroll() != null) {

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

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -135,18 +135,17 @@ void executeFetchPhase() {
135135
}
136136

137137
void innerExecuteFetchPhase() throws Exception {
138-
boolean useScroll = request.scroll() != null;
139-
sortedShardList = searchPhaseController.sortDocs(useScroll, queryResults);
140-
searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardList);
138+
final boolean isScrollRequest = request.scroll() != null;
139+
sortedShardDocs = searchPhaseController.sortDocs(isScrollRequest, queryResults);
140+
searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardDocs);
141141

142142
if (docIdsToLoad.asList().isEmpty()) {
143143
finishHim();
144144
return;
145145
}
146146

147-
final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(
148-
request, sortedShardList, firstResults.length()
149-
);
147+
final ScoreDoc[] lastEmittedDocPerShard = (request.scroll() != null) ?
148+
searchPhaseController.getLastEmittedDocPerShard(queryResults.asList(), sortedShardDocs, firstResults.length()) : null;
150149
final AtomicInteger counter = new AtomicInteger(docIdsToLoad.asList().size());
151150
for (final AtomicArray.Entry<IntArrayList> entry : docIdsToLoad.asList()) {
152151
QuerySearchResult queryResult = queryResults.get(entry.index);
@@ -196,12 +195,10 @@ private void finishHim() {
196195
threadPool.executor(ThreadPool.Names.SEARCH).execute(new ActionRunnable<SearchResponse>(listener) {
197196
@Override
198197
public void doRun() throws IOException {
199-
final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, queryResults,
198+
final boolean isScrollRequest = request.scroll() != null;
199+
final InternalSearchResponse internalResponse = searchPhaseController.merge(isScrollRequest, sortedShardDocs, queryResults,
200200
fetchResults);
201-
String scrollId = null;
202-
if (request.scroll() != null) {
203-
scrollId = TransportSearchHelper.buildScrollId(request.searchType(), firstResults);
204-
}
201+
String scrollId = isScrollRequest ? TransportSearchHelper.buildScrollId(request.searchType(), firstResults) : null;
205202
listener.onResponse(new SearchResponse(internalResponse, scrollId, expectedSuccessfulOps, successfulOps.get(),
206203
buildTookInMillis(), buildShardFailures()));
207204
releaseIrrelevantSearchContexts(queryResults, docIdsToLoad);

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

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -60,14 +60,11 @@ protected void moveToSecondPhase() throws Exception {
6060
threadPool.executor(ThreadPool.Names.SEARCH).execute(new ActionRunnable<SearchResponse>(listener) {
6161
@Override
6262
public void doRun() throws IOException {
63-
boolean useScroll = request.scroll() != null;
64-
sortedShardList = searchPhaseController.sortDocs(useScroll, firstResults);
65-
final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, firstResults,
63+
final boolean isScrollRequest = request.scroll() != null;
64+
sortedShardDocs = searchPhaseController.sortDocs(isScrollRequest, firstResults);
65+
final InternalSearchResponse internalResponse = searchPhaseController.merge(isScrollRequest, sortedShardDocs, firstResults,
6666
firstResults);
67-
String scrollId = null;
68-
if (request.scroll() != null) {
69-
scrollId = TransportSearchHelper.buildScrollId(request.searchType(), firstResults);
70-
}
67+
String scrollId = isScrollRequest ? TransportSearchHelper.buildScrollId(request.searchType(), firstResults) : null;
7168
listener.onResponse(new SearchResponse(internalResponse, scrollId, expectedSuccessfulOps, successfulOps.get(),
7269
buildTookInMillis(), buildShardFailures()));
7370
}

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

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -68,18 +68,17 @@ protected void sendExecuteFirstPhase(DiscoveryNode node, ShardSearchTransportReq
6868

6969
@Override
7070
protected void moveToSecondPhase() throws Exception {
71-
boolean useScroll = request.scroll() != null;
72-
sortedShardList = searchPhaseController.sortDocs(useScroll, firstResults);
73-
searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardList);
71+
final boolean isScrollRequest = request.scroll() != null;
72+
sortedShardDocs = searchPhaseController.sortDocs(isScrollRequest, firstResults);
73+
searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardDocs);
7474

7575
if (docIdsToLoad.asList().isEmpty()) {
7676
finishHim();
7777
return;
7878
}
7979

80-
final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(
81-
request, sortedShardList, firstResults.length()
82-
);
80+
final ScoreDoc[] lastEmittedDocPerShard = isScrollRequest ?
81+
searchPhaseController.getLastEmittedDocPerShard(firstResults.asList(), sortedShardDocs, firstResults.length()) : null;
8382
final AtomicInteger counter = new AtomicInteger(docIdsToLoad.asList().size());
8483
for (AtomicArray.Entry<IntArrayList> entry : docIdsToLoad.asList()) {
8584
QuerySearchResultProvider queryResult = firstResults.get(entry.index);
@@ -129,12 +128,10 @@ private void finishHim() {
129128
threadPool.executor(ThreadPool.Names.SEARCH).execute(new ActionRunnable<SearchResponse>(listener) {
130129
@Override
131130
public void doRun() throws IOException {
132-
final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, firstResults,
131+
final boolean isScrollRequest = request.scroll() != null;
132+
final InternalSearchResponse internalResponse = searchPhaseController.merge(isScrollRequest, sortedShardDocs, firstResults,
133133
fetchResults);
134-
String scrollId = null;
135-
if (request.scroll() != null) {
136-
scrollId = TransportSearchHelper.buildScrollId(request.searchType(), firstResults);
137-
}
134+
String scrollId = isScrollRequest ? TransportSearchHelper.buildScrollId(request.searchType(), firstResults) : null;
138135
listener.onResponse(new SearchResponse(internalResponse, scrollId, expectedSuccessfulOps,
139136
successfulOps.get(), buildTookInMillis(), buildShardFailures()));
140137
releaseIrrelevantSearchContexts(firstResults, docIdsToLoad);

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -168,8 +168,8 @@ private void finishHim() {
168168
}
169169

170170
private void innerFinishHim() throws Exception {
171-
ScoreDoc[] sortedShardList = searchPhaseController.sortDocs(true, queryFetchResults);
172-
final InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, queryFetchResults,
171+
ScoreDoc[] sortedShardDocs = searchPhaseController.sortDocs(true, queryFetchResults);
172+
final InternalSearchResponse internalResponse = searchPhaseController.merge(true, sortedShardDocs, queryFetchResults,
173173
queryFetchResults);
174174
String scrollId = null;
175175
if (request.scroll() != null) {

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

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ class SearchScrollQueryThenFetchAsyncAction extends AbstractAsyncAction {
5353
private volatile AtomicArray<ShardSearchFailure> shardFailures;
5454
final AtomicArray<QuerySearchResult> queryResults;
5555
final AtomicArray<FetchSearchResult> fetchResults;
56-
private volatile ScoreDoc[] sortedShardList;
56+
private volatile ScoreDoc[] sortedShardDocs;
5757
private final AtomicInteger successfulOps;
5858

5959
SearchScrollQueryThenFetchAsyncAction(ESLogger logger, ClusterService clusterService,
@@ -165,17 +165,18 @@ void onQueryPhaseFailure(final int shardIndex, final AtomicInteger counter, fina
165165
}
166166

167167
private void executeFetchPhase() throws Exception {
168-
sortedShardList = searchPhaseController.sortDocs(true, queryResults);
168+
sortedShardDocs = searchPhaseController.sortDocs(true, queryResults);
169169
AtomicArray<IntArrayList> docIdsToLoad = new AtomicArray<>(queryResults.length());
170-
searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardList);
170+
searchPhaseController.fillDocIdsToLoad(docIdsToLoad, sortedShardDocs);
171171

172172
if (docIdsToLoad.asList().isEmpty()) {
173173
finishHim();
174174
return;
175175
}
176176

177177

178-
final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(sortedShardList, queryResults.length());
178+
final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(queryResults.asList(),
179+
sortedShardDocs, queryResults.length());
179180
final AtomicInteger counter = new AtomicInteger(docIdsToLoad.asList().size());
180181
for (final AtomicArray.Entry<IntArrayList> entry : docIdsToLoad.asList()) {
181182
IntArrayList docIds = entry.value;
@@ -216,7 +217,7 @@ private void finishHim() {
216217
}
217218

218219
private void innerFinishHim() {
219-
InternalSearchResponse internalResponse = searchPhaseController.merge(sortedShardList, queryResults, fetchResults);
220+
InternalSearchResponse internalResponse = searchPhaseController.merge(true, sortedShardDocs, queryResults, fetchResults);
220221
String scrollId = null;
221222
if (request.scroll() != null) {
222223
scrollId = request.scrollId();

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

Lines changed: 39 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import com.carrotsearch.hppc.ObjectFloatHashMap;
2323
import org.apache.lucene.search.FieldDoc;
24+
import org.apache.lucene.search.ScoreDoc;
2425
import org.apache.lucene.search.TopDocs;
2526
import org.elasticsearch.ElasticsearchException;
2627
import org.elasticsearch.ExceptionsHelper;
@@ -87,13 +88,16 @@
8788
import org.elasticsearch.search.searchafter.SearchAfterBuilder;
8889
import org.elasticsearch.search.sort.SortAndFormats;
8990
import org.elasticsearch.search.sort.SortBuilder;
91+
import org.elasticsearch.search.suggest.Suggest;
92+
import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
9093
import org.elasticsearch.threadpool.ThreadPool;
9194
import org.elasticsearch.threadpool.ThreadPool.Cancellable;
9295
import org.elasticsearch.threadpool.ThreadPool.Names;
9396

9497
import java.io.IOException;
9598
import java.util.Collections;
9699
import java.util.HashMap;
100+
import java.util.List;
97101
import java.util.Map;
98102
import java.util.Optional;
99103
import java.util.concurrent.ExecutionException;
@@ -265,7 +269,7 @@ public QuerySearchResultProvider executeQueryPhase(ShardSearchRequest request) t
265269

266270
loadOrExecuteQueryPhase(request, context);
267271

268-
if (context.queryResult().topDocs().scoreDocs.length == 0 && context.scrollContext() == null) {
272+
if (hasHits(context.queryResult()) == false && context.scrollContext() == null) {
269273
freeContext(context.id());
270274
} else {
271275
contextProcessedSuccessfully(context);
@@ -320,7 +324,7 @@ public QuerySearchResult executeQueryPhase(QuerySearchRequest request) {
320324
operationListener.onPreQueryPhase(context);
321325
long time = System.nanoTime();
322326
queryPhase.execute(context);
323-
if (context.queryResult().topDocs().scoreDocs.length == 0 && context.scrollContext() == null) {
327+
if (hasHits(context.queryResult()) == false && context.scrollContext() == null) {
324328
// no hits, we can release the context since there will be no fetch phase
325329
freeContext(context.id());
326330
} else {
@@ -811,40 +815,55 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
811815
}
812816
}
813817

814-
private static final int[] EMPTY_DOC_IDS = new int[0];
815-
816818
/**
817819
* Shortcut ids to load, we load only "from" and up to "size". The phase controller
818820
* handles this as well since the result is always size * shards for Q_A_F
819821
*/
820822
private void shortcutDocIdsToLoad(SearchContext context) {
823+
final int[] docIdsToLoad;
824+
int docsOffset = 0;
825+
final Suggest suggest = context.queryResult().suggest();
826+
int numSuggestDocs = 0;
827+
final List<CompletionSuggestion> completionSuggestions;
828+
if (suggest != null && suggest.hasScoreDocs()) {
829+
completionSuggestions = suggest.filter(CompletionSuggestion.class);
830+
for (CompletionSuggestion completionSuggestion : completionSuggestions) {
831+
numSuggestDocs += completionSuggestion.getOptions().size();
832+
}
833+
} else {
834+
completionSuggestions = Collections.emptyList();
835+
}
821836
if (context.request().scroll() != null) {
822837
TopDocs topDocs = context.queryResult().topDocs();
823-
int[] docIdsToLoad = new int[topDocs.scoreDocs.length];
838+
docIdsToLoad = new int[topDocs.scoreDocs.length + numSuggestDocs];
824839
for (int i = 0; i < topDocs.scoreDocs.length; i++) {
825-
docIdsToLoad[i] = topDocs.scoreDocs[i].doc;
840+
docIdsToLoad[docsOffset++] = topDocs.scoreDocs[i].doc;
826841
}
827-
context.docIdsToLoad(docIdsToLoad, 0, docIdsToLoad.length);
828842
} else {
829843
TopDocs topDocs = context.queryResult().topDocs();
830844
if (topDocs.scoreDocs.length < context.from()) {
831845
// no more docs...
832-
context.docIdsToLoad(EMPTY_DOC_IDS, 0, 0);
833-
return;
834-
}
835-
int totalSize = context.from() + context.size();
836-
int[] docIdsToLoad = new int[Math.min(topDocs.scoreDocs.length - context.from(), context.size())];
837-
int counter = 0;
838-
for (int i = context.from(); i < totalSize; i++) {
839-
if (i < topDocs.scoreDocs.length) {
840-
docIdsToLoad[counter] = topDocs.scoreDocs[i].doc;
841-
} else {
842-
break;
846+
docIdsToLoad = new int[numSuggestDocs];
847+
} else {
848+
int totalSize = context.from() + context.size();
849+
docIdsToLoad = new int[Math.min(topDocs.scoreDocs.length - context.from(), context.size()) +
850+
numSuggestDocs];
851+
for (int i = context.from(); i < Math.min(totalSize, topDocs.scoreDocs.length); i++) {
852+
docIdsToLoad[docsOffset++] = topDocs.scoreDocs[i].doc;
843853
}
844-
counter++;
845854
}
846-
context.docIdsToLoad(docIdsToLoad, 0, counter);
847855
}
856+
for (CompletionSuggestion completionSuggestion : completionSuggestions) {
857+
for (CompletionSuggestion.Entry.Option option : completionSuggestion.getOptions()) {
858+
docIdsToLoad[docsOffset++] = option.getDoc().doc;
859+
}
860+
}
861+
context.docIdsToLoad(docIdsToLoad, 0, docIdsToLoad.length);
862+
}
863+
864+
private static boolean hasHits(final QuerySearchResult searchResult) {
865+
return searchResult.topDocs().scoreDocs.length > 0 ||
866+
(searchResult.suggest() != null && searchResult.suggest().hasScoreDocs());
848867
}
849868

850869
private void processScroll(InternalScrollSearchRequest request, SearchContext context) {

0 commit comments

Comments
 (0)