use of org.opensearch.search.SearchPhaseResult 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;
}
use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.
the class SearchPhaseControllerTests method testConsumerConcurrently.
public void testConsumerConcurrently() throws Exception {
int expectedNumResults = randomIntBetween(1, 100);
int bufferSize = randomIntBetween(2, 200);
SearchRequest request = randomSearchRequest();
request.source(new SearchSourceBuilder().aggregation(AggregationBuilders.avg("foo")));
request.setBatchedReduceSize(bufferSize);
ArraySearchPhaseResults<SearchPhaseResult> consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, expectedNumResults, exc -> {
});
AtomicInteger max = new AtomicInteger();
Thread[] threads = new Thread[expectedNumResults];
CountDownLatch latch = new CountDownLatch(expectedNumResults);
for (int i = 0; i < expectedNumResults; i++) {
int id = i;
threads[i] = new Thread(() -> {
int number = randomIntBetween(1, 1000);
max.updateAndGet(prev -> Math.max(prev, number));
QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", id), new SearchShardTarget("node", new ShardId("a", "b", id), null, OriginalIndices.NONE), null);
result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[] { new ScoreDoc(0, number) }), number), new DocValueFormat[0]);
InternalAggregations aggs = InternalAggregations.from(Collections.singletonList(new InternalMax("test", (double) number, DocValueFormat.RAW, Collections.emptyMap())));
result.aggregations(aggs);
result.setShardIndex(id);
result.size(1);
consumer.consumeResult(result, latch::countDown);
});
threads[i].start();
}
for (int i = 0; i < expectedNumResults; i++) {
threads[i].join();
}
latch.await();
SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
assertAggReduction(request);
InternalMax internalMax = (InternalMax) reduce.aggregations.asList().get(0);
assertEquals(max.get(), internalMax.getValue(), 0.0D);
assertEquals(1, reduce.sortedTopDocs.scoreDocs.length);
assertEquals(max.get(), reduce.maxScore, 0.0f);
assertEquals(expectedNumResults, reduce.totalHits.value);
assertEquals(max.get(), reduce.sortedTopDocs.scoreDocs[0].score, 0.0f);
assertFalse(reduce.sortedTopDocs.isSortedByField);
assertNull(reduce.sortedTopDocs.sortFields);
assertNull(reduce.sortedTopDocs.collapseField);
assertNull(reduce.sortedTopDocs.collapseValues);
}
use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.
the class SearchPhaseControllerTests method testSortDocs.
public void testSortDocs() {
List<CompletionSuggestion> suggestions = new ArrayList<>();
for (int i = 0; i < randomIntBetween(1, 5); i++) {
suggestions.add(new CompletionSuggestion(randomAlphaOfLength(randomIntBetween(1, 5)), randomIntBetween(1, 20), false));
}
int nShards = randomIntBetween(1, 20);
int queryResultSize = randomBoolean() ? 0 : randomIntBetween(1, nShards * 2);
AtomicArray<SearchPhaseResult> results = generateQueryResults(nShards, suggestions, queryResultSize, false);
Optional<SearchPhaseResult> first = results.asList().stream().findFirst();
int from = 0, size = 0;
if (first.isPresent()) {
from = first.get().queryResult().from();
size = first.get().queryResult().size();
}
int accumulatedLength = Math.min(queryResultSize, getTotalQueryHits(results));
List<CompletionSuggestion> reducedCompletionSuggestions = reducedSuggest(results);
for (Suggest.Suggestion<?> suggestion : reducedCompletionSuggestions) {
int suggestionSize = suggestion.getEntries().get(0).getOptions().size();
accumulatedLength += suggestionSize;
}
List<TopDocs> topDocsList = new ArrayList<>();
for (SearchPhaseResult result : results.asList()) {
QuerySearchResult queryResult = result.queryResult();
TopDocs topDocs = queryResult.consumeTopDocs().topDocs;
SearchPhaseController.setShardIndex(topDocs, result.getShardIndex());
topDocsList.add(topDocs);
}
ScoreDoc[] sortedDocs = SearchPhaseController.sortDocs(true, topDocsList, from, size, reducedCompletionSuggestions).scoreDocs;
assertThat(sortedDocs.length, equalTo(accumulatedLength));
}
use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.
the class SearchPhaseControllerTests method generateFetchResults.
private static AtomicArray<SearchPhaseResult> generateFetchResults(int nShards, ScoreDoc[] mergedSearchDocs, Suggest mergedSuggest) {
AtomicArray<SearchPhaseResult> fetchResults = new AtomicArray<>(nShards);
for (int shardIndex = 0; shardIndex < nShards; shardIndex++) {
float maxScore = -1F;
String clusterAlias = randomBoolean() ? null : "remote";
SearchShardTarget shardTarget = new SearchShardTarget("", new ShardId("", "", shardIndex), clusterAlias, OriginalIndices.NONE);
FetchSearchResult fetchSearchResult = new FetchSearchResult(new ShardSearchContextId("", shardIndex), shardTarget);
List<SearchHit> searchHits = new ArrayList<>();
for (ScoreDoc scoreDoc : mergedSearchDocs) {
if (scoreDoc.shardIndex == shardIndex) {
searchHits.add(new SearchHit(scoreDoc.doc, "", Collections.emptyMap(), Collections.emptyMap()));
if (scoreDoc.score > maxScore) {
maxScore = scoreDoc.score;
}
}
}
for (Suggest.Suggestion<?> suggestion : mergedSuggest) {
if (suggestion instanceof CompletionSuggestion) {
for (CompletionSuggestion.Entry.Option option : ((CompletionSuggestion) suggestion).getOptions()) {
ScoreDoc doc = option.getDoc();
if (doc.shardIndex == shardIndex) {
searchHits.add(new SearchHit(doc.doc, "", Collections.emptyMap(), Collections.emptyMap()));
if (doc.score > maxScore) {
maxScore = doc.score;
}
}
}
}
}
SearchHit[] hits = searchHits.toArray(new SearchHit[0]);
fetchSearchResult.hits(new SearchHits(hits, new TotalHits(hits.length, Relation.EQUAL_TO), maxScore));
fetchResults.set(shardIndex, fetchSearchResult);
}
return fetchResults;
}
use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.
the class SearchPhaseControllerTests method reducedSuggest.
private static List<CompletionSuggestion> reducedSuggest(AtomicArray<SearchPhaseResult> results) {
Map<String, List<Suggest.Suggestion<CompletionSuggestion.Entry>>> groupedSuggestion = new HashMap<>();
for (SearchPhaseResult entry : results.asList()) {
for (Suggest.Suggestion<?> suggestion : entry.queryResult().suggest()) {
List<Suggest.Suggestion<CompletionSuggestion.Entry>> suggests = groupedSuggestion.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>());
suggests.add((CompletionSuggestion) suggestion);
}
}
CompletionSuggestion completionSuggestion = new CompletionSuggestion(null, -1, randomBoolean());
return groupedSuggestion.values().stream().map(completionSuggestion::reduce).collect(Collectors.toList());
}
Aggregations