Examples with TopDocsAndMaxScore org.opensearch.common.lucene.search.TopDocsAndMaxScore used on opensource projects

Example 1 with TopDocsAndMaxScore

use of org.opensearch.common.lucene.search.TopDocsAndMaxScore in project OpenSearch by opensearch-project.

the class SearchResponseMerger method getMergedResponse.

/**
 * Returns the merged response. To be called once all responses have been added through {@link #add(SearchResponse)}
 * so that all responses are merged into a single one.
 */
SearchResponse getMergedResponse(SearchResponse.Clusters clusters) {
    // we end up calling merge without anything to merge, we just return an empty search response
    if (searchResponses.size() == 0) {
        return SearchResponse.empty(searchTimeProvider::buildTookInMillis, clusters);
    }
    int totalShards = 0;
    int skippedShards = 0;
    int successfulShards = 0;
    // the current reduce phase counts as one
    int numReducePhases = 1;
    List<ShardSearchFailure> failures = new ArrayList<>();
    Map<String, ProfileShardResult> profileResults = new HashMap<>();
    List<InternalAggregations> aggs = new ArrayList<>();
    Map<ShardIdAndClusterAlias, Integer> shards = new TreeMap<>();
    List<TopDocs> topDocsList = new ArrayList<>(searchResponses.size());
    Map<String, List<Suggest.Suggestion>> groupedSuggestions = new HashMap<>();
    Boolean trackTotalHits = null;
    SearchPhaseController.TopDocsStats topDocsStats = new SearchPhaseController.TopDocsStats(trackTotalHitsUpTo);
    for (SearchResponse searchResponse : searchResponses) {
        totalShards += searchResponse.getTotalShards();
        skippedShards += searchResponse.getSkippedShards();
        successfulShards += searchResponse.getSuccessfulShards();
        numReducePhases += searchResponse.getNumReducePhases();
        Collections.addAll(failures, searchResponse.getShardFailures());
        profileResults.putAll(searchResponse.getProfileResults());
        if (searchResponse.getAggregations() != null) {
            InternalAggregations internalAggs = (InternalAggregations) searchResponse.getAggregations();
            aggs.add(internalAggs);
        }
        Suggest suggest = searchResponse.getSuggest();
        if (suggest != null) {
            for (Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> entries : suggest) {
                List<Suggest.Suggestion> suggestionList = groupedSuggestions.computeIfAbsent(entries.getName(), s -> new ArrayList<>());
                suggestionList.add(entries);
            }
            List<CompletionSuggestion> completionSuggestions = suggest.filter(CompletionSuggestion.class);
            for (CompletionSuggestion completionSuggestion : completionSuggestions) {
                for (CompletionSuggestion.Entry options : completionSuggestion) {
                    for (CompletionSuggestion.Entry.Option option : options) {
                        SearchShardTarget shard = option.getHit().getShard();
                        ShardIdAndClusterAlias shardId = new ShardIdAndClusterAlias(shard.getShardId(), shard.getClusterAlias());
                        shards.putIfAbsent(shardId, null);
                    }
                }
            }
        }
        SearchHits searchHits = searchResponse.getHits();
        final TotalHits totalHits;
        if (searchHits.getTotalHits() == null) {
            // in case we didn't track total hits, we get null from each cluster, but we need to set 0 eq to the TopDocs
            totalHits = new TotalHits(0, TotalHits.Relation.EQUAL_TO);
            assert trackTotalHits == null || trackTotalHits == false;
            trackTotalHits = false;
        } else {
            totalHits = searchHits.getTotalHits();
            assert trackTotalHits == null || trackTotalHits;
            trackTotalHits = true;
        }
        TopDocs topDocs = searchHitsToTopDocs(searchHits, totalHits, shards);
        topDocsStats.add(new TopDocsAndMaxScore(topDocs, searchHits.getMaxScore()), searchResponse.isTimedOut(), searchResponse.isTerminatedEarly());
        if (searchHits.getHits().length > 0) {
            // there is no point in adding empty search hits and merging them with the others. Also, empty search hits always come
            // without sort fields and collapse info, despite sort by field and/or field collapsing was requested, which causes
            // issues reconstructing the proper TopDocs instance and breaks mergeTopDocs which expects the same type for each result.
            topDocsList.add(topDocs);
        }
    }
    // after going through all the hits and collecting all their distinct shards, we assign shardIndex and set it to the ScoreDocs
    setTopDocsShardIndex(shards, topDocsList);
    TopDocs topDocs = SearchPhaseController.mergeTopDocs(topDocsList, size, from);
    SearchHits mergedSearchHits = topDocsToSearchHits(topDocs, topDocsStats);
    setSuggestShardIndex(shards, groupedSuggestions);
    Suggest suggest = groupedSuggestions.isEmpty() ? null : new Suggest(Suggest.reduce(groupedSuggestions));
    InternalAggregations reducedAggs = InternalAggregations.topLevelReduce(aggs, aggReduceContextBuilder.forFinalReduction());
    ShardSearchFailure[] shardFailures = failures.toArray(ShardSearchFailure.EMPTY_ARRAY);
    SearchProfileShardResults profileShardResults = profileResults.isEmpty() ? null : new SearchProfileShardResults(profileResults);
    // make failures ordering consistent between ordinary search and CCS by looking at the shard they come from
    Arrays.sort(shardFailures, FAILURES_COMPARATOR);
    InternalSearchResponse response = new InternalSearchResponse(mergedSearchHits, reducedAggs, suggest, profileShardResults, topDocsStats.timedOut, topDocsStats.terminatedEarly, numReducePhases);
    long tookInMillis = searchTimeProvider.buildTookInMillis();
    return new SearchResponse(response, null, totalShards, successfulShards, skippedShards, tookInMillis, shardFailures, clusters, null);
}

Example 2 with TopDocsAndMaxScore

use of org.opensearch.common.lucene.search.TopDocsAndMaxScore in project OpenSearch by opensearch-project.

the class Lucene method readTopDocs.

public static TopDocsAndMaxScore readTopDocs(StreamInput in) throws IOException {
    byte type = in.readByte();
    if (type == 0) {
        TotalHits totalHits = readTotalHits(in);
        float maxScore = in.readFloat();
        final int scoreDocCount = in.readVInt();
        final ScoreDoc[] scoreDocs;
        if (scoreDocCount == 0) {
            scoreDocs = EMPTY_SCORE_DOCS;
        } else {
            scoreDocs = new ScoreDoc[scoreDocCount];
            for (int i = 0; i < scoreDocs.length; i++) {
                scoreDocs[i] = new ScoreDoc(in.readVInt(), in.readFloat());
            }
        }
        return new TopDocsAndMaxScore(new TopDocs(totalHits, scoreDocs), maxScore);
    } else if (type == 1) {
        TotalHits totalHits = readTotalHits(in);
        float maxScore = in.readFloat();
        SortField[] fields = in.readArray(Lucene::readSortField, SortField[]::new);
        FieldDoc[] fieldDocs = new FieldDoc[in.readVInt()];
        for (int i = 0; i < fieldDocs.length; i++) {
            fieldDocs[i] = readFieldDoc(in);
        }
        return new TopDocsAndMaxScore(new TopFieldDocs(totalHits, fieldDocs, fields), maxScore);
    } else if (type == 2) {
        TotalHits totalHits = readTotalHits(in);
        float maxScore = in.readFloat();
        String field = in.readString();
        SortField[] fields = in.readArray(Lucene::readSortField, SortField[]::new);
        int size = in.readVInt();
        Object[] collapseValues = new Object[size];
        FieldDoc[] fieldDocs = new FieldDoc[size];
        for (int i = 0; i < fieldDocs.length; i++) {
            fieldDocs[i] = readFieldDoc(in);
            collapseValues[i] = readSortValue(in);
        }
        return new TopDocsAndMaxScore(new CollapseTopFieldDocs(field, totalHits, fieldDocs, fields, collapseValues), maxScore);
    } else {
        throw new IllegalStateException("Unknown type " + type);
    }
}

Example 3 with TopDocsAndMaxScore

use of org.opensearch.common.lucene.search.TopDocsAndMaxScore 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 4 with TopDocsAndMaxScore

use of org.opensearch.common.lucene.search.TopDocsAndMaxScore 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 5 with TopDocsAndMaxScore

use of org.opensearch.common.lucene.search.TopDocsAndMaxScore 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;
}