Examples with QuerySearchResultProvider org.elasticsearch.search.query.QuerySearchResultProvider used on opensource projects

Example 16 with QuerySearchResultProvider

use of org.elasticsearch.search.query.QuerySearchResultProvider in project elasticsearch by elastic.

the class SearchPhaseControllerTests method testConsumer.

public void testConsumer() {
    int bufferSize = randomIntBetween(2, 3);
    SearchRequest request = new SearchRequest();
    request.source(new SearchSourceBuilder().aggregation(AggregationBuilders.avg("foo")));
    request.setBatchedReduceSize(bufferSize);
    InitialSearchPhase.SearchPhaseResults<QuerySearchResultProvider> consumer = searchPhaseController.newSearchPhaseResults(request, 3);
    QuerySearchResult result = new QuerySearchResult(0, new SearchShardTarget("node", new Index("a", "b"), 0));
    result.topDocs(new TopDocs(0, new ScoreDoc[0], 0.0F), new DocValueFormat[0]);
    InternalAggregations aggs = new InternalAggregations(Arrays.asList(new InternalMax("test", 1.0D, DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap())));
    result.aggregations(aggs);
    consumer.consumeResult(0, result);
    result = new QuerySearchResult(1, new SearchShardTarget("node", new Index("a", "b"), 0));
    result.topDocs(new TopDocs(0, new ScoreDoc[0], 0.0F), new DocValueFormat[0]);
    aggs = new InternalAggregations(Arrays.asList(new InternalMax("test", 3.0D, DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap())));
    result.aggregations(aggs);
    consumer.consumeResult(2, result);
    result = new QuerySearchResult(1, new SearchShardTarget("node", new Index("a", "b"), 0));
    result.topDocs(new TopDocs(0, new ScoreDoc[0], 0.0F), new DocValueFormat[0]);
    aggs = new InternalAggregations(Arrays.asList(new InternalMax("test", 2.0D, DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap())));
    result.aggregations(aggs);
    consumer.consumeResult(1, result);
    int numTotalReducePhases = 1;
    if (bufferSize == 2) {
        assertThat(consumer, instanceOf(SearchPhaseController.QueryPhaseResultConsumer.class));
        assertEquals(1, ((SearchPhaseController.QueryPhaseResultConsumer) consumer).getNumReducePhases());
        assertEquals(2, ((SearchPhaseController.QueryPhaseResultConsumer) consumer).getNumBuffered());
        numTotalReducePhases++;
    } else {
        assertThat(consumer, not(instanceOf(SearchPhaseController.QueryPhaseResultConsumer.class)));
    }
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    assertEquals(numTotalReducePhases, reduce.numReducePhases);
    InternalMax max = (InternalMax) reduce.aggregations.asList().get(0);
    assertEquals(3.0D, max.getValue(), 0.0D);
}

Example 17 with QuerySearchResultProvider

use of org.elasticsearch.search.query.QuerySearchResultProvider 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;
}