use of org.opensearch.search.suggest.Suggest in project OpenSearch by opensearch-project.
the class SearchDocumentationIT method testSearchRequestSuggestions.
@SuppressWarnings({ "unused", "rawtypes" })
public void testSearchRequestSuggestions() throws IOException {
RestHighLevelClient client = highLevelClient();
{
BulkRequest request = new BulkRequest();
request.add(new IndexRequest("posts").id("1").source(XContentType.JSON, "user", "foobar"));
request.add(new IndexRequest("posts").id("2").source(XContentType.JSON, "user", "quxx"));
request.add(new IndexRequest("posts").id("3").source(XContentType.JSON, "user", "quzz"));
request.add(new IndexRequest("posts").id("4").source(XContentType.JSON, "user", "corge"));
request.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
BulkResponse bulkResponse = client.bulk(request, RequestOptions.DEFAULT);
assertSame(RestStatus.OK, bulkResponse.status());
assertFalse(bulkResponse.hasFailures());
}
{
SearchRequest searchRequest = new SearchRequest();
// tag::search-request-suggestion
SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
SuggestionBuilder termSuggestionBuilder = // <1>
SuggestBuilders.termSuggestion("user").text("fooarb");
SuggestBuilder suggestBuilder = new SuggestBuilder();
// <2>
suggestBuilder.addSuggestion("suggest_user", termSuggestionBuilder);
searchSourceBuilder.suggest(suggestBuilder);
// end::search-request-suggestion
searchRequest.source(searchSourceBuilder);
SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
{
// tag::search-request-suggestion-get
// <1>
Suggest suggest = searchResponse.getSuggest();
// <2>
TermSuggestion termSuggestion = suggest.getSuggestion("suggest_user");
for (TermSuggestion.Entry entry : termSuggestion.getEntries()) {
// <3>
for (TermSuggestion.Entry.Option option : entry) {
// <4>
String suggestText = option.getText().string();
}
}
// end::search-request-suggestion-get
assertEquals(1, termSuggestion.getEntries().size());
assertEquals(1, termSuggestion.getEntries().get(0).getOptions().size());
assertEquals("foobar", termSuggestion.getEntries().get(0).getOptions().get(0).getText().string());
}
}
}
use of org.opensearch.search.suggest.Suggest 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);
}
use of org.opensearch.search.suggest.Suggest in project OpenSearch by opensearch-project.
the class SearchResponseTests method createTestItem.
/**
* if minimal is set, don't include search hits, aggregations, suggest etc... to make test simpler
*/
private SearchResponse createTestItem(boolean minimal, ShardSearchFailure... shardSearchFailures) {
boolean timedOut = randomBoolean();
Boolean terminatedEarly = randomBoolean() ? null : randomBoolean();
int numReducePhases = randomIntBetween(1, 10);
long tookInMillis = randomNonNegativeLong();
int totalShards = randomIntBetween(1, Integer.MAX_VALUE);
int successfulShards = randomIntBetween(0, totalShards);
int skippedShards = randomIntBetween(0, totalShards);
InternalSearchResponse internalSearchResponse;
if (minimal == false) {
SearchHits hits = SearchHitsTests.createTestItem(true, true);
InternalAggregations aggregations = aggregationsTests.createTestInstance();
Suggest suggest = SuggestTests.createTestItem();
SearchProfileShardResults profileShardResults = SearchProfileShardResultsTests.createTestItem();
internalSearchResponse = new InternalSearchResponse(hits, aggregations, suggest, profileShardResults, timedOut, terminatedEarly, numReducePhases);
} else {
internalSearchResponse = InternalSearchResponse.empty();
}
return new SearchResponse(internalSearchResponse, null, totalShards, successfulShards, skippedShards, tookInMillis, shardSearchFailures, randomBoolean() ? randomClusters() : SearchResponse.Clusters.EMPTY);
}
use of org.opensearch.search.suggest.Suggest 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);
}
use of org.opensearch.search.suggest.Suggest 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;
}
Aggregations