List of usage examples for org.apache.lucene.search TopDocs merge
public static TopFieldDocs merge(Sort sort, int start, int topN, TopFieldDocs[] shardHits, boolean setShardIndex)
From source file:org.elasticsearch.action.search.SearchPhaseController.java
License:Apache License
/** * Returns a score doc array of top N search docs across all shards, followed by top suggest docs for each * named completion suggestion across all shards. If more than one named completion suggestion is specified in the * request, the suggest docs for a named suggestion are ordered by the suggestion name. * * Note: The order of the sorted score docs depends on the shard index in the result array if the merge process needs to disambiguate * the result. In oder to obtain stable results the shard index (index of the result in the result array) must be the same. * * @param ignoreFrom Whether to ignore the from and sort all hits in each shard result. * Enabled only for scroll search, because that only retrieves hits of length 'size' in the query phase. * @param resultsArr Shard result holder */// ww w . j a va 2 s . c o m public ScoreDoc[] sortDocs(boolean ignoreFrom, AtomicArray<? extends QuerySearchResultProvider> resultsArr) throws IOException { List<? extends AtomicArray.Entry<? extends QuerySearchResultProvider>> results = resultsArr.asList(); if (results.isEmpty()) { return EMPTY_DOCS; } final QuerySearchResult result; boolean canOptimize = false; int shardIndex = -1; if (results.size() == 1) { canOptimize = true; result = results.get(0).value.queryResult(); shardIndex = results.get(0).index; } else { boolean hasResult = false; QuerySearchResult resultToOptimize = null; // lets see if we only got hits from a single shard, if so, we can optimize... for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : results) { if (entry.value.queryResult().hasHits()) { if (hasResult) { // we already have one, can't really optimize canOptimize = false; break; } canOptimize = true; hasResult = true; resultToOptimize = entry.value.queryResult(); shardIndex = entry.index; } } result = canOptimize ? resultToOptimize : results.get(0).value.queryResult(); assert result != null; } if (canOptimize) { int offset = result.from(); if (ignoreFrom) { offset = 0; } ScoreDoc[] scoreDocs = result.topDocs().scoreDocs; ScoreDoc[] docs; int numSuggestDocs = 0; final Suggest suggest = result.queryResult().suggest(); final List<CompletionSuggestion> completionSuggestions; if (suggest != null) { completionSuggestions = suggest.filter(CompletionSuggestion.class); for (CompletionSuggestion suggestion : completionSuggestions) { numSuggestDocs += suggestion.getOptions().size(); } } else { completionSuggestions = Collections.emptyList(); } int docsOffset = 0; if (scoreDocs.length == 0 || scoreDocs.length < offset) { docs = new ScoreDoc[numSuggestDocs]; } else { int resultDocsSize = result.size(); if ((scoreDocs.length - offset) < resultDocsSize) { resultDocsSize = scoreDocs.length - offset; } docs = new ScoreDoc[resultDocsSize + numSuggestDocs]; for (int i = 0; i < resultDocsSize; i++) { ScoreDoc scoreDoc = scoreDocs[offset + i]; scoreDoc.shardIndex = shardIndex; docs[i] = scoreDoc; docsOffset++; } } for (CompletionSuggestion suggestion : completionSuggestions) { for (CompletionSuggestion.Entry.Option option : suggestion.getOptions()) { ScoreDoc doc = option.getDoc(); doc.shardIndex = shardIndex; docs[docsOffset++] = doc; } } return docs; } final int topN = result.queryResult().size(); final int from = ignoreFrom ? 0 : result.queryResult().from(); final TopDocs mergedTopDocs; final int numShards = resultsArr.length(); if (result.queryResult().topDocs() instanceof CollapseTopFieldDocs) { CollapseTopFieldDocs firstTopDocs = (CollapseTopFieldDocs) result.queryResult().topDocs(); final Sort sort = new Sort(firstTopDocs.fields); final CollapseTopFieldDocs[] shardTopDocs = new CollapseTopFieldDocs[numShards]; fillTopDocs(shardTopDocs, results, new CollapseTopFieldDocs(firstTopDocs.field, 0, new FieldDoc[0], sort.getSort(), new Object[0], Float.NaN)); mergedTopDocs = CollapseTopFieldDocs.merge(sort, from, topN, shardTopDocs); } else if (result.queryResult().topDocs() instanceof TopFieldDocs) { TopFieldDocs firstTopDocs = (TopFieldDocs) result.queryResult().topDocs(); final Sort sort = new Sort(firstTopDocs.fields); final TopFieldDocs[] shardTopDocs = new TopFieldDocs[resultsArr.length()]; fillTopDocs(shardTopDocs, results, new TopFieldDocs(0, new FieldDoc[0], sort.getSort(), Float.NaN)); mergedTopDocs = TopDocs.merge(sort, from, topN, shardTopDocs, true); } else { final TopDocs[] shardTopDocs = new TopDocs[resultsArr.length()]; fillTopDocs(shardTopDocs, results, Lucene.EMPTY_TOP_DOCS); mergedTopDocs = TopDocs.merge(from, topN, shardTopDocs, true); } ScoreDoc[] scoreDocs = mergedTopDocs.scoreDocs; final Map<String, List<Suggestion<CompletionSuggestion.Entry>>> groupedCompletionSuggestions = new HashMap<>(); // group suggestions and assign shard index for (AtomicArray.Entry<? extends QuerySearchResultProvider> sortedResult : results) { Suggest shardSuggest = sortedResult.value.queryResult().suggest(); if (shardSuggest != null) { for (CompletionSuggestion suggestion : shardSuggest.filter(CompletionSuggestion.class)) { suggestion.setShardIndex(sortedResult.index); List<Suggestion<CompletionSuggestion.Entry>> suggestions = groupedCompletionSuggestions .computeIfAbsent(suggestion.getName(), s -> new ArrayList<>()); suggestions.add(suggestion); } } } if (groupedCompletionSuggestions.isEmpty() == false) { int numSuggestDocs = 0; List<Suggestion<? extends Entry<? extends Entry.Option>>> completionSuggestions = new ArrayList<>( groupedCompletionSuggestions.size()); for (List<Suggestion<CompletionSuggestion.Entry>> groupedSuggestions : groupedCompletionSuggestions .values()) { final CompletionSuggestion completionSuggestion = CompletionSuggestion.reduceTo(groupedSuggestions); assert completionSuggestion != null; numSuggestDocs += completionSuggestion.getOptions().size(); completionSuggestions.add(completionSuggestion); } scoreDocs = new ScoreDoc[mergedTopDocs.scoreDocs.length + numSuggestDocs]; System.arraycopy(mergedTopDocs.scoreDocs, 0, scoreDocs, 0, mergedTopDocs.scoreDocs.length); int offset = mergedTopDocs.scoreDocs.length; Suggest suggestions = new Suggest(completionSuggestions); for (CompletionSuggestion completionSuggestion : suggestions.filter(CompletionSuggestion.class)) { for (CompletionSuggestion.Entry.Option option : completionSuggestion.getOptions()) { scoreDocs[offset++] = option.getDoc(); } } } return scoreDocs; }