Examples with DocMap org.apache.lucene.index.MergeState.DocMap used on opensource projects

Example 1 with DocMap

use of org.apache.lucene.index.MergeState.DocMap in project lucene-solr by apache.

the class MultiSorter method sort.

/** Does a merge sort of the leaves of the incoming reader, returning {@link DocMap} to map each leaf's
   *  documents into the merged segment.  The documents for each incoming leaf reader must already be sorted by the same sort!
   *  Returns null if the merge sort is not needed (segments are already in index sort order).
   **/
static MergeState.DocMap[] sort(Sort sort, List<CodecReader> readers) throws IOException {
    // TODO: optimize if only 1 reader is incoming, though that's a rare case
    SortField[] fields = sort.getSort();
    final ComparableProvider[][] comparables = new ComparableProvider[fields.length][];
    for (int i = 0; i < fields.length; i++) {
        comparables[i] = getComparableProviders(readers, fields[i]);
    }
    int leafCount = readers.size();
    PriorityQueue<LeafAndDocID> queue = new PriorityQueue<LeafAndDocID>(leafCount) {

        @Override
        public boolean lessThan(LeafAndDocID a, LeafAndDocID b) {
            for (int i = 0; i < comparables.length; i++) {
                int cmp = a.values[i].compareTo(b.values[i]);
                if (cmp != 0) {
                    return cmp < 0;
                }
            }
            // tie-break by docID natural order:
            if (a.readerIndex != b.readerIndex) {
                return a.readerIndex < b.readerIndex;
            } else {
                return a.docID < b.docID;
            }
        }
    };
    PackedLongValues.Builder[] builders = new PackedLongValues.Builder[leafCount];
    for (int i = 0; i < leafCount; i++) {
        CodecReader reader = readers.get(i);
        LeafAndDocID leaf = new LeafAndDocID(i, reader.getLiveDocs(), reader.maxDoc(), comparables.length);
        for (int j = 0; j < comparables.length; j++) {
            leaf.values[j] = comparables[j][i].getComparable(leaf.docID);
            assert leaf.values[j] != null;
        }
        queue.add(leaf);
        builders[i] = PackedLongValues.monotonicBuilder(PackedInts.COMPACT);
    }
    // merge sort:
    int mappedDocID = 0;
    int lastReaderIndex = 0;
    boolean isSorted = true;
    while (queue.size() != 0) {
        LeafAndDocID top = queue.top();
        if (lastReaderIndex > top.readerIndex) {
            // merge sort is needed
            isSorted = false;
        }
        lastReaderIndex = top.readerIndex;
        builders[top.readerIndex].add(mappedDocID);
        if (top.liveDocs == null || top.liveDocs.get(top.docID)) {
            mappedDocID++;
        }
        top.docID++;
        if (top.docID < top.maxDoc) {
            for (int j = 0; j < comparables.length; j++) {
                top.values[j] = comparables[j][top.readerIndex].getComparable(top.docID);
                assert top.values[j] != null;
            }
            queue.updateTop();
        } else {
            queue.pop();
        }
    }
    if (isSorted) {
        return null;
    }
    MergeState.DocMap[] docMaps = new MergeState.DocMap[leafCount];
    for (int i = 0; i < leafCount; i++) {
        final PackedLongValues remapped = builders[i].build();
        final Bits liveDocs = readers.get(i).getLiveDocs();
        docMaps[i] = new MergeState.DocMap() {

            @Override
            public int get(int docID) {
                if (liveDocs == null || liveDocs.get(docID)) {
                    return (int) remapped.get(docID);
                } else {
                    return -1;
                }
            }
        };
    }
    return docMaps;
}