Examples with CompletionSuggestion org.elasticsearch.search.suggest.completion.CompletionSuggestion used on opensource projects

Example 6 with CompletionSuggestion

use of org.elasticsearch.search.suggest.completion.CompletionSuggestion in project elasticsearch by elastic.

the class SearchService method shortcutDocIdsToLoad.

/**
     * Shortcut ids to load, we load only "from" and up to "size". The phase controller
     * handles this as well since the result is always size * shards for Q_A_F
     */
private void shortcutDocIdsToLoad(SearchContext context) {
    final int[] docIdsToLoad;
    int docsOffset = 0;
    final Suggest suggest = context.queryResult().suggest();
    int numSuggestDocs = 0;
    final List<CompletionSuggestion> completionSuggestions;
    if (suggest != null && suggest.hasScoreDocs()) {
        completionSuggestions = suggest.filter(CompletionSuggestion.class);
        for (CompletionSuggestion completionSuggestion : completionSuggestions) {
            numSuggestDocs += completionSuggestion.getOptions().size();
        }
    } else {
        completionSuggestions = Collections.emptyList();
    }
    if (context.request().scroll() != null) {
        TopDocs topDocs = context.queryResult().topDocs();
        docIdsToLoad = new int[topDocs.scoreDocs.length + numSuggestDocs];
        for (int i = 0; i < topDocs.scoreDocs.length; i++) {
            docIdsToLoad[docsOffset++] = topDocs.scoreDocs[i].doc;
        }
    } else {
        TopDocs topDocs = context.queryResult().topDocs();
        if (topDocs.scoreDocs.length < context.from()) {
            // no more docs...
            docIdsToLoad = new int[numSuggestDocs];
        } else {
            int totalSize = context.from() + context.size();
            docIdsToLoad = new int[Math.min(topDocs.scoreDocs.length - context.from(), context.size()) + numSuggestDocs];
            for (int i = context.from(); i < Math.min(totalSize, topDocs.scoreDocs.length); i++) {
                docIdsToLoad[docsOffset++] = topDocs.scoreDocs[i].doc;
            }
        }
    }
    for (CompletionSuggestion completionSuggestion : completionSuggestions) {
        for (CompletionSuggestion.Entry.Option option : completionSuggestion.getOptions()) {
            docIdsToLoad[docsOffset++] = option.getDoc().doc;
        }
    }
    context.docIdsToLoad(docIdsToLoad, 0, docIdsToLoad.length);
}

Example 7 with CompletionSuggestion

use of org.elasticsearch.search.suggest.completion.CompletionSuggestion in project elasticsearch by elastic.

the class SearchPhaseController method merge.

/**
     * Enriches search hits and completion suggestion hits from <code>sortedDocs</code> using <code>fetchResultsArr</code>,
     * merges suggestions, aggregations and profile results
     *
     * Expects sortedDocs to have top search docs across all shards, optionally followed by top suggest docs for each named
     * completion suggestion ordered by suggestion name
     */
public InternalSearchResponse merge(boolean ignoreFrom, ScoreDoc[] sortedDocs, ReducedQueryPhase reducedQueryPhase, AtomicArray<? extends QuerySearchResultProvider> fetchResultsArr) {
    if (reducedQueryPhase.isEmpty()) {
        return InternalSearchResponse.empty();
    }
    List<? extends AtomicArray.Entry<? extends QuerySearchResultProvider>> fetchResults = fetchResultsArr.asList();
    SearchHits hits = getHits(reducedQueryPhase, ignoreFrom, sortedDocs, fetchResultsArr);
    if (reducedQueryPhase.suggest != null) {
        if (!fetchResults.isEmpty()) {
            int currentOffset = hits.getHits().length;
            for (CompletionSuggestion suggestion : reducedQueryPhase.suggest.filter(CompletionSuggestion.class)) {
                final List<CompletionSuggestion.Entry.Option> suggestionOptions = suggestion.getOptions();
                for (int scoreDocIndex = currentOffset; scoreDocIndex < currentOffset + suggestionOptions.size(); scoreDocIndex++) {
                    ScoreDoc shardDoc = sortedDocs[scoreDocIndex];
                    QuerySearchResultProvider searchResultProvider = fetchResultsArr.get(shardDoc.shardIndex);
                    if (searchResultProvider == null) {
                        continue;
                    }
                    FetchSearchResult fetchResult = searchResultProvider.fetchResult();
                    int fetchResultIndex = fetchResult.counterGetAndIncrement();
                    if (fetchResultIndex < fetchResult.hits().internalHits().length) {
                        SearchHit hit = fetchResult.hits().internalHits()[fetchResultIndex];
                        CompletionSuggestion.Entry.Option suggestOption = suggestionOptions.get(scoreDocIndex - currentOffset);
                        hit.score(shardDoc.score);
                        hit.shard(fetchResult.shardTarget());
                        suggestOption.setHit(hit);
                    }
                }
                currentOffset += suggestionOptions.size();
            }
            assert currentOffset == sortedDocs.length : "expected no more score doc slices";
        }
    }
    return reducedQueryPhase.buildResponse(hits);
}

Example 8 with CompletionSuggestion

use of org.elasticsearch.search.suggest.completion.CompletionSuggestion in project elasticsearch by elastic.

the class SearchPhaseController method sortDocs.

/**
     * 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
     */
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;
}

Example 9 with CompletionSuggestion

use of org.elasticsearch.search.suggest.completion.CompletionSuggestion in project elasticsearch by elastic.

the class SearchPhaseControllerTests method generateQueryResults.

private AtomicArray<QuerySearchResultProvider> generateQueryResults(int nShards, List<CompletionSuggestion> suggestions, int searchHitsSize, boolean useConstantScore) {
    AtomicArray<QuerySearchResultProvider> queryResults = new AtomicArray<>(nShards);
    for (int shardIndex = 0; shardIndex < nShards; shardIndex++) {
        QuerySearchResult querySearchResult = new QuerySearchResult(shardIndex, new SearchShardTarget("", new Index("", ""), shardIndex));
        TopDocs topDocs = new TopDocs(0, new ScoreDoc[0], 0);
        if (searchHitsSize > 0) {
            int nDocs = randomIntBetween(0, searchHitsSize);
            ScoreDoc[] scoreDocs = new ScoreDoc[nDocs];
            float maxScore = 0F;
            for (int i = 0; i < nDocs; i++) {
                float score = useConstantScore ? 1.0F : Math.abs(randomFloat());
                scoreDocs[i] = new ScoreDoc(i, score);
                if (score > maxScore) {
                    maxScore = score;
                }
            }
            topDocs = new TopDocs(scoreDocs.length, scoreDocs, maxScore);
        }
        List<CompletionSuggestion> shardSuggestion = new ArrayList<>();
        for (CompletionSuggestion completionSuggestion : suggestions) {
            CompletionSuggestion suggestion = new CompletionSuggestion(completionSuggestion.getName(), completionSuggestion.getSize());
            final CompletionSuggestion.Entry completionEntry = new CompletionSuggestion.Entry(new Text(""), 0, 5);
            suggestion.addTerm(completionEntry);
            int optionSize = randomIntBetween(1, suggestion.getSize());
            float maxScore = randomIntBetween(suggestion.getSize(), (int) Float.MAX_VALUE);
            for (int i = 0; i < optionSize; i++) {
                completionEntry.addOption(new CompletionSuggestion.Entry.Option(i, new Text(""), maxScore, Collections.emptyMap()));
                float dec = randomIntBetween(0, optionSize);
                if (dec <= maxScore) {
                    maxScore -= dec;
                }
            }
            suggestion.setShardIndex(shardIndex);
            shardSuggestion.add(suggestion);
        }
        querySearchResult.topDocs(topDocs, null);
        querySearchResult.size(searchHitsSize);
        querySearchResult.suggest(new Suggest(new ArrayList<>(shardSuggestion)));
        queryResults.set(shardIndex, querySearchResult);
    }
    return queryResults;
}

Example 10 with CompletionSuggestion

use of org.elasticsearch.search.suggest.completion.CompletionSuggestion in project elasticsearch by elastic.

the class CompletionSuggestSearchIT method testSuggestDocument.

public void testSuggestDocument() throws Exception {
    final CompletionMappingBuilder mapping = new CompletionMappingBuilder();
    createIndexAndMapping(mapping);
    int numDocs = randomIntBetween(10, 100);
    List<IndexRequestBuilder> indexRequestBuilders = new ArrayList<>();
    for (int i = 1; i <= numDocs; i++) {
        indexRequestBuilders.add(client().prepareIndex(INDEX, TYPE, "" + i).setSource(jsonBuilder().startObject().startObject(FIELD).field("input", "suggestion" + i).field("weight", i).endObject().endObject()));
    }
    indexRandom(true, indexRequestBuilders);
    CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg").size(numDocs);
    SearchResponse searchResponse = client().prepareSearch(INDEX).suggest(new SuggestBuilder().addSuggestion("foo", prefix)).get();
    CompletionSuggestion completionSuggestion = searchResponse.getSuggest().getSuggestion("foo");
    CompletionSuggestion.Entry options = completionSuggestion.getEntries().get(0);
    assertThat(options.getOptions().size(), equalTo(numDocs));
    int id = numDocs;
    for (CompletionSuggestion.Entry.Option option : options) {
        assertThat(option.getText().toString(), equalTo("suggestion" + id));
        assertSearchHit(option.getHit(), hasId("" + id));
        assertSearchHit(option.getHit(), hasScore((id)));
        assertNotNull(option.getHit().getSourceAsMap());
        id--;
    }
}