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

Example 1 with QuerySearchResult

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

the class SearchPhaseControllerTests method testConsumerConcurrently.

public void testConsumerConcurrently() throws InterruptedException {
    int expectedNumResults = randomIntBetween(1, 100);
    int bufferSize = randomIntBetween(2, 200);
    SearchRequest request = new SearchRequest();
    request.source(new SearchSourceBuilder().aggregation(AggregationBuilders.avg("foo")));
    request.setBatchedReduceSize(bufferSize);
    InitialSearchPhase.SearchPhaseResults<QuerySearchResultProvider> consumer = searchPhaseController.newSearchPhaseResults(request, expectedNumResults);
    AtomicInteger max = new AtomicInteger();
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    for (int i = 0; i < expectedNumResults; i++) {
        int id = i;
        Thread t = new Thread(() -> {
            int number = randomIntBetween(1, 1000);
            max.updateAndGet(prev -> Math.max(prev, number));
            QuerySearchResult result = new QuerySearchResult(id, new SearchShardTarget("node", new Index("a", "b"), id));
            result.topDocs(new TopDocs(id, new ScoreDoc[0], 0.0F), new DocValueFormat[0]);
            InternalAggregations aggs = new InternalAggregations(Arrays.asList(new InternalMax("test", (double) number, DocValueFormat.RAW, Collections.emptyList(), Collections.emptyMap())));
            result.aggregations(aggs);
            consumer.consumeResult(id, result);
            latch.countDown();
        });
        t.start();
    }
    latch.await();
    SearchPhaseController.ReducedQueryPhase reduce = consumer.reduce();
    InternalMax internalMax = (InternalMax) reduce.aggregations.asList().get(0);
    assertEquals(max.get(), internalMax.getValue(), 0.0D);
}

Example 2 with QuerySearchResult

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

the class DfsQueryPhaseTests method testFailPhaseOnException.

public void testFailPhaseOnException() throws IOException {
    AtomicArray<DfsSearchResult> results = new AtomicArray<>(2);
    AtomicReference<AtomicArray<QuerySearchResultProvider>> responseRef = new AtomicReference<>();
    results.set(0, new DfsSearchResult(1, new SearchShardTarget("node1", new Index("test", "na"), 0)));
    results.set(1, new DfsSearchResult(2, new SearchShardTarget("node2", new Index("test", "na"), 0)));
    results.get(0).termsStatistics(new Term[0], new TermStatistics[0]);
    results.get(1).termsStatistics(new Term[0], new TermStatistics[0]);
    SearchPhaseController controller = new SearchPhaseController(Settings.EMPTY, BigArrays.NON_RECYCLING_INSTANCE, null);
    SearchTransportService searchTransportService = new SearchTransportService(Settings.builder().put("search.remote.connect", false).build(), null, null) {

        @Override
        public void sendExecuteQuery(Transport.Connection connection, QuerySearchRequest request, SearchTask task, ActionListener<QuerySearchResult> listener) {
            if (request.id() == 1) {
                QuerySearchResult queryResult = new QuerySearchResult(123, new SearchShardTarget("node1", new Index("test", "na"), 0));
                queryResult.topDocs(new TopDocs(1, new ScoreDoc[] { new ScoreDoc(42, 1.0F) }, 2.0F), new DocValueFormat[0]);
                // the size of the result set
                queryResult.size(2);
                listener.onResponse(queryResult);
            } else if (request.id() == 2) {
                throw new UncheckedIOException(new MockDirectoryWrapper.FakeIOException());
            } else {
                fail("no such request ID: " + request.id());
            }
        }
    };
    MockSearchPhaseContext mockSearchPhaseContext = new MockSearchPhaseContext(2);
    mockSearchPhaseContext.searchTransport = searchTransportService;
    DfsQueryPhase phase = new DfsQueryPhase(results, controller, (response) -> new SearchPhase("test") {

        @Override
        public void run() throws IOException {
            responseRef.set(response.results);
        }
    }, mockSearchPhaseContext);
    assertEquals("dfs_query", phase.getName());
    expectThrows(UncheckedIOException.class, () -> phase.run());
    // phase execution will clean up on the contexts
    assertTrue(mockSearchPhaseContext.releasedSearchContexts.isEmpty());
}

Example 3 with QuerySearchResult

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

the class SearchScrollQueryThenFetchAsyncAction method executeFetchPhase.

private void executeFetchPhase() throws Exception {
    sortedShardDocs = searchPhaseController.sortDocs(true, queryResults);
    if (sortedShardDocs.length == 0) {
        finishHim(searchPhaseController.reducedQueryPhase(queryResults.asList()));
        return;
    }
    final IntArrayList[] docIdsToLoad = searchPhaseController.fillDocIdsToLoad(queryResults.length(), sortedShardDocs);
    SearchPhaseController.ReducedQueryPhase reducedQueryPhase = searchPhaseController.reducedQueryPhase(queryResults.asList());
    final ScoreDoc[] lastEmittedDocPerShard = searchPhaseController.getLastEmittedDocPerShard(reducedQueryPhase, sortedShardDocs, queryResults.length());
    final AtomicInteger counter = new AtomicInteger(docIdsToLoad.length);
    for (int i = 0; i < docIdsToLoad.length; i++) {
        final int index = i;
        final IntArrayList docIds = docIdsToLoad[index];
        if (docIds != null) {
            final QuerySearchResult querySearchResult = queryResults.get(index);
            ScoreDoc lastEmittedDoc = lastEmittedDocPerShard[index];
            ShardFetchRequest shardFetchRequest = new ShardFetchRequest(querySearchResult.id(), docIds, lastEmittedDoc);
            DiscoveryNode node = nodes.get(querySearchResult.shardTarget().getNodeId());
            searchTransportService.sendExecuteFetchScroll(node, shardFetchRequest, task, new ActionListener<FetchSearchResult>() {

                @Override
                public void onResponse(FetchSearchResult result) {
                    result.shardTarget(querySearchResult.shardTarget());
                    fetchResults.set(index, result);
                    if (counter.decrementAndGet() == 0) {
                        finishHim(reducedQueryPhase);
                    }
                }

                @Override
                public void onFailure(Exception t) {
                    if (logger.isDebugEnabled()) {
                        logger.debug("Failed to execute fetch phase", t);
                    }
                    successfulOps.decrementAndGet();
                    if (counter.decrementAndGet() == 0) {
                        finishHim(reducedQueryPhase);
                    }
                }
            });
        } else {
            // the counter is set to the total size of docIdsToLoad which can have null values so we have to count them down too
            if (counter.decrementAndGet() == 0) {
                finishHim(reducedQueryPhase);
            }
        }
    }
}

Example 4 with QuerySearchResult

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

the class SearchPhaseController method reducedQueryPhase.

/**
     * Reduces the given query results and consumes all aggregations and profile results.
     * @param queryResults a list of non-null query shard results
     * @param bufferdAggs a list of pre-collected / buffered aggregations. if this list is non-null all aggregations have been consumed
     *                    from all non-null query results.
     * @param numReducePhases the number of non-final reduce phases applied to the query results.
     * @see QuerySearchResult#consumeAggs()
     * @see QuerySearchResult#consumeProfileResult()
     */
private ReducedQueryPhase reducedQueryPhase(List<? extends AtomicArray.Entry<? extends QuerySearchResultProvider>> queryResults, List<InternalAggregations> bufferdAggs, int numReducePhases) {
    assert numReducePhases >= 0 : "num reduce phases must be >= 0 but was: " + numReducePhases;
    // increment for this phase
    numReducePhases++;
    long totalHits = 0;
    long fetchHits = 0;
    float maxScore = Float.NEGATIVE_INFINITY;
    boolean timedOut = false;
    Boolean terminatedEarly = null;
    if (queryResults.isEmpty()) {
        // early terminate we have nothing to reduce
        return new ReducedQueryPhase(totalHits, fetchHits, maxScore, timedOut, terminatedEarly, null, null, null, null, numReducePhases);
    }
    final QuerySearchResult firstResult = queryResults.get(0).value.queryResult();
    final boolean hasSuggest = firstResult.suggest() != null;
    final boolean hasProfileResults = firstResult.hasProfileResults();
    final boolean consumeAggs;
    final List<InternalAggregations> aggregationsList;
    if (bufferdAggs != null) {
        consumeAggs = false;
        // we already have results from intermediate reduces and just need to perform the final reduce
        assert firstResult.hasAggs() : "firstResult has no aggs but we got non null buffered aggs?";
        aggregationsList = bufferdAggs;
    } else if (firstResult.hasAggs()) {
        // the number of shards was less than the buffer size so we reduce agg results directly
        aggregationsList = new ArrayList<>(queryResults.size());
        consumeAggs = true;
    } else {
        // no aggregations
        aggregationsList = Collections.emptyList();
        consumeAggs = false;
    }
    // count the total (we use the query result provider here, since we might not get any hits (we scrolled past them))
    final Map<String, List<Suggestion>> groupedSuggestions = hasSuggest ? new HashMap<>() : Collections.emptyMap();
    final Map<String, ProfileShardResult> profileResults = hasProfileResults ? new HashMap<>(queryResults.size()) : Collections.emptyMap();
    for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults) {
        QuerySearchResult result = entry.value.queryResult();
        if (result.searchTimedOut()) {
            timedOut = true;
        }
        if (result.terminatedEarly() != null) {
            if (terminatedEarly == null) {
                terminatedEarly = result.terminatedEarly();
            } else if (result.terminatedEarly()) {
                terminatedEarly = true;
            }
        }
        totalHits += result.topDocs().totalHits;
        fetchHits += result.topDocs().scoreDocs.length;
        if (!Float.isNaN(result.topDocs().getMaxScore())) {
            maxScore = Math.max(maxScore, result.topDocs().getMaxScore());
        }
        if (hasSuggest) {
            assert result.suggest() != null;
            for (Suggestion<? extends Suggestion.Entry<? extends Suggestion.Entry.Option>> suggestion : result.suggest()) {
                List<Suggestion> suggestionList = groupedSuggestions.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>());
                suggestionList.add(suggestion);
            }
        }
        if (consumeAggs) {
            aggregationsList.add((InternalAggregations) result.consumeAggs());
        }
        if (hasProfileResults) {
            String key = result.shardTarget().toString();
            profileResults.put(key, result.consumeProfileResult());
        }
    }
    final Suggest suggest = groupedSuggestions.isEmpty() ? null : new Suggest(Suggest.reduce(groupedSuggestions));
    ReduceContext reduceContext = new ReduceContext(bigArrays, scriptService, true);
    final InternalAggregations aggregations = aggregationsList.isEmpty() ? null : reduceAggs(aggregationsList, firstResult.pipelineAggregators(), reduceContext);
    final SearchProfileShardResults shardResults = profileResults.isEmpty() ? null : new SearchProfileShardResults(profileResults);
    return new ReducedQueryPhase(totalHits, fetchHits, maxScore, timedOut, terminatedEarly, firstResult, suggest, aggregations, shardResults, numReducePhases);
}

Example 5 with QuerySearchResult

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

the class SearchPhaseController method sortDocs.

/**
     * Returns a score doc array of top N search docs across all shards, followed by top suggest docs for each
     * named completion suggestion across all shards. If more than one named completion suggestion is specified in the
     * request, the suggest docs for a named suggestion are ordered by the suggestion name.
     *
     * Note: The order of the sorted score docs depends on the shard index in the result array if the merge process needs to disambiguate
     * the result. In oder to obtain stable results the shard index (index of the result in the result array) must be the same.
     *
     * @param ignoreFrom Whether to ignore the from and sort all hits in each shard result.
     *                   Enabled only for scroll search, because that only retrieves hits of length 'size' in the query phase.
     * @param resultsArr Shard result holder
     */
public ScoreDoc[] sortDocs(boolean ignoreFrom, AtomicArray<? extends QuerySearchResultProvider> resultsArr) throws IOException {
    List<? extends AtomicArray.Entry<? extends QuerySearchResultProvider>> results = resultsArr.asList();
    if (results.isEmpty()) {
        return EMPTY_DOCS;
    }
    final QuerySearchResult result;
    boolean canOptimize = false;
    int shardIndex = -1;
    if (results.size() == 1) {
        canOptimize = true;
        result = results.get(0).value.queryResult();
        shardIndex = results.get(0).index;
    } else {
        boolean hasResult = false;
        QuerySearchResult resultToOptimize = null;
        // lets see if we only got hits from a single shard, if so, we can optimize...
        for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : results) {
            if (entry.value.queryResult().hasHits()) {
                if (hasResult) {
                    // we already have one, can't really optimize
                    canOptimize = false;
                    break;
                }
                canOptimize = true;
                hasResult = true;
                resultToOptimize = entry.value.queryResult();
                shardIndex = entry.index;
            }
        }
        result = canOptimize ? resultToOptimize : results.get(0).value.queryResult();
        assert result != null;
    }
    if (canOptimize) {
        int offset = result.from();
        if (ignoreFrom) {
            offset = 0;
        }
        ScoreDoc[] scoreDocs = result.topDocs().scoreDocs;
        ScoreDoc[] docs;
        int numSuggestDocs = 0;
        final Suggest suggest = result.queryResult().suggest();
        final List<CompletionSuggestion> completionSuggestions;
        if (suggest != null) {
            completionSuggestions = suggest.filter(CompletionSuggestion.class);
            for (CompletionSuggestion suggestion : completionSuggestions) {
                numSuggestDocs += suggestion.getOptions().size();
            }
        } else {
            completionSuggestions = Collections.emptyList();
        }
        int docsOffset = 0;
        if (scoreDocs.length == 0 || scoreDocs.length < offset) {
            docs = new ScoreDoc[numSuggestDocs];
        } else {
            int resultDocsSize = result.size();
            if ((scoreDocs.length - offset) < resultDocsSize) {
                resultDocsSize = scoreDocs.length - offset;
            }
            docs = new ScoreDoc[resultDocsSize + numSuggestDocs];
            for (int i = 0; i < resultDocsSize; i++) {
                ScoreDoc scoreDoc = scoreDocs[offset + i];
                scoreDoc.shardIndex = shardIndex;
                docs[i] = scoreDoc;
                docsOffset++;
            }
        }
        for (CompletionSuggestion suggestion : completionSuggestions) {
            for (CompletionSuggestion.Entry.Option option : suggestion.getOptions()) {
                ScoreDoc doc = option.getDoc();
                doc.shardIndex = shardIndex;
                docs[docsOffset++] = doc;
            }
        }
        return docs;
    }
    final int topN = result.queryResult().size();
    final int from = ignoreFrom ? 0 : result.queryResult().from();
    final TopDocs mergedTopDocs;
    final int numShards = resultsArr.length();
    if (result.queryResult().topDocs() instanceof CollapseTopFieldDocs) {
        CollapseTopFieldDocs firstTopDocs = (CollapseTopFieldDocs) result.queryResult().topDocs();
        final Sort sort = new Sort(firstTopDocs.fields);
        final CollapseTopFieldDocs[] shardTopDocs = new CollapseTopFieldDocs[numShards];
        fillTopDocs(shardTopDocs, results, new CollapseTopFieldDocs(firstTopDocs.field, 0, new FieldDoc[0], sort.getSort(), new Object[0], Float.NaN));
        mergedTopDocs = CollapseTopFieldDocs.merge(sort, from, topN, shardTopDocs);
    } else if (result.queryResult().topDocs() instanceof TopFieldDocs) {
        TopFieldDocs firstTopDocs = (TopFieldDocs) result.queryResult().topDocs();
        final Sort sort = new Sort(firstTopDocs.fields);
        final TopFieldDocs[] shardTopDocs = new TopFieldDocs[resultsArr.length()];
        fillTopDocs(shardTopDocs, results, new TopFieldDocs(0, new FieldDoc[0], sort.getSort(), Float.NaN));
        mergedTopDocs = TopDocs.merge(sort, from, topN, shardTopDocs, true);
    } else {
        final TopDocs[] shardTopDocs = new TopDocs[resultsArr.length()];
        fillTopDocs(shardTopDocs, results, Lucene.EMPTY_TOP_DOCS);
        mergedTopDocs = TopDocs.merge(from, topN, shardTopDocs, true);
    }
    ScoreDoc[] scoreDocs = mergedTopDocs.scoreDocs;
    final Map<String, List<Suggestion<CompletionSuggestion.Entry>>> groupedCompletionSuggestions = new HashMap<>();
    // group suggestions and assign shard index
    for (AtomicArray.Entry<? extends QuerySearchResultProvider> sortedResult : results) {
        Suggest shardSuggest = sortedResult.value.queryResult().suggest();
        if (shardSuggest != null) {
            for (CompletionSuggestion suggestion : shardSuggest.filter(CompletionSuggestion.class)) {
                suggestion.setShardIndex(sortedResult.index);
                List<Suggestion<CompletionSuggestion.Entry>> suggestions = groupedCompletionSuggestions.computeIfAbsent(suggestion.getName(), s -> new ArrayList<>());
                suggestions.add(suggestion);
            }
        }
    }
    if (groupedCompletionSuggestions.isEmpty() == false) {
        int numSuggestDocs = 0;
        List<Suggestion<? extends Entry<? extends Entry.Option>>> completionSuggestions = new ArrayList<>(groupedCompletionSuggestions.size());
        for (List<Suggestion<CompletionSuggestion.Entry>> groupedSuggestions : groupedCompletionSuggestions.values()) {
            final CompletionSuggestion completionSuggestion = CompletionSuggestion.reduceTo(groupedSuggestions);
            assert completionSuggestion != null;
            numSuggestDocs += completionSuggestion.getOptions().size();
            completionSuggestions.add(completionSuggestion);
        }
        scoreDocs = new ScoreDoc[mergedTopDocs.scoreDocs.length + numSuggestDocs];
        System.arraycopy(mergedTopDocs.scoreDocs, 0, scoreDocs, 0, mergedTopDocs.scoreDocs.length);
        int offset = mergedTopDocs.scoreDocs.length;
        Suggest suggestions = new Suggest(completionSuggestions);
        for (CompletionSuggestion completionSuggestion : suggestions.filter(CompletionSuggestion.class)) {
            for (CompletionSuggestion.Entry.Option option : completionSuggestion.getOptions()) {
                scoreDocs[offset++] = option.getDoc();
            }
        }
    }
    return scoreDocs;
}