Examples with TopFieldDocs org.apache.lucene.search.TopFieldDocs used on opensource projects

Example 21 with TopFieldDocs

use of org.apache.lucene.search.TopFieldDocs in project lucene-solr by apache.

the class TopGroupsResultTransformer method transformToNative.

/**
   * {@inheritDoc}
   */
@Override
public Map<String, ?> transformToNative(NamedList<NamedList> shardResponse, Sort groupSort, Sort withinGroupSort, String shard) {
    Map<String, Object> result = new HashMap<>();
    final IndexSchema schema = rb.req.getSearcher().getSchema();
    for (Map.Entry<String, NamedList> entry : shardResponse) {
        String key = entry.getKey();
        NamedList commandResult = entry.getValue();
        Integer totalGroupedHitCount = (Integer) commandResult.get("totalGroupedHitCount");
        Integer totalHits = (Integer) commandResult.get("totalHits");
        if (totalHits != null) {
            Integer matches = (Integer) commandResult.get("matches");
            Float maxScore = (Float) commandResult.get("maxScore");
            if (maxScore == null) {
                maxScore = Float.NaN;
            }
            @SuppressWarnings("unchecked") List<NamedList<Object>> documents = (List<NamedList<Object>>) commandResult.get("documents");
            ScoreDoc[] scoreDocs = transformToNativeShardDoc(documents, groupSort, shard, schema);
            final TopDocs topDocs;
            if (withinGroupSort.equals(Sort.RELEVANCE)) {
                topDocs = new TopDocs(totalHits, scoreDocs, maxScore);
            } else {
                topDocs = new TopFieldDocs(totalHits, scoreDocs, withinGroupSort.getSort(), maxScore);
            }
            result.put(key, new QueryCommandResult(topDocs, matches));
            continue;
        }
        Integer totalHitCount = (Integer) commandResult.get("totalHitCount");
        List<GroupDocs<BytesRef>> groupDocs = new ArrayList<>();
        for (int i = 2; i < commandResult.size(); i++) {
            String groupValue = commandResult.getName(i);
            @SuppressWarnings("unchecked") NamedList<Object> groupResult = (NamedList<Object>) commandResult.getVal(i);
            Integer totalGroupHits = (Integer) groupResult.get("totalHits");
            Float maxScore = (Float) groupResult.get("maxScore");
            if (maxScore == null) {
                maxScore = Float.NaN;
            }
            @SuppressWarnings("unchecked") List<NamedList<Object>> documents = (List<NamedList<Object>>) groupResult.get("documents");
            ScoreDoc[] scoreDocs = transformToNativeShardDoc(documents, withinGroupSort, shard, schema);
            BytesRef groupValueRef = groupValue != null ? new BytesRef(groupValue) : null;
            groupDocs.add(new GroupDocs<>(Float.NaN, maxScore, totalGroupHits, scoreDocs, groupValueRef, null));
        }
        @SuppressWarnings("unchecked") GroupDocs<BytesRef>[] groupDocsArr = groupDocs.toArray(new GroupDocs[groupDocs.size()]);
        TopGroups<BytesRef> topGroups = new TopGroups<>(groupSort.getSort(), withinGroupSort.getSort(), totalHitCount, totalGroupedHitCount, groupDocsArr, Float.NaN);
        result.put(key, topGroups);
    }
    return result;
}

Example 22 with TopFieldDocs

use of org.apache.lucene.search.TopFieldDocs 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;
}

Example 23 with TopFieldDocs

use of org.apache.lucene.search.TopFieldDocs in project elasticsearch by elastic.

the class InternalTopHits method doReduce.

@Override
public InternalAggregation doReduce(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 instanceof TopFieldDocs) {
        Sort sort = new Sort(((TopFieldDocs) 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;
            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;
            shardHits[i] = topHitsAgg.searchHits;
        }
        reducedTopDocs = TopDocs.merge(from, size, shardDocs, true);
    }
    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);
    }
    return new InternalTopHits(name, this.from, this.size, reducedTopDocs, new SearchHits(hits, reducedTopDocs.totalHits, reducedTopDocs.getMaxScore()), pipelineAggregators(), getMetaData());
}

Example 24 with TopFieldDocs

use of org.apache.lucene.search.TopFieldDocs in project elasticsearch by elastic.

the class Lucene method writeTopDocs.

public static void writeTopDocs(StreamOutput out, TopDocs topDocs) throws IOException {
    if (topDocs instanceof CollapseTopFieldDocs) {
        out.writeByte((byte) 2);
        CollapseTopFieldDocs collapseDocs = (CollapseTopFieldDocs) topDocs;
        out.writeVInt(topDocs.totalHits);
        out.writeFloat(topDocs.getMaxScore());
        out.writeString(collapseDocs.field);
        out.writeVInt(collapseDocs.fields.length);
        for (SortField sortField : collapseDocs.fields) {
            writeSortField(out, sortField);
        }
        out.writeVInt(topDocs.scoreDocs.length);
        for (int i = 0; i < topDocs.scoreDocs.length; i++) {
            ScoreDoc doc = collapseDocs.scoreDocs[i];
            writeFieldDoc(out, (FieldDoc) doc);
            writeSortValue(out, collapseDocs.collapseValues[i]);
        }
    } else if (topDocs instanceof TopFieldDocs) {
        out.writeByte((byte) 1);
        TopFieldDocs topFieldDocs = (TopFieldDocs) topDocs;
        out.writeVInt(topDocs.totalHits);
        out.writeFloat(topDocs.getMaxScore());
        out.writeVInt(topFieldDocs.fields.length);
        for (SortField sortField : topFieldDocs.fields) {
            writeSortField(out, sortField);
        }
        out.writeVInt(topDocs.scoreDocs.length);
        for (ScoreDoc doc : topFieldDocs.scoreDocs) {
            writeFieldDoc(out, (FieldDoc) doc);
        }
    } else {
        out.writeByte((byte) 0);
        out.writeVInt(topDocs.totalHits);
        out.writeFloat(topDocs.getMaxScore());
        out.writeVInt(topDocs.scoreDocs.length);
        for (ScoreDoc doc : topDocs.scoreDocs) {
            writeScoreDoc(out, doc);
        }
    }
}

Example 25 with TopFieldDocs

use of org.apache.lucene.search.TopFieldDocs in project elasticsearch by elastic.

the class Lucene method readTopDocs.

public static TopDocs readTopDocs(StreamInput in) throws IOException {
    byte type = in.readByte();
    if (type == 0) {
        int totalHits = in.readVInt();
        float maxScore = in.readFloat();
        ScoreDoc[] scoreDocs = new ScoreDoc[in.readVInt()];
        for (int i = 0; i < scoreDocs.length; i++) {
            scoreDocs[i] = new ScoreDoc(in.readVInt(), in.readFloat());
        }
        return new TopDocs(totalHits, scoreDocs, maxScore);
    } else if (type == 1) {
        int totalHits = in.readVInt();
        float maxScore = in.readFloat();
        SortField[] fields = new SortField[in.readVInt()];
        for (int i = 0; i < fields.length; i++) {
            fields[i] = readSortField(in);
        }
        FieldDoc[] fieldDocs = new FieldDoc[in.readVInt()];
        for (int i = 0; i < fieldDocs.length; i++) {
            fieldDocs[i] = readFieldDoc(in);
        }
        return new TopFieldDocs(totalHits, fieldDocs, fields, maxScore);
    } else if (type == 2) {
        int totalHits = in.readVInt();
        float maxScore = in.readFloat();
        String field = in.readString();
        SortField[] fields = new SortField[in.readVInt()];
        for (int i = 0; i < fields.length; i++) {
            fields[i] = readSortField(in);
        }
        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 CollapseTopFieldDocs(field, totalHits, fieldDocs, fields, collapseValues, maxScore);
    } else {
        throw new IllegalStateException("Unknown type " + type);
    }
}