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

Example 1 with CompletionSuggestion

use of org.opensearch.search.suggest.completion.CompletionSuggestion in project OpenSearch by opensearch-project.

the class SearchResponseMerger method getMergedResponse.

/**
 * Returns the merged response. To be called once all responses have been added through {@link #add(SearchResponse)}
 * so that all responses are merged into a single one.
 */
SearchResponse getMergedResponse(SearchResponse.Clusters clusters) {
    // we end up calling merge without anything to merge, we just return an empty search response
    if (searchResponses.size() == 0) {
        return SearchResponse.empty(searchTimeProvider::buildTookInMillis, clusters);
    }
    int totalShards = 0;
    int skippedShards = 0;
    int successfulShards = 0;
    // the current reduce phase counts as one
    int numReducePhases = 1;
    List<ShardSearchFailure> failures = new ArrayList<>();
    Map<String, ProfileShardResult> profileResults = new HashMap<>();
    List<InternalAggregations> aggs = new ArrayList<>();
    Map<ShardIdAndClusterAlias, Integer> shards = new TreeMap<>();
    List<TopDocs> topDocsList = new ArrayList<>(searchResponses.size());
    Map<String, List<Suggest.Suggestion>> groupedSuggestions = new HashMap<>();
    Boolean trackTotalHits = null;
    SearchPhaseController.TopDocsStats topDocsStats = new SearchPhaseController.TopDocsStats(trackTotalHitsUpTo);
    for (SearchResponse searchResponse : searchResponses) {
        totalShards += searchResponse.getTotalShards();
        skippedShards += searchResponse.getSkippedShards();
        successfulShards += searchResponse.getSuccessfulShards();
        numReducePhases += searchResponse.getNumReducePhases();
        Collections.addAll(failures, searchResponse.getShardFailures());
        profileResults.putAll(searchResponse.getProfileResults());
        if (searchResponse.getAggregations() != null) {
            InternalAggregations internalAggs = (InternalAggregations) searchResponse.getAggregations();
            aggs.add(internalAggs);
        }
        Suggest suggest = searchResponse.getSuggest();
        if (suggest != null) {
            for (Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> entries : suggest) {
                List<Suggest.Suggestion> suggestionList = groupedSuggestions.computeIfAbsent(entries.getName(), s -> new ArrayList<>());
                suggestionList.add(entries);
            }
            List<CompletionSuggestion> completionSuggestions = suggest.filter(CompletionSuggestion.class);
            for (CompletionSuggestion completionSuggestion : completionSuggestions) {
                for (CompletionSuggestion.Entry options : completionSuggestion) {
                    for (CompletionSuggestion.Entry.Option option : options) {
                        SearchShardTarget shard = option.getHit().getShard();
                        ShardIdAndClusterAlias shardId = new ShardIdAndClusterAlias(shard.getShardId(), shard.getClusterAlias());
                        shards.putIfAbsent(shardId, null);
                    }
                }
            }
        }
        SearchHits searchHits = searchResponse.getHits();
        final TotalHits totalHits;
        if (searchHits.getTotalHits() == null) {
            // in case we didn't track total hits, we get null from each cluster, but we need to set 0 eq to the TopDocs
            totalHits = new TotalHits(0, TotalHits.Relation.EQUAL_TO);
            assert trackTotalHits == null || trackTotalHits == false;
            trackTotalHits = false;
        } else {
            totalHits = searchHits.getTotalHits();
            assert trackTotalHits == null || trackTotalHits;
            trackTotalHits = true;
        }
        TopDocs topDocs = searchHitsToTopDocs(searchHits, totalHits, shards);
        topDocsStats.add(new TopDocsAndMaxScore(topDocs, searchHits.getMaxScore()), searchResponse.isTimedOut(), searchResponse.isTerminatedEarly());
        if (searchHits.getHits().length > 0) {
            // there is no point in adding empty search hits and merging them with the others. Also, empty search hits always come
            // without sort fields and collapse info, despite sort by field and/or field collapsing was requested, which causes
            // issues reconstructing the proper TopDocs instance and breaks mergeTopDocs which expects the same type for each result.
            topDocsList.add(topDocs);
        }
    }
    // after going through all the hits and collecting all their distinct shards, we assign shardIndex and set it to the ScoreDocs
    setTopDocsShardIndex(shards, topDocsList);
    TopDocs topDocs = SearchPhaseController.mergeTopDocs(topDocsList, size, from);
    SearchHits mergedSearchHits = topDocsToSearchHits(topDocs, topDocsStats);
    setSuggestShardIndex(shards, groupedSuggestions);
    Suggest suggest = groupedSuggestions.isEmpty() ? null : new Suggest(Suggest.reduce(groupedSuggestions));
    InternalAggregations reducedAggs = InternalAggregations.topLevelReduce(aggs, aggReduceContextBuilder.forFinalReduction());
    ShardSearchFailure[] shardFailures = failures.toArray(ShardSearchFailure.EMPTY_ARRAY);
    SearchProfileShardResults profileShardResults = profileResults.isEmpty() ? null : new SearchProfileShardResults(profileResults);
    // make failures ordering consistent between ordinary search and CCS by looking at the shard they come from
    Arrays.sort(shardFailures, FAILURES_COMPARATOR);
    InternalSearchResponse response = new InternalSearchResponse(mergedSearchHits, reducedAggs, suggest, profileShardResults, topDocsStats.timedOut, topDocsStats.terminatedEarly, numReducePhases);
    long tookInMillis = searchTimeProvider.buildTookInMillis();
    return new SearchResponse(response, null, totalShards, successfulShards, skippedShards, tookInMillis, shardFailures, clusters, null);
}

Example 2 with CompletionSuggestion

use of org.opensearch.search.suggest.completion.CompletionSuggestion in project OpenSearch by opensearch-project.

the class CompletionSuggestSearchIT method testSuggestDocumentNoSource.

public void testSuggestDocumentNoSource() 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).setId("" + 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)).setFetchSource(false).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));
        assertThat(option.getHit(), hasId("" + id));
        assertThat(option.getHit(), hasScore((id)));
        assertNull(option.getHit().getSourceAsMap());
        id--;
    }
}

Example 3 with CompletionSuggestion

use of org.opensearch.search.suggest.completion.CompletionSuggestion in project OpenSearch by opensearch-project.

the class CompletionSuggestSearchIT method testSuggestDocumentSourceFiltering.

public void testSuggestDocumentSourceFiltering() 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).setId("" + i).setSource(jsonBuilder().startObject().startObject(FIELD).field("input", "suggestion" + i).field("weight", i).endObject().field("a", "include").field("b", "exclude").endObject()));
    }
    indexRandom(true, indexRequestBuilders);
    CompletionSuggestionBuilder prefix = SuggestBuilders.completionSuggestion(FIELD).prefix("sugg").size(numDocs);
    SearchResponse searchResponse = client().prepareSearch(INDEX).suggest(new SuggestBuilder().addSuggestion("foo", prefix)).setFetchSource("a", "b").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));
        assertThat(option.getHit(), hasId("" + id));
        assertThat(option.getHit(), hasScore((id)));
        assertNotNull(option.getHit().getSourceAsMap());
        Set<String> sourceFields = option.getHit().getSourceAsMap().keySet();
        assertThat(sourceFields, contains("a"));
        assertThat(sourceFields, not(contains("b")));
        id--;
    }
}

Example 4 with CompletionSuggestion

use of org.opensearch.search.suggest.completion.CompletionSuggestion in project OpenSearch by opensearch-project.

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_T_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().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().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 5 with CompletionSuggestion

use of org.opensearch.search.suggest.completion.CompletionSuggestion in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method generateQueryResults.

/**
 * Generate random query results received from the provided number of shards, including the provided
 * number of search hits and randomly generated completion suggestions based on the name and size of the provided ones.
 * Note that <code>shardIndex</code> is already set to the generated completion suggestions to simulate what
 * {@link SearchPhaseController#reducedQueryPhase} does,
 * meaning that the returned query results can be fed directly to {@link SearchPhaseController#sortDocs}
 */
private static AtomicArray<SearchPhaseResult> generateQueryResults(int nShards, List<CompletionSuggestion> suggestions, int searchHitsSize, boolean useConstantScore) {
    AtomicArray<SearchPhaseResult> queryResults = new AtomicArray<>(nShards);
    for (int shardIndex = 0; shardIndex < nShards; shardIndex++) {
        String clusterAlias = randomBoolean() ? null : "remote";
        SearchShardTarget searchShardTarget = new SearchShardTarget("", new ShardId("", "", shardIndex), clusterAlias, OriginalIndices.NONE);
        QuerySearchResult querySearchResult = new QuerySearchResult(new ShardSearchContextId("", shardIndex), searchShardTarget, null);
        final TopDocs topDocs;
        float maxScore = 0;
        if (searchHitsSize == 0) {
            topDocs = Lucene.EMPTY_TOP_DOCS;
        } else {
            int nDocs = randomIntBetween(0, searchHitsSize);
            ScoreDoc[] scoreDocs = new ScoreDoc[nDocs];
            for (int i = 0; i < nDocs; i++) {
                float score = useConstantScore ? 1.0F : Math.abs(randomFloat());
                scoreDocs[i] = new ScoreDoc(i, score);
                maxScore = Math.max(score, maxScore);
            }
            topDocs = new TopDocs(new TotalHits(scoreDocs.length, TotalHits.Relation.EQUAL_TO), scoreDocs);
        }
        List<CompletionSuggestion> shardSuggestion = new ArrayList<>();
        for (CompletionSuggestion completionSuggestion : suggestions) {
            CompletionSuggestion suggestion = new CompletionSuggestion(completionSuggestion.getName(), completionSuggestion.getSize(), false);
            final CompletionSuggestion.Entry completionEntry = new CompletionSuggestion.Entry(new Text(""), 0, 5);
            suggestion.addTerm(completionEntry);
            int optionSize = randomIntBetween(1, suggestion.getSize());
            float maxScoreValue = randomIntBetween(suggestion.getSize(), (int) Float.MAX_VALUE);
            for (int i = 0; i < optionSize; i++) {
                completionEntry.addOption(new CompletionSuggestion.Entry.Option(i, new Text(""), maxScoreValue, Collections.emptyMap()));
                float dec = randomIntBetween(0, optionSize);
                if (dec <= maxScoreValue) {
                    maxScoreValue -= dec;
                }
            }
            suggestion.setShardIndex(shardIndex);
            shardSuggestion.add(suggestion);
        }
        querySearchResult.topDocs(new TopDocsAndMaxScore(topDocs, maxScore), null);
        querySearchResult.size(searchHitsSize);
        querySearchResult.suggest(new Suggest(new ArrayList<>(shardSuggestion)));
        querySearchResult.setShardIndex(shardIndex);
        queryResults.set(shardIndex, querySearchResult);
    }
    return queryResults;
}