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

Example 21 with TopDocsAndMaxScore

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

the class InternalTopHits method reduce.

@Override
public InternalAggregation reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
    final SearchHits[] shardHits = new SearchHits[aggregations.size()];
    final int from;
    final int size;
    if (reduceContext.isFinalReduce()) {
        from = this.from;
        size = this.size;
    } else {
        // if we are not in the final reduce we need to ensure we maintain all possible elements during reduce
        // hence for pagination we need to maintain all hits until we are in the final phase.
        from = 0;
        size = this.from + this.size;
    }
    final TopDocs reducedTopDocs;
    final TopDocs[] shardDocs;
    if (topDocs.topDocs instanceof TopFieldDocs) {
        Sort sort = new Sort(((TopFieldDocs) topDocs.topDocs).fields);
        shardDocs = new TopFieldDocs[aggregations.size()];
        for (int i = 0; i < shardDocs.length; i++) {
            InternalTopHits topHitsAgg = (InternalTopHits) aggregations.get(i);
            shardDocs[i] = topHitsAgg.topDocs.topDocs;
            shardHits[i] = topHitsAgg.searchHits;
        }
        reducedTopDocs = TopDocs.merge(sort, from, size, (TopFieldDocs[]) shardDocs, true);
    } else {
        shardDocs = new TopDocs[aggregations.size()];
        for (int i = 0; i < shardDocs.length; i++) {
            InternalTopHits topHitsAgg = (InternalTopHits) aggregations.get(i);
            shardDocs[i] = topHitsAgg.topDocs.topDocs;
            shardHits[i] = topHitsAgg.searchHits;
        }
        reducedTopDocs = TopDocs.merge(from, size, shardDocs, true);
    }
    float maxScore = Float.NaN;
    for (InternalAggregation agg : aggregations) {
        InternalTopHits topHitsAgg = (InternalTopHits) agg;
        if (Float.isNaN(topHitsAgg.topDocs.maxScore) == false) {
            if (Float.isNaN(maxScore)) {
                maxScore = topHitsAgg.topDocs.maxScore;
            } else {
                maxScore = Math.max(maxScore, topHitsAgg.topDocs.maxScore);
            }
        }
    }
    final int[] tracker = new int[shardHits.length];
    SearchHit[] hits = new SearchHit[reducedTopDocs.scoreDocs.length];
    for (int i = 0; i < reducedTopDocs.scoreDocs.length; i++) {
        ScoreDoc scoreDoc = reducedTopDocs.scoreDocs[i];
        int position;
        do {
            position = tracker[scoreDoc.shardIndex]++;
        } while (shardDocs[scoreDoc.shardIndex].scoreDocs[position] != scoreDoc);
        hits[i] = shardHits[scoreDoc.shardIndex].getAt(position);
    }
    assert reducedTopDocs.totalHits.relation == Relation.EQUAL_TO;
    return new InternalTopHits(name, this.from, this.size, new TopDocsAndMaxScore(reducedTopDocs, maxScore), new SearchHits(hits, reducedTopDocs.totalHits, maxScore), getMetadata());
}

Example 22 with TopDocsAndMaxScore

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

the class InnerHitsPhase method hitExecute.

private void hitExecute(Map<String, InnerHitsContext.InnerHitSubContext> innerHits, SearchHit hit, SourceLookup rootLookup) throws IOException {
    for (Map.Entry<String, InnerHitsContext.InnerHitSubContext> entry : innerHits.entrySet()) {
        InnerHitsContext.InnerHitSubContext innerHitsContext = entry.getValue();
        TopDocsAndMaxScore topDoc = innerHitsContext.topDocs(hit);
        Map<String, SearchHits> results = hit.getInnerHits();
        if (results == null) {
            hit.setInnerHits(results = new HashMap<>());
        }
        innerHitsContext.queryResult().topDocs(topDoc, innerHitsContext.sort() == null ? null : innerHitsContext.sort().formats);
        int[] docIdsToLoad = new int[topDoc.topDocs.scoreDocs.length];
        for (int j = 0; j < topDoc.topDocs.scoreDocs.length; j++) {
            docIdsToLoad[j] = topDoc.topDocs.scoreDocs[j].doc;
        }
        innerHitsContext.docIdsToLoad(docIdsToLoad, 0, docIdsToLoad.length);
        innerHitsContext.setRootId(new Uid(MapperService.SINGLE_MAPPING_NAME, hit.getId()));
        innerHitsContext.setRootLookup(rootLookup);
        fetchPhase.execute(innerHitsContext);
        FetchSearchResult fetchResult = innerHitsContext.fetchResult();
        SearchHit[] internalHits = fetchResult.fetchResult().hits().getHits();
        for (int j = 0; j < internalHits.length; j++) {
            ScoreDoc scoreDoc = topDoc.topDocs.scoreDocs[j];
            SearchHit searchHitFields = internalHits[j];
            searchHitFields.score(scoreDoc.score);
            if (scoreDoc instanceof FieldDoc) {
                FieldDoc fieldDoc = (FieldDoc) scoreDoc;
                searchHitFields.sortValues(fieldDoc.fields, innerHitsContext.sort().formats);
            }
        }
        results.put(entry.getKey(), fetchResult.hits());
    }
}

Example 23 with TopDocsAndMaxScore

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

the class TopHitsAggregator method buildAggregation.

@Override
public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOException {
    Collectors collectors = topDocsCollectors.get(owningBucketOrdinal);
    if (collectors == null) {
        return buildEmptyAggregation();
    }
    TopDocsCollector<?> topDocsCollector = collectors.topDocsCollector;
    TopDocs topDocs = topDocsCollector.topDocs();
    float maxScore = Float.NaN;
    if (subSearchContext.sort() == null) {
        for (RescoreContext ctx : context().rescore()) {
            try {
                topDocs = ctx.rescorer().rescore(topDocs, context.searcher(), ctx);
            } catch (IOException e) {
                throw new OpenSearchException("Rescore TopHits Failed", e);
            }
        }
        if (topDocs.scoreDocs.length > 0) {
            maxScore = topDocs.scoreDocs[0].score;
        }
    } else if (subSearchContext.trackScores()) {
        TopFieldCollector.populateScores(topDocs.scoreDocs, subSearchContext.searcher(), subSearchContext.query());
        maxScore = collectors.maxScoreCollector.getMaxScore();
    }
    final TopDocsAndMaxScore topDocsAndMaxScore = new TopDocsAndMaxScore(topDocs, maxScore);
    subSearchContext.queryResult().topDocs(topDocsAndMaxScore, subSearchContext.sort() == null ? null : subSearchContext.sort().formats);
    int[] docIdsToLoad = new int[topDocs.scoreDocs.length];
    for (int i = 0; i < topDocs.scoreDocs.length; i++) {
        docIdsToLoad[i] = topDocs.scoreDocs[i].doc;
    }
    subSearchContext.docIdsToLoad(docIdsToLoad, 0, docIdsToLoad.length);
    fetchPhase.execute(subSearchContext);
    FetchSearchResult fetchResult = subSearchContext.fetchResult();
    SearchHit[] internalHits = fetchResult.fetchResult().hits().getHits();
    for (int i = 0; i < internalHits.length; i++) {
        ScoreDoc scoreDoc = topDocs.scoreDocs[i];
        SearchHit searchHitFields = internalHits[i];
        searchHitFields.shard(subSearchContext.shardTarget());
        searchHitFields.score(scoreDoc.score);
        if (scoreDoc instanceof FieldDoc) {
            FieldDoc fieldDoc = (FieldDoc) scoreDoc;
            searchHitFields.sortValues(fieldDoc.fields, subSearchContext.sort().formats);
        }
    }
    return new InternalTopHits(name, subSearchContext.from(), subSearchContext.size(), topDocsAndMaxScore, fetchResult.hits(), metadata());
}

Example 24 with TopDocsAndMaxScore

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

the class QuerySearchResult method consumeTopDocs.

/**
 * Returns and nulls out the top docs for this search results. This allows to free up memory once the top docs are consumed.
 * @throws IllegalStateException if the top docs have already been consumed.
 */
public TopDocsAndMaxScore consumeTopDocs() {
    TopDocsAndMaxScore topDocsAndMaxScore = this.topDocsAndMaxScore;
    if (topDocsAndMaxScore == null) {
        throw new IllegalStateException("topDocs already consumed");
    }
    this.topDocsAndMaxScore = null;
    return topDocsAndMaxScore;
}

Example 25 with TopDocsAndMaxScore

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

the class ContextIndexSearcher method search.

public void search(List<LeafReaderContext> leaves, Weight weight, CollectorManager manager, QuerySearchResult result, DocValueFormat[] formats, TotalHits totalHits) throws IOException {
    final List<Collector> collectors = new ArrayList<>(leaves.size());
    for (LeafReaderContext ctx : leaves) {
        final Collector collector = manager.newCollector();
        searchLeaf(ctx, weight, collector);
        collectors.add(collector);
    }
    TopFieldDocs mergedTopDocs = (TopFieldDocs) manager.reduce(collectors);
    // We need to reset shard index; OpenSearch will set shard index later during reduce stage
    for (ScoreDoc scoreDoc : mergedTopDocs.scoreDocs) {
        scoreDoc.shardIndex = -1;
    }
    if (totalHits != null) {
        // we have already precalculated totalHits for the whole index
        mergedTopDocs = new TopFieldDocs(totalHits, mergedTopDocs.scoreDocs, mergedTopDocs.fields);
    }
    result.topDocs(new TopDocsAndMaxScore(mergedTopDocs, Float.NaN), formats);
}