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

Example 1 with QuerySearchResult

use of org.opensearch.search.query.QuerySearchResult in project OpenSearch by opensearch-project.

the class QueryPhaseResultConsumer method partialReduce.

private MergeResult partialReduce(QuerySearchResult[] toConsume, List<SearchShard> emptyResults, SearchPhaseController.TopDocsStats topDocsStats, MergeResult lastMerge, int numReducePhases) {
    // ensure consistent ordering
    Arrays.sort(toConsume, Comparator.comparingInt(QuerySearchResult::getShardIndex));
    for (QuerySearchResult result : toConsume) {
        topDocsStats.add(result.topDocs(), result.searchTimedOut(), result.terminatedEarly());
    }
    final TopDocs newTopDocs;
    if (hasTopDocs) {
        List<TopDocs> topDocsList = new ArrayList<>();
        if (lastMerge != null) {
            topDocsList.add(lastMerge.reducedTopDocs);
        }
        for (QuerySearchResult result : toConsume) {
            TopDocsAndMaxScore topDocs = result.consumeTopDocs();
            SearchPhaseController.setShardIndex(topDocs.topDocs, result.getShardIndex());
            topDocsList.add(topDocs.topDocs);
        }
        newTopDocs = SearchPhaseController.mergeTopDocs(topDocsList, // we have to merge here in the same way we collect on a shard
        topNSize, 0);
    } else {
        newTopDocs = null;
    }
    final InternalAggregations newAggs;
    if (hasAggs) {
        List<InternalAggregations> aggsList = new ArrayList<>();
        if (lastMerge != null) {
            aggsList.add(lastMerge.reducedAggs);
        }
        for (QuerySearchResult result : toConsume) {
            aggsList.add(result.consumeAggs().expand());
        }
        newAggs = InternalAggregations.topLevelReduce(aggsList, aggReduceContextBuilder.forPartialReduction());
    } else {
        newAggs = null;
    }
    List<SearchShard> processedShards = new ArrayList<>(emptyResults);
    if (lastMerge != null) {
        processedShards.addAll(lastMerge.processedShards);
    }
    for (QuerySearchResult result : toConsume) {
        SearchShardTarget target = result.getSearchShardTarget();
        processedShards.add(new SearchShard(target.getClusterAlias(), target.getShardId()));
    }
    progressListener.notifyPartialReduce(processedShards, topDocsStats.getTotalHits(), newAggs, numReducePhases);
    // we leave the results un-serialized because serializing is slow but we compute the serialized
    // size as an estimate of the memory used by the newly reduced aggregations.
    long serializedSize = hasAggs ? newAggs.getSerializedSize() : 0;
    return new MergeResult(processedShards, newTopDocs, newAggs, hasAggs ? serializedSize : 0);
}

Example 2 with QuerySearchResult

use of org.opensearch.search.query.QuerySearchResult in project OpenSearch by opensearch-project.

the class DfsQueryPhase method run.

@Override
public void run() throws IOException {
    // TODO we can potentially also consume the actual per shard results from the initial phase here in the aggregateDfs
    // to free up memory early
    final CountedCollector<SearchPhaseResult> counter = new CountedCollector<>(queryResult, searchResults.size(), () -> context.executeNextPhase(this, nextPhaseFactory.apply(queryResult)), context);
    for (final DfsSearchResult dfsResult : searchResults) {
        final SearchShardTarget searchShardTarget = dfsResult.getSearchShardTarget();
        Transport.Connection connection = context.getConnection(searchShardTarget.getClusterAlias(), searchShardTarget.getNodeId());
        QuerySearchRequest querySearchRequest = new QuerySearchRequest(searchShardTarget.getOriginalIndices(), dfsResult.getContextId(), dfsResult.getShardSearchRequest(), dfs);
        final int shardIndex = dfsResult.getShardIndex();
        searchTransportService.sendExecuteQuery(connection, querySearchRequest, context.getTask(), new SearchActionListener<QuerySearchResult>(searchShardTarget, shardIndex) {

            @Override
            protected void innerOnResponse(QuerySearchResult response) {
                try {
                    counter.onResult(response);
                } catch (Exception e) {
                    context.onPhaseFailure(DfsQueryPhase.this, "", e);
                }
            }

            @Override
            public void onFailure(Exception exception) {
                try {
                    context.getLogger().debug(() -> new ParameterizedMessage("[{}] Failed to execute query phase", querySearchRequest.contextId()), exception);
                    progressListener.notifyQueryFailure(shardIndex, searchShardTarget, exception);
                    counter.onFailure(shardIndex, searchShardTarget, exception);
                } finally {
                    if (context.getRequest().pointInTimeBuilder() == null) {
                        // the query might not have been executed at all (for example because thread pool rejected
                        // execution) and the search context that was created in dfs phase might not be released.
                        // release it again to be in the safe side
                        context.sendReleaseSearchContext(querySearchRequest.contextId(), connection, searchShardTarget.getOriginalIndices());
                    }
                }
            }
        });
    }
}

Example 3 with QuerySearchResult

use of org.opensearch.search.query.QuerySearchResult 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);
}

Example 4 with QuerySearchResult

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

Example 5 with QuerySearchResult

use of org.opensearch.search.query.QuerySearchResult in project OpenSearch by opensearch-project.

the class SearchPhaseControllerTests method testConsumerOnlyAggs.

public void testConsumerOnlyAggs() throws Exception {
    int expectedNumResults = randomIntBetween(1, 100);
    int bufferSize = randomIntBetween(2, 200);
    SearchRequest request = randomSearchRequest();
    request.source(new SearchSourceBuilder().aggregation(AggregationBuilders.avg("foo")).size(0));
    request.setBatchedReduceSize(bufferSize);
    QueryPhaseResultConsumer consumer = searchPhaseController.newSearchPhaseResults(fixedExecutor, new NoopCircuitBreaker(CircuitBreaker.REQUEST), SearchProgressListener.NOOP, request, expectedNumResults, exc -> {
    });
    AtomicInteger max = new AtomicInteger();
    CountDownLatch latch = new CountDownLatch(expectedNumResults);
    for (int i = 0; i < expectedNumResults; i++) {
        int number = randomIntBetween(1, 1000);
        max.updateAndGet(prev -> Math.max(prev, number));
        QuerySearchResult result = new QuerySearchResult(new ShardSearchContextId("", i), new SearchShardTarget("node", new ShardId("a", "b", i), null, OriginalIndices.NONE), null);
        result.topDocs(new TopDocsAndMaxScore(new TopDocs(new TotalHits(1, TotalHits.Relation.EQUAL_TO), new ScoreDoc[0]), number), new DocValueFormat[0]);
        InternalAggregations aggs = InternalAggregations.from(Collections.singletonList(new InternalMax("test", (double) number, DocValueFormat.RAW, Collections.emptyMap())));
        result.aggregations(aggs);
        result.setShardIndex(i);
        result.size(1);
        consumer.consumeResult(result, latch::countDown);
    }
    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(0, reduce.sortedTopDocs.scoreDocs.length);
    assertEquals(max.get(), reduce.maxScore, 0.0f);
    assertEquals(expectedNumResults, reduce.totalHits.value);
    assertFalse(reduce.sortedTopDocs.isSortedByField);
    assertNull(reduce.sortedTopDocs.sortFields);
    assertNull(reduce.sortedTopDocs.collapseField);
    assertNull(reduce.sortedTopDocs.collapseValues);
}