Examples with SearchPhaseResult org.opensearch.search.SearchPhaseResult used on opensource projects

Example 26 with SearchPhaseResult

use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method testSortDocsIsIdempotent.

public void testSortDocsIsIdempotent() throws Exception {
    int nShards = randomIntBetween(1, 20);
    int queryResultSize = randomBoolean() ? 0 : randomIntBetween(1, nShards * 2);
    long randomSeed = randomLong();
    boolean useConstantScore = randomBoolean();
    AtomicArray<SearchPhaseResult> results = generateSeededQueryResults(randomSeed, nShards, Collections.emptyList(), queryResultSize, useConstantScore);
    boolean ignoreFrom = randomBoolean();
    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();
    }
    List<TopDocs> topDocsList = new ArrayList<>();
    for (SearchPhaseResult result : results.asList()) {
        QuerySearchResult queryResult = result.queryResult();
        TopDocs topDocs = queryResult.consumeTopDocs().topDocs;
        topDocsList.add(topDocs);
        SearchPhaseController.setShardIndex(topDocs, result.getShardIndex());
    }
    ScoreDoc[] sortedDocs = SearchPhaseController.sortDocs(ignoreFrom, topDocsList, from, size, Collections.emptyList()).scoreDocs;
    results = generateSeededQueryResults(randomSeed, nShards, Collections.emptyList(), queryResultSize, useConstantScore);
    topDocsList = new ArrayList<>();
    for (SearchPhaseResult result : results.asList()) {
        QuerySearchResult queryResult = result.queryResult();
        TopDocs topDocs = queryResult.consumeTopDocs().topDocs;
        topDocsList.add(topDocs);
        SearchPhaseController.setShardIndex(topDocs, result.getShardIndex());
    }
    ScoreDoc[] sortedDocs2 = SearchPhaseController.sortDocs(ignoreFrom, topDocsList, from, size, Collections.emptyList()).scoreDocs;
    assertEquals(sortedDocs.length, sortedDocs2.length);
    for (int i = 0; i < sortedDocs.length; i++) {
        assertEquals(sortedDocs[i].doc, sortedDocs2[i].doc);
        assertEquals(sortedDocs[i].shardIndex, sortedDocs2[i].shardIndex);
        assertEquals(sortedDocs[i].score, sortedDocs2[i].score, 0.0f);
    }
}

Example 27 with SearchPhaseResult

use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method testMerge.

public void testMerge() {
    List<CompletionSuggestion> suggestions = new ArrayList<>();
    int maxSuggestSize = 0;
    for (int i = 0; i < randomIntBetween(1, 5); i++) {
        int size = randomIntBetween(1, 20);
        maxSuggestSize += size;
        suggestions.add(new CompletionSuggestion(randomAlphaOfLength(randomIntBetween(1, 5)), size, false));
    }
    int nShards = randomIntBetween(1, 20);
    int queryResultSize = randomBoolean() ? 0 : randomIntBetween(1, nShards * 2);
    AtomicArray<SearchPhaseResult> queryResults = generateQueryResults(nShards, suggestions, queryResultSize, false);
    for (int trackTotalHits : new int[] { SearchContext.TRACK_TOTAL_HITS_DISABLED, SearchContext.TRACK_TOTAL_HITS_ACCURATE }) {
        SearchPhaseController.ReducedQueryPhase reducedQueryPhase = searchPhaseController.reducedQueryPhase(queryResults.asList(), new ArrayList<>(), new ArrayList<>(), new SearchPhaseController.TopDocsStats(trackTotalHits), 0, true, InternalAggregationTestCase.emptyReduceContextBuilder(), true);
        AtomicArray<SearchPhaseResult> fetchResults = generateFetchResults(nShards, reducedQueryPhase.sortedTopDocs.scoreDocs, reducedQueryPhase.suggest);
        InternalSearchResponse mergedResponse = searchPhaseController.merge(false, reducedQueryPhase, fetchResults.asList(), fetchResults::get);
        if (trackTotalHits == SearchContext.TRACK_TOTAL_HITS_DISABLED) {
            assertNull(mergedResponse.hits.getTotalHits());
        } else {
            assertThat(mergedResponse.hits.getTotalHits().value, equalTo(0L));
            assertEquals(mergedResponse.hits.getTotalHits().relation, Relation.EQUAL_TO);
        }
        for (SearchHit hit : mergedResponse.hits().getHits()) {
            SearchPhaseResult searchPhaseResult = fetchResults.get(hit.getShard().getShardId().id());
            assertSame(searchPhaseResult.getSearchShardTarget(), hit.getShard());
        }
        int suggestSize = 0;
        for (Suggest.Suggestion s : reducedQueryPhase.suggest) {
            Stream<CompletionSuggestion.Entry> stream = s.getEntries().stream();
            suggestSize += stream.collect(Collectors.summingInt(e -> e.getOptions().size()));
        }
        assertThat(suggestSize, lessThanOrEqualTo(maxSuggestSize));
        assertThat(mergedResponse.hits().getHits().length, equalTo(reducedQueryPhase.sortedTopDocs.scoreDocs.length - suggestSize));
        Suggest suggestResult = mergedResponse.suggest();
        for (Suggest.Suggestion<?> suggestion : reducedQueryPhase.suggest) {
            assertThat(suggestion, instanceOf(CompletionSuggestion.class));
            if (suggestion.getEntries().get(0).getOptions().size() > 0) {
                CompletionSuggestion suggestionResult = suggestResult.getSuggestion(suggestion.getName());
                assertNotNull(suggestionResult);
                List<CompletionSuggestion.Entry.Option> options = suggestionResult.getEntries().get(0).getOptions();
                assertThat(options.size(), equalTo(suggestion.getEntries().get(0).getOptions().size()));
                for (CompletionSuggestion.Entry.Option option : options) {
                    assertNotNull(option.getHit());
                    SearchPhaseResult searchPhaseResult = fetchResults.get(option.getHit().getShard().getShardId().id());
                    assertSame(searchPhaseResult.getSearchShardTarget(), option.getHit().getShard());
                }
            }
        }
    }
}

Example 28 with SearchPhaseResult

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

Example 29 with SearchPhaseResult

use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method getTotalQueryHits.

private static int getTotalQueryHits(AtomicArray<SearchPhaseResult> results) {
    int resultCount = 0;
    for (SearchPhaseResult shardResult : results.asList()) {
        TopDocs topDocs = shardResult.queryResult().topDocs().topDocs;
        assert topDocs.totalHits.relation == Relation.EQUAL_TO;
        resultCount += (int) topDocs.totalHits.value;
    }
    return resultCount;
}

Example 30 with SearchPhaseResult

use of org.opensearch.search.SearchPhaseResult in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method consumerTestCase.

private void consumerTestCase(int numEmptyResponses) throws Exception {
    long beforeCompletedTasks = fixedExecutor.getCompletedTaskCount();
    int numShards = 3 + numEmptyResponses;
    int bufferSize = randomIntBetween(2, 3);
    CountDownLatch latch = new CountDownLatch(numShards);
    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, 3 + numEmptyResponses, exc -> {
    });
    if (numEmptyResponses == 0) {
        assertEquals(0, reductions.size());
    }
    if (numEmptyResponses > 0) {
        QuerySearchResult empty = QuerySearchResult.nullInstance();
        int shardId = 2 + numEmptyResponses;
        empty.setShardIndex(2 + numEmptyResponses);
        empty.setSearchShardTarget(new SearchShardTarget("node", new ShardId("a", "b", shardId), null, OriginalIndices.NONE));
        consumer.consumeResult(empty, latch::countDown);
        numEmptyResponses--;
    }
    QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", 0), new SearchShardTarget("node", new ShardId("a", "b", 0), null, OriginalIndices.NONE), null);
    result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[0]);
    InternalAggregations aggs = InternalAggregations.from(singletonList(new InternalMax("test", 1.0D, DocValueFormat.RAW, emptyMap())));
    result.aggregations(aggs);
    result.setShardIndex(0);
    consumer.consumeResult(result, latch::countDown);
    result = new QuerySearchResult(new ShardSearchContextId("", 1), new SearchShardTarget("node", new ShardId("a", "b", 0), null, OriginalIndices.NONE), null);
    result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[0]);
    aggs = InternalAggregations.from(singletonList(new InternalMax("test", 3.0D, DocValueFormat.RAW, emptyMap())));
    result.aggregations(aggs);
    result.setShardIndex(2);
    consumer.consumeResult(result, latch::countDown);
    result = new QuerySearchResult(new ShardSearchContextId("", 1), new SearchShardTarget("node", new ShardId("a", "b", 0), null, OriginalIndices.NONE), null);
    result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(0, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), Float.NaN), new DocValueFormat[0]);
    aggs = InternalAggregations.from(singletonList(new InternalMax("test", 2.0D, DocValueFormat.RAW, emptyMap())));
    result.aggregations(aggs);
    result.setShardIndex(1);
    consumer.consumeResult(result, latch::countDown);
    while (numEmptyResponses > 0) {
        result = QuerySearchResult.nullInstance();
        int shardId = 2 + numEmptyResponses;
        result.setShardIndex(shardId);
        result.setSearchShardTarget(new SearchShardTarget("node", new ShardId("a", "b", shardId), null, OriginalIndices.NONE));
        consumer.consumeResult(result, latch::countDown);
        numEmptyResponses--;
    }
    latch.await();
    final int numTotalReducePhases;
    if (numShards > bufferSize) {
        if (bufferSize == 2) {
            assertEquals(1, ((QueryPhaseResultConsumer) consumer).getNumReducePhases());
            assertEquals(1, reductions.size());
            assertEquals(false, reductions.get(0));
            numTotalReducePhases = 2;
        } else {
            assertEquals(0, ((QueryPhaseResultConsumer) consumer).getNumReducePhases());
            assertEquals(0, reductions.size());
            numTotalReducePhases = 1;
        }
    } else {
        assertEquals(0, reductions.size());
        numTotalReducePhases = 1;
    }
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    assertEquals(numTotalReducePhases, reduce.numReducePhases);
    assertEquals(numTotalReducePhases, reductions.size());
    assertAggReduction(request);
    InternalMax max = (InternalMax) reduce.aggregations.asList().get(0);
    assertEquals(3.0D, max.getValue(), 0.0D);
    assertFalse(reduce.sortedTopDocs.isSortedByField);
    assertNull(reduce.sortedTopDocs.sortFields);
    assertNull(reduce.sortedTopDocs.collapseField);
    assertNull(reduce.sortedTopDocs.collapseValues);
}