org.elasticsearch.search.scan.ScanContext Maven / Gradle / Ivy
package org.elasticsearch.search.scan;
/*
* Licensed to ElasticSearch and Shay Banon under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. ElasticSearch licenses this
* file to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
import com.google.common.collect.Maps;
import org.apache.lucene.index.AtomicReaderContext;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.*;
import org.apache.lucene.util.Bits;
import org.elasticsearch.common.lucene.docset.AllDocIdSet;
import org.elasticsearch.common.lucene.search.XFilteredQuery;
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 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 XFilteredQuery(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 readerStates;
private final int from;
private final int to;
private final ArrayList docs;
private final boolean trackScores;
private Scorer scorer;
private int docBase;
private int counter;
private IndexReader currentReader;
private ReaderState readerState;
ScanCollector(Map readerStates, int from, int size, boolean trackScores) {
this.readerStates = readerStates;
this.from = from;
this.to = from + size;
this.trackScores = trackScores;
this.docs = new ArrayList(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(AtomicReaderContext context) 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 = context.reader();
this.docBase = context.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 readerStates;
private final ScanCollector scanCollector;
public ScanFilter(Map readerStates, ScanCollector scanCollector) {
this.readerStates = readerStates;
this.scanCollector = scanCollector;
}
@Override
public DocIdSet getDocIdSet(AtomicReaderContext context, Bits acceptedDocs) throws IOException {
ReaderState readerState = readerStates.get(context.reader());
if (readerState != null && readerState.done) {
scanCollector.incCounter(readerState.count);
return null;
}
return new AllDocIdSet(context.reader().maxDoc());
}
}
static class ReaderState {
public int count;
public boolean done;
}
}