Example usage for org.apache.lucene.search TopDocs merge

List of usage examples for org.apache.lucene.search TopDocs merge

Introduction

In this page you can find the example usage for org.apache.lucene.search TopDocs merge.

Prototype

public static TopDocs merge(int start, int topN, TopDocs[] shardHits, boolean setShardIndex) 

Source Link

Document

Same as #merge(int,TopDocs[]) but also ignores the top start top docs.

Usage

From source file:io.crate.action.sql.query.CrateResultSorter.java

License:Apache License

/**
 * copied from SearchPhaseController, to manually set offset
 *
 * @param resultsArr Shard result holder
 * @param offset the number of results to skip
 * @param limit the number of results to return at max
 *///from www  .jav a 2s.c  o  m
public ScoreDoc[] sortDocs(AtomicArray<? extends QuerySearchResultProvider> resultsArr, int offset, int limit)
        throws IOException {
    List<? extends AtomicArray.Entry<? extends QuerySearchResultProvider>> results = resultsArr.asList();
    if (results.isEmpty()) {
        return SearchPhaseController.EMPTY_DOCS;
    }

    if (optimizeSingleShard) {
        boolean canOptimize = false;
        QuerySearchResult result = null;
        int shardIndex = -1;
        if (results.size() == 1) {
            canOptimize = true;
            result = results.get(0).value.queryResult();
            shardIndex = results.get(0).index;
        } else {
            // 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().topDocs().scoreDocs.length > 0) {
                    if (result != null) { // we already have one, can't really optimize
                        canOptimize = false;
                        break;
                    }
                    canOptimize = true;
                    result = entry.value.queryResult();
                    shardIndex = entry.index;
                }
            }
        }
        if (canOptimize) {
            ScoreDoc[] scoreDocs = result.topDocs().scoreDocs;
            if (scoreDocs.length == 0 || scoreDocs.length < offset) {
                return SearchPhaseController.EMPTY_DOCS;
            }

            int resultDocsSize = result.size();
            if ((scoreDocs.length - offset) < resultDocsSize) {
                resultDocsSize = scoreDocs.length - offset;
            }
            ScoreDoc[] docs = new ScoreDoc[resultDocsSize];
            for (int i = 0; i < resultDocsSize; i++) {
                ScoreDoc scoreDoc = scoreDocs[offset + i];
                scoreDoc.shardIndex = shardIndex;
                docs[i] = scoreDoc;
            }
            return docs;
        }
    }

    @SuppressWarnings("unchecked")
    AtomicArray.Entry<? extends QuerySearchResultProvider>[] sortedResults = results
            .toArray(new AtomicArray.Entry[results.size()]);
    Arrays.sort(sortedResults, SearchPhaseController.QUERY_RESULT_ORDERING);
    QuerySearchResultProvider firstResult = sortedResults[0].value;

    final Sort sort;
    if (firstResult.queryResult().topDocs() instanceof TopFieldDocs) {
        TopFieldDocs firstTopDocs = (TopFieldDocs) firstResult.queryResult().topDocs();
        sort = new Sort(firstTopDocs.fields);
    } else {
        sort = null;
    }

    // Need to use the length of the resultsArr array, since the slots will be based on the position in the resultsArr array
    TopDocs[] shardTopDocs = new TopDocs[resultsArr.length()];
    if (firstResult.includeFetch()) {
        // if we did both query and fetch on the same go, we have fetched all the docs from each shards already, use them...
        // this is also important since we shortcut and fetch only docs from "from" and up to "size"
        limit *= sortedResults.length;
    }
    for (AtomicArray.Entry<? extends QuerySearchResultProvider> sortedResult : sortedResults) {
        TopDocs topDocs = sortedResult.value.queryResult().topDocs();
        // the 'index' field is the position in the resultsArr atomic array
        shardTopDocs[sortedResult.index] = topDocs;
    }
    // TopDocs#merge can't deal with null shard TopDocs
    for (int i = 0; i < shardTopDocs.length; i++) {
        if (shardTopDocs[i] == null) {
            shardTopDocs[i] = Lucene.EMPTY_TOP_DOCS;
        }
    }
    TopDocs mergedTopDocs = TopDocs.merge(sort, offset, limit, shardTopDocs);
    return mergedTopDocs.scoreDocs;
}

From source file:org.elasticsearch.search.aggregations.bucket.tophits.InternalTopHits.java

License:Apache License

@Override
public InternalAggregation reduce(ReduceContext reduceContext) {
    List<InternalAggregation> aggregations = reduceContext.aggregations();
    if (aggregations.size() == 1 && from == 0) {
        return aggregations.get(0);
    }//w  w  w. j  a  v a 2s.c  o  m

    TopDocs[] shardDocs = new TopDocs[aggregations.size()];
    InternalSearchHits[] shardHits = new InternalSearchHits[aggregations.size()];
    for (int i = 0; i < shardDocs.length; i++) {
        InternalTopHits topHitsAgg = (InternalTopHits) aggregations.get(i);
        shardDocs[i] = topHitsAgg.topDocs;
        shardHits[i] = topHitsAgg.searchHits;
    }

    try {
        int[] tracker = new int[shardHits.length];
        TopDocs reducedTopDocs = TopDocs.merge(sort, from, size, shardDocs);
        InternalSearchHit[] hits = new InternalSearchHit[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] = (InternalSearchHit) shardHits[scoreDoc.shardIndex].getAt(position);
        }
        return new InternalTopHits(name,
                new InternalSearchHits(hits, reducedTopDocs.totalHits, reducedTopDocs.getMaxScore()));
    } catch (IOException e) {
        throw ExceptionsHelper.convertToElastic(e);
    }
}

From source file:org.elasticsearch.search.aggregations.metrics.tophits.InternalTopHits.java

License:Apache License

@Override
public InternalAggregation doReduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
    InternalSearchHits[] shardHits = new InternalSearchHits[aggregations.size()];

    final TopDocs reducedTopDocs;
    final TopDocs[] shardDocs;

    try {/*from www . j a  va  2s  .  c  o  m*/
        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] = (TopFieldDocs) topHitsAgg.topDocs;
                shardHits[i] = topHitsAgg.searchHits;
            }
            reducedTopDocs = TopDocs.merge(sort, from, size, (TopFieldDocs[]) shardDocs);
        } 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);
        }

        final int[] tracker = new int[shardHits.length];
        InternalSearchHit[] hits = new InternalSearchHit[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] = (InternalSearchHit) shardHits[scoreDoc.shardIndex].getAt(position);
        }
        return new InternalTopHits(name, from, size, reducedTopDocs,
                new InternalSearchHits(hits, reducedTopDocs.totalHits, reducedTopDocs.getMaxScore()),
                pipelineAggregators(), getMetaData());
    } catch (IOException e) {
        throw ExceptionsHelper.convertToElastic(e);
    }
}

From source file:org.meresco.lucene.search.TopDocSuperCollector.java

License:Open Source License

public TopDocs topDocs(int start) throws IOException {
    TopDocs[] topdocs = new TopDocs[this.subs.size()];
    for (int i = 0; i < topdocs.length; i++)
        topdocs[i] = this.subs.get(i).topdocs;
    return TopDocs.merge(this.sort, start, this.numHits - start, topdocs);
}

From source file:sh.isaac.provider.query.lucene.IsaacFilteredCollectorManager.java

License:Apache License

@Override
public TopDocs reduce(Collection<IsaacFilteredCollectorManager.IsaacFilteredCollector> collectors)
        throws IOException {
    final TopDocs[] topDocs = new TopDocs[collectors.size()];
    int i = 0;/*from   www . j av  a2 s .c o  m*/
    int totalSize = 0;
    for (IsaacFilteredCollectorManager.IsaacFilteredCollector collector : collectors) {
        topDocs[i] = collector.topDocs();
        totalSize = totalSize + topDocs[i++].scoreDocs.length;
    }
    return TopDocs.merge(0, Math.min(sizeLimit, totalSize), topDocs, true);
}