ScanContext.java

package org.elasticsearch.search.scan;

import com.google.common.collect.Maps;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.*;
import org.elasticsearch.common.lucene.docset.AllDocSet;
import org.elasticsearch.search.internal.SearchContext;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Map;

/**
 * The scan context allows to optimize readers we already processed during scanning. We do that by keeping track
 * of the count per reader, and if we are done with it, we no longer process it by using a filter that returns
 * null docIdSet for this reader.
 */
public class ScanContext {

    private final Map<IndexReader, ReaderState> readerStates = Maps.newHashMap();

    public void clear() {
        readerStates.clear();
    }

    public TopDocs execute(SearchContext context) throws IOException {
        ScanCollector collector = new ScanCollector(readerStates, context.from(), context.size(), context.trackScores());
        Query query = new FilteredQuery(context.query(), new ScanFilter(readerStates, collector));
        try {
            context.searcher().search(query, collector);
        } catch (ScanCollector.StopCollectingException e) {
            // all is well
        }
        return collector.topDocs();
    }

    static class ScanCollector extends Collector {

        private final Map<IndexReader, ReaderState> readerStates;

        private final int from;

        private final int to;

        private final ArrayList<ScoreDoc> docs;

        private final boolean trackScores;

        private Scorer scorer;

        private int docBase;

        private int counter;

        private IndexReader currentReader;
        private ReaderState readerState;

        ScanCollector(Map<IndexReader, ReaderState> readerStates, int from, int size, boolean trackScores) {
            this.readerStates = readerStates;
            this.from = from;
            this.to = from + size;
            this.trackScores = trackScores;
            this.docs = new ArrayList<ScoreDoc>(size);
        }

        void incCounter(int count) {
            this.counter += count;
        }

        public TopDocs topDocs() {
            return new TopDocs(docs.size(), docs.toArray(new ScoreDoc[docs.size()]), 0f);
        }

        @Override
        public void setScorer(Scorer scorer) throws IOException {
            this.scorer = scorer;
        }

        @Override
        public void collect(int doc) throws IOException {
            if (counter >= from) {
                docs.add(new ScoreDoc(docBase + doc, trackScores ? scorer.score() : 0f));
            }
            readerState.count++;
            counter++;
            if (counter >= to) {
                throw StopCollectingException;
            }
        }

        @Override
        public void setNextReader(IndexReader reader, int docBase) throws IOException {
            // if we have a reader state, and we haven't registered one already, register it
            // we need to check in readersState since even when the filter return null, setNextReader is still
            // called for that reader (before)
            if (currentReader != null && !readerStates.containsKey(currentReader)) {
                assert readerState != null;
                readerState.done = true;
                readerStates.put(currentReader, readerState);
            }
            this.currentReader = reader;
            this.docBase = docBase;
            this.readerState = new ReaderState();
        }

        @Override
        public boolean acceptsDocsOutOfOrder() {
            return true;
        }

        public static final RuntimeException StopCollectingException = new StopCollectingException();

        static class StopCollectingException extends RuntimeException {
            @Override
            public Throwable fillInStackTrace() {
                return null;
            }
        }
    }

    public static class ScanFilter extends Filter {

        private final Map<IndexReader, ReaderState> readerStates;

        private final ScanCollector scanCollector;

        public ScanFilter(Map<IndexReader, ReaderState> readerStates, ScanCollector scanCollector) {
            this.readerStates = readerStates;
            this.scanCollector = scanCollector;
        }

        @Override
        public DocIdSet getDocIdSet(IndexReader reader) throws IOException {
            ReaderState readerState = readerStates.get(reader);
            if (readerState != null && readerState.done) {
                scanCollector.incCounter(readerState.count);
                return null;
            }
            return new AllDocSet(reader.maxDoc());
        }
    }

    static class ReaderState {
        public int count;
        public boolean done;
    }
}