honeymoose
/
OpenSearch
mirror of https://github.com/honeymoose/OpenSearch.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Move early termination based on index sort to TopDocs collector (#27666) 

Lucene TopDocs collector are now able to early terminate the collection
based on the index sort. This change plugs this new functionality directly in the
query phase instead of relying on a dedicated early terminating collector.
This commit is contained in:
Jim Ferenczi 2017-12-21 08:57:06 +01:00 committed by GitHub
parent 7d88851766
commit 5ac5fd95ae
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 197 additions and 239 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
core/src/main/java/org/apache/lucene/queries/SearchAfterSortedDocQuery.java
View File

@ -79,7 +79,7 @@ public class SearchAfterSortedDocQuery extends Query {
throw new IOException("search sort :[" + sort.getSort() + "] does not match the index sort:[" + segmentSort + "]");
}
final int afterDoc = after.doc - context.docBase;
TopComparator comparator= getTopComparator(fieldComparators, reverseMuls, context, afterDoc);
TopComparator comparator = getTopComparator(fieldComparators, reverseMuls, context, afterDoc);
final int maxDoc = context.reader().maxDoc();
final int firstDoc = searchAfterDoc(comparator, 0, context.reader().maxDoc());
if (firstDoc >= maxDoc) {
@ -143,7 +143,7 @@ public class SearchAfterSortedDocQuery extends Query {
}
}
if (topDoc <= doc) {
if (doc <= topDoc) {
return false;
}
return true;

69
core/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java
View File

@ -19,15 +19,10 @@
package org.elasticsearch.search.query;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.search.Collector;
import org.apache.lucene.search.EarlyTerminatingSortingCollector;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.MultiCollector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.TotalHitCountCollector;
import org.apache.lucene.search.Weight;
import org.elasticsearch.common.lucene.MinimumScoreCollector;
import org.elasticsearch.common.lucene.search.FilteredCollector;
@ -40,14 +35,12 @@ import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.function.BooleanSupplier;
import java.util.function.IntSupplier;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_CANCELLED;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MIN_SCORE;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MULTI;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_POST_FILTER;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TERMINATE_AFTER_COUNT;
import static org.elasticsearch.search.query.TopDocsCollectorContext.shortcutTotalHitCount;
abstract class QueryCollectorContext {
private String profilerName;
@ -70,22 +63,12 @@ abstract class QueryCollectorContext {
return new InternalProfileCollector(collector, profilerName, in != null ? Collections.singletonList(in) : Collections.emptyList());
}
/**
* A value of <code>false</code> indicates that the underlying collector can infer
* its results directly from the context (search is not needed).
* Default to true (search is needed).
*/
boolean shouldCollect() {
return true;
}
/**
* Post-process <code>result</code> after search execution.
*
* @param result The query search result to populate
* @param hasCollected True if search was executed
*/
void postProcess(QuerySearchResult result, boolean hasCollected) throws IOException {}
void postProcess(QuerySearchResult result) throws IOException {}
/**
* Creates the collector tree from the provided <code>collectors</code>
@ -175,11 +158,6 @@ abstract class QueryCollectorContext {
Collector create(Collector in) throws IOException {
return new CancellableCollector(cancelled, in);
}
@Override
boolean shouldCollect() {
return false;
}
};
}
@ -198,52 +176,11 @@ abstract class QueryCollectorContext {
}
@Override
void postProcess(QuerySearchResult result, boolean hasCollected) throws IOException {
if (hasCollected && collector.terminatedEarly()) {
void postProcess(QuerySearchResult result) throws IOException {
if (collector.terminatedEarly()) {
result.terminatedEarly(true);
}
}
};
}
/**
* Creates a sorting termination collector limiting the collection to the first <code>numHits</code> per segment.
* The total hit count matching the query is also computed if <code>trackTotalHits</code> is true.
*/
static QueryCollectorContext createEarlySortingTerminationCollectorContext(IndexReader reader,
Query query,
Sort indexSort,
int numHits,
boolean trackTotalHits,
boolean shouldCollect) {
return new QueryCollectorContext(REASON_SEARCH_TERMINATE_AFTER_COUNT) {
private IntSupplier countSupplier = null;
@Override
Collector create(Collector in) throws IOException {
EarlyTerminatingSortingCollector sortingCollector = new EarlyTerminatingSortingCollector(in, indexSort, numHits);
Collector collector = sortingCollector;
if (trackTotalHits) {
int count = shouldCollect ? -1 : shortcutTotalHitCount(reader, query);
if (count == -1) {
TotalHitCountCollector countCollector = new TotalHitCountCollector();
collector = MultiCollector.wrap(sortingCollector, countCollector);
this.countSupplier = countCollector::getTotalHits;
} else {
this.countSupplier = () -> count;
}
}
return collector;
}
@Override
void postProcess(QuerySearchResult result, boolean hasCollected) throws IOException {
if (countSupplier != null) {
final TopDocs topDocs = result.topDocs();
topDocs.totalHits = countSupplier.getAsInt();
result.topDocs(topDocs, result.sortValueFormats());
}
}
};
}
}

66
core/src/main/java/org/elasticsearch/search/query/QueryPhase.java
View File

@ -20,6 +20,7 @@
package org.elasticsearch.search.query;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.queries.MinDocQuery;
import org.apache.lucene.queries.SearchAfterSortedDocQuery;
import org.apache.lucene.search.BooleanClause;
@ -36,7 +37,6 @@ import org.apache.lucene.search.Sort;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.util.Counter;
import org.elasticsearch.action.search.SearchTask;
import org.elasticsearch.common.Nullable;
import org.elasticsearch.common.lucene.Lucene;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.concurrent.EsThreadPoolExecutor;
@ -61,7 +61,6 @@ import java.util.LinkedList;
import java.util.function.Consumer;
import static org.elasticsearch.search.query.QueryCollectorContext.createCancellableCollectorContext;
import static org.elasticsearch.search.query.QueryCollectorContext.createEarlySortingTerminationCollectorContext;
import static org.elasticsearch.search.query.QueryCollectorContext.createEarlyTerminationCollectorContext;
import static org.elasticsearch.search.query.QueryCollectorContext.createFilteredCollectorContext;
import static org.elasticsearch.search.query.QueryCollectorContext.createMinScoreCollectorContext;
@ -104,10 +103,8 @@ public class QueryPhase implements SearchPhase {
// request, preProcess is called on the DFS phase phase, this is why we pre-process them
// here to make sure it happens during the QUERY phase
aggregationPhase.preProcess(searchContext);
Sort indexSort = searchContext.mapperService().getIndexSettings().getIndexSortConfig()
.buildIndexSort(searchContext.mapperService()::fullName, searchContext::getForField);
final ContextIndexSearcher searcher = searchContext.searcher();
boolean rescore = execute(searchContext, searchContext.searcher(), searcher::setCheckCancelled, indexSort);
boolean rescore = execute(searchContext, searchContext.searcher(), searcher::setCheckCancelled);
if (rescore) { // only if we do a regular search
rescorePhase.execute(searchContext);
@ -127,11 +124,12 @@ public class QueryPhase implements SearchPhase {
* wire everything (mapperService, etc.)
* @return whether the rescoring phase should be executed
*/
static boolean execute(SearchContext searchContext, final IndexSearcher searcher,
Consumer<Runnable> checkCancellationSetter, @Nullable Sort indexSort) throws QueryPhaseExecutionException {
static boolean execute(SearchContext searchContext,
final IndexSearcher searcher,
Consumer<Runnable> checkCancellationSetter) throws QueryPhaseExecutionException {
final IndexReader reader = searcher.getIndexReader();
QuerySearchResult queryResult = searchContext.queryResult();
queryResult.searchTimedOut(false);
try {
queryResult.from(searchContext.from());
queryResult.size(searchContext.size());
@ -161,7 +159,7 @@ public class QueryPhase implements SearchPhase {
// ... and stop collecting after ${size} matches
searchContext.terminateAfter(searchContext.size());
searchContext.trackTotalHits(false);
} else if (canEarlyTerminate(indexSort, searchContext)) {
} else if (canEarlyTerminate(reader, searchContext.sort())) {
// now this gets interesting: since the search sort is a prefix of the index sort, we can directly
// skip to the desired doc
if (after != null) {
@ -177,10 +175,14 @@ public class QueryPhase implements SearchPhase {
}
final LinkedList<QueryCollectorContext> collectors = new LinkedList<>();
// whether the chain contains a collector that filters documents
boolean hasFilterCollector = false;
if (searchContext.parsedPostFilter() != null) {
// add post filters before aggregations
// it will only be applied to top hits
collectors.add(createFilteredCollectorContext(searcher, searchContext.parsedPostFilter().query()));
// this collector can filter documents during the collection
hasFilterCollector = true;
}
if (searchContext.queryCollectors().isEmpty() == false) {
// plug in additional collectors, like aggregations
@ -189,10 +191,14 @@ public class QueryPhase implements SearchPhase {
if (searchContext.minimumScore() != null) {
// apply the minimum score after multi collector so we filter aggs as well
collectors.add(createMinScoreCollectorContext(searchContext.minimumScore()));
// this collector can filter documents during the collection
hasFilterCollector = true;
}
if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) {
// apply terminate after after all filters collectors
collectors.add(createEarlyTerminationCollectorContext(searchContext.terminateAfter()));
// this collector can filter documents during the collection
hasFilterCollector = true;
}
boolean timeoutSet = scrollContext == null && searchContext.timeout() != null &&
@ -240,21 +246,9 @@ public class QueryPhase implements SearchPhase {
// searchContext.lowLevelCancellation()
collectors.add(createCancellableCollectorContext(searchContext.getTask()::isCancelled));
final IndexReader reader = searcher.getIndexReader();
final boolean doProfile = searchContext.getProfilers() != null;
// create the top docs collector last when the other collectors are known
final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, reader,
collectors.stream().anyMatch(QueryCollectorContext::shouldCollect));
final boolean shouldCollect = topDocsFactory.shouldCollect();
if (topDocsFactory.numHits() > 0 &&
(scrollContext == null || scrollContext.totalHits != -1) &&
canEarlyTerminate(indexSort, searchContext)) {
// top docs collection can be early terminated based on index sort
// add the collector context first so we don't early terminate aggs but only top docs
collectors.addFirst(createEarlySortingTerminationCollectorContext(reader, searchContext.query(), indexSort,
topDocsFactory.numHits(), searchContext.trackTotalHits(), shouldCollect));
}
final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, reader, hasFilterCollector);
// add the top docs collector, the first collector context in the chain
collectors.addFirst(topDocsFactory);
@ -268,9 +262,7 @@ public class QueryPhase implements SearchPhase {
}
try {
if (shouldCollect) {
searcher.search(query, queryCollector);
}
searcher.search(query, queryCollector);
} catch (TimeExceededException e) {
assert timeoutSet : "TimeExceededException thrown even though timeout wasn't set";
queryResult.searchTimedOut(true);
@ -280,7 +272,7 @@ public class QueryPhase implements SearchPhase {
final QuerySearchResult result = searchContext.queryResult();
for (QueryCollectorContext ctx : collectors) {
ctx.postProcess(result, shouldCollect);
ctx.postProcess(result);
}
EsThreadPoolExecutor executor = (EsThreadPoolExecutor)
searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH);
@ -317,13 +309,21 @@ public class QueryPhase implements SearchPhase {
}
/**
* Returns true if the provided <code>searchContext</code> can early terminate based on <code>indexSort</code>
* @param indexSort The index sort specification
* @param context The search context for the request
*/
static boolean canEarlyTerminate(Sort indexSort, SearchContext context) {
final Sort sort = context.sort() == null ? Sort.RELEVANCE : context.sort().sort;
return indexSort != null && EarlyTerminatingSortingCollector.canEarlyTerminate(sort, indexSort);
* Returns whether collection within the provided <code>reader</code> can be early-terminated if it sorts
* with <code>sortAndFormats</code>.
**/
static boolean canEarlyTerminate(IndexReader reader, SortAndFormats sortAndFormats) {
if (sortAndFormats == null || sortAndFormats.sort == null) {
return false;
}
final Sort sort = sortAndFormats.sort;
for (LeafReaderContext ctx : reader.leaves()) {
Sort indexSort = ctx.reader().getMetaData().getSort();
if (indexSort == null || EarlyTerminatingSortingCollector.canEarlyTerminate(sort, indexSort) == false) {
return false;
}
}
return true;
}
private static class TimeExceededException extends RuntimeException {}

144
core/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java
View File

@ -27,6 +27,7 @@ import org.apache.lucene.search.Collector;
import org.apache.lucene.search.ConstantScoreQuery;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.MatchAllDocsQuery;
import org.apache.lucene.search.MultiCollector;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.Sort;
@ -48,9 +49,12 @@ import org.elasticsearch.search.sort.SortAndFormats;
import java.io.IOException;
import java.util.Objects;
import java.util.function.IntSupplier;
import java.util.function.Supplier;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_COUNT;
import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TOP_HITS;
import static org.elasticsearch.search.query.QueryPhase.canEarlyTerminate;
/**
* A {@link QueryCollectorContext} that creates top docs collector
@ -77,27 +81,36 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
return false;
}
static class TotalHitCountCollectorContext extends TopDocsCollectorContext {
private final TotalHitCountCollector collector;
private final int hitCount;
static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext {
private final Collector collector;
private final IntSupplier hitCountSupplier;
/**
* Ctr
* @param reader The index reader
* @param query The query to execute
* @param shouldCollect True if any previous collector context in the chain forces the search to be executed, false otherwise
* @param trackTotalHits True if the total number of hits should be tracked
* @param hasFilterCollector True if the collector chain contains a filter
*/
private TotalHitCountCollectorContext(IndexReader reader, Query query, boolean shouldCollect) throws IOException {
private EmptyTopDocsCollectorContext(IndexReader reader, Query query,
boolean trackTotalHits, boolean hasFilterCollector) throws IOException {
super(REASON_SEARCH_COUNT, 0);
this.collector = new TotalHitCountCollector();
// implicit total hit counts are valid only when there is no filter collector in the chain
// so we check the shortcut only if shouldCollect is true
this.hitCount = shouldCollect ? -1 : shortcutTotalHitCount(reader, query);
}
@Override
boolean shouldCollect() {
return hitCount == -1;
if (trackTotalHits) {
TotalHitCountCollector hitCountCollector = new TotalHitCountCollector();
// implicit total hit counts are valid only when there is no filter collector in the chain
int hitCount = hasFilterCollector ? -1 : shortcutTotalHitCount(reader, query);
if (hitCount == -1) {
this.collector = hitCountCollector;
this.hitCountSupplier = hitCountCollector::getTotalHits;
} else {
this.collector = new EarlyTerminatingCollector(hitCountCollector, 0);
this.hitCountSupplier = () -> hitCount;
}
} else {
this.collector = new EarlyTerminatingCollector(new TotalHitCountCollector(), 0);
// for bwc hit count is set to 0, it will be converted to -1 by the coordinating node
this.hitCountSupplier = () -> 0;
}
}
Collector create(Collector in) {
@ -106,14 +119,8 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
}
@Override
void postProcess(QuerySearchResult result, boolean hasCollected) {
final int totalHitCount;
if (hasCollected) {
totalHitCount = collector.getTotalHits();
} else {
assert hitCount != -1;
totalHitCount = hitCount;
}
void postProcess(QuerySearchResult result) {
final int totalHitCount = hitCountSupplier.getAsInt();
result.topDocs(new TopDocs(totalHitCount, Lucene.EMPTY_SCORE_DOCS, 0), null);
}
}
@ -148,47 +155,83 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
}
@Override
void postProcess(QuerySearchResult result, boolean hasCollected) throws IOException {
assert hasCollected;
void postProcess(QuerySearchResult result) throws IOException {
result.topDocs(topDocsCollector.getTopDocs(), sortFmt);
}
}
abstract static class SimpleTopDocsCollectorContext extends TopDocsCollectorContext {
private final @Nullable SortAndFormats sortAndFormats;
private final TopDocsCollector<?> topDocsCollector;
private final Collector collector;
private final IntSupplier totalHitsSupplier;
private final Supplier<TopDocs> topDocsSupplier;
/**
* Ctr
* @param reader The index reader
* @param query The Lucene query
* @param sortAndFormats The query sort
* @param numHits The number of top hits to retrieve
* @param searchAfter The doc this request should "search after"
* @param trackMaxScore True if max score should be tracked
* @param trackTotalHits True if the total number of hits should be tracked
* @param hasFilterCollector True if the collector chain contains at least one collector that can filters document
*/
private SimpleTopDocsCollectorContext(@Nullable SortAndFormats sortAndFormats,
private SimpleTopDocsCollectorContext(IndexReader reader,
Query query,
@Nullable SortAndFormats sortAndFormats,
@Nullable ScoreDoc searchAfter,
int numHits,
boolean trackMaxScore) throws IOException {
boolean trackMaxScore,
boolean trackTotalHits,
boolean hasFilterCollector) throws IOException {
super(REASON_SEARCH_TOP_HITS, numHits);
this.sortAndFormats = sortAndFormats;
if (sortAndFormats == null) {
this.topDocsCollector = TopScoreDocCollector.create(numHits, searchAfter);
final TopDocsCollector<?> topDocsCollector = TopScoreDocCollector.create(numHits, searchAfter);
this.collector = topDocsCollector;
this.topDocsSupplier = topDocsCollector::topDocs;
this.totalHitsSupplier = topDocsCollector::getTotalHits;
} else {
this.topDocsCollector = TopFieldCollector.create(sortAndFormats.sort, numHits,
(FieldDoc) searchAfter, true, trackMaxScore, trackMaxScore, true);
/**
* We explicitly don't track total hits in the topdocs collector, it can early terminate
* if the sort matches the index sort.
*/
final TopDocsCollector<?> topDocsCollector = TopFieldCollector.create(sortAndFormats.sort, numHits,
(FieldDoc) searchAfter, true, trackMaxScore, trackMaxScore, false);
this.topDocsSupplier = topDocsCollector::topDocs;
if (trackTotalHits) {
// implicit total hit counts are valid only when there is no filter collector in the chain
int count = hasFilterCollector ? -1 : shortcutTotalHitCount(reader, query);
if (count != -1) {
// we can extract the total count from the shard statistics directly
this.totalHitsSupplier = () -> count;
this.collector = topDocsCollector;
} else {
// wrap a collector that counts the total number of hits even
// if the top docs collector terminates early
final TotalHitCountCollector countingCollector = new TotalHitCountCollector();
this.collector = MultiCollector.wrap(topDocsCollector, countingCollector);
this.totalHitsSupplier = countingCollector::getTotalHits;
}
} else {
// total hit count is not needed
this.collector = topDocsCollector;
this.totalHitsSupplier = topDocsCollector::getTotalHits;
}
}
}
@Override
Collector create(Collector in) {
assert in == null;
return topDocsCollector;
return collector;
}
@Override
void postProcess(QuerySearchResult result, boolean hasCollected) throws IOException {
assert hasCollected;
final TopDocs topDocs = topDocsCollector.topDocs();
void postProcess(QuerySearchResult result) throws IOException {
final TopDocs topDocs = topDocsSupplier.get();
topDocs.totalHits = totalHitsSupplier.getAsInt();
result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats);
}
}
@ -197,19 +240,24 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
private final ScrollContext scrollContext;
private final int numberOfShards;
private ScrollingTopDocsCollectorContext(ScrollContext scrollContext,
private ScrollingTopDocsCollectorContext(IndexReader reader,
Query query,
ScrollContext scrollContext,
@Nullable SortAndFormats sortAndFormats,
int numHits,
boolean trackMaxScore,
int numberOfShards) throws IOException {
super(sortAndFormats, scrollContext.lastEmittedDoc, numHits, trackMaxScore);
int numberOfShards,
boolean trackTotalHits,
boolean hasFilterCollector) throws IOException {
super(reader, query, sortAndFormats, scrollContext.lastEmittedDoc, numHits, trackMaxScore,
trackTotalHits, hasFilterCollector);
this.scrollContext = Objects.requireNonNull(scrollContext);
this.numberOfShards = numberOfShards;
}
@Override
void postProcess(QuerySearchResult result, boolean hasCollected) throws IOException {
super.postProcess(result, hasCollected);
void postProcess(QuerySearchResult result) throws IOException {
super.postProcess(result);
final TopDocs topDocs = result.topDocs();
if (scrollContext.totalHits == -1) {
// first round
@ -266,22 +314,24 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
}
/**
* Creates a {@link TopDocsCollectorContext} from the provided <code>searchContext</code>
* Creates a {@link TopDocsCollectorContext} from the provided <code>searchContext</code>.
* @param hasFilterCollector True if the collector chain contains at least one collector that can filters document.
*/
static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searchContext,
IndexReader reader,
boolean shouldCollect) throws IOException {
boolean hasFilterCollector) throws IOException {
final Query query = searchContext.query();
// top collectors don't like a size of 0
final int totalNumDocs = Math.max(1, reader.numDocs());
if (searchContext.size() == 0) {
// no matter what the value of from is
return new TotalHitCountCollectorContext(reader, query, shouldCollect);
return new EmptyTopDocsCollectorContext(reader, query, searchContext.trackTotalHits(), hasFilterCollector);
} else if (searchContext.scrollContext() != null) {
// no matter what the value of from is
int numDocs = Math.min(searchContext.size(), totalNumDocs);
return new ScrollingTopDocsCollectorContext(searchContext.scrollContext(),
searchContext.sort(), numDocs, searchContext.trackScores(), searchContext.numberOfShards());
return new ScrollingTopDocsCollectorContext(reader, query, searchContext.scrollContext(),
searchContext.sort(), numDocs, searchContext.trackScores(), searchContext.numberOfShards(),
searchContext.trackTotalHits(), hasFilterCollector);
} else if (searchContext.collapse() != null) {
boolean trackScores = searchContext.sort() == null ? true : searchContext.trackScores();
int numDocs = Math.min(searchContext.from() + searchContext.size(), totalNumDocs);
@ -296,10 +346,8 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
numDocs = Math.max(numDocs, rescoreContext.getWindowSize());
}
}
return new SimpleTopDocsCollectorContext(searchContext.sort(),
searchContext.searchAfter(),
numDocs,
searchContext.trackScores()) {
return new SimpleTopDocsCollectorContext(reader, query, searchContext.sort(), searchContext.searchAfter(), numDocs,
searchContext.trackScores(), searchContext.trackTotalHits(), hasFilterCollector) {
@Override
boolean shouldRescore() {
return rescore;

153
core/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java
View File

@ -27,7 +27,6 @@ import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.MultiReader;
import org.apache.lucene.index.NoMergePolicy;
import org.apache.lucene.index.RandomIndexWriter;
import org.apache.lucene.index.Term;
@ -64,7 +63,6 @@ import org.elasticsearch.test.TestSearchContext;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.hamcrest.Matchers.anyOf;
import static org.hamcrest.Matchers.equalTo;
@ -98,19 +96,12 @@ public class QueryPhaseTests extends IndexShardTestCase {
context.setSize(0);
context.setTask(new SearchTask(123L, "", "", "", null));
IndexSearcher searcher = new IndexSearcher(reader);
final AtomicBoolean collected = new AtomicBoolean();
IndexSearcher contextSearcher = new IndexSearcher(reader) {
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
collected.set(true);
super.search(leaves, weight, collector);
}
};
final IndexSearcher searcher = shouldCollect ? new IndexSearcher(reader) :
getAssertingEarlyTerminationSearcher(reader, 0);
final boolean rescore = QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
final boolean rescore = QueryPhase.execute(context, searcher, checkCancelled -> {});
assertFalse(rescore);
assertEquals(searcher.count(query), context.queryResult().topDocs().totalHits);
assertEquals(shouldCollect, collected.get());
}
private void countTestCase(boolean withDeletions) throws Exception {
@ -163,51 +154,57 @@ public class QueryPhaseTests extends IndexShardTestCase {
}
public void testPostFilterDisablesCountOptimization() throws Exception {
Directory dir = newDirectory();
final Sort sort = new Sort(new SortField("rank", SortField.Type.INT));
IndexWriterConfig iwc = newIndexWriterConfig()
.setIndexSort(sort);
RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc);
Document doc = new Document();
w.addDocument(doc);
w.close();
IndexReader reader = DirectoryReader.open(dir);
IndexSearcher contextSearcher = getAssertingEarlyTerminationSearcher(reader, 0);
TestSearchContext context = new TestSearchContext(null, indexShard);
context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
context.setSize(0);
context.setTask(new SearchTask(123L, "", "", "", null));
context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
final AtomicBoolean collected = new AtomicBoolean();
IndexSearcher contextSearcher = new IndexSearcher(new MultiReader()) {
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
collected.set(true);
super.search(leaves, weight, collector);
}
};
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertEquals(0, context.queryResult().topDocs().totalHits);
assertFalse(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertEquals(1, context.queryResult().topDocs().totalHits);
contextSearcher = new IndexSearcher(reader);
context.parsedPostFilter(new ParsedQuery(new MatchNoDocsQuery()));
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertEquals(0, context.queryResult().topDocs().totalHits);
assertTrue(collected.get());
reader.close();
dir.close();
}
public void testMinScoreDisablesCountOptimization() throws Exception {
Directory dir = newDirectory();
final Sort sort = new Sort(new SortField("rank", SortField.Type.INT));
IndexWriterConfig iwc = newIndexWriterConfig()
.setIndexSort(sort);
RandomIndexWriter w = new RandomIndexWriter(random(), dir, iwc);
Document doc = new Document();
w.addDocument(doc);
w.close();
IndexReader reader = DirectoryReader.open(dir);
IndexSearcher contextSearcher = getAssertingEarlyTerminationSearcher(reader, 0);
TestSearchContext context = new TestSearchContext(null, indexShard);
context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
context.setSize(0);
context.setTask(new SearchTask(123L, "", "", "", null));
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertEquals(1, context.queryResult().topDocs().totalHits);
final AtomicBoolean collected = new AtomicBoolean();
IndexSearcher contextSearcher = new IndexSearcher(new MultiReader()) {
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
collected.set(true);
super.search(leaves, weight, collector);
}
};
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
contextSearcher = new IndexSearcher(reader);
context.minimumScore(100);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertEquals(0, context.queryResult().topDocs().totalHits);
assertFalse(collected.get());
context.minimumScore(1);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertEquals(0, context.queryResult().topDocs().totalHits);
assertTrue(collected.get());
reader.close();
dir.close();
}
public void testQueryCapturesThreadPoolStats() throws Exception {
@ -226,7 +223,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
IndexReader reader = DirectoryReader.open(dir);
IndexSearcher contextSearcher = new IndexSearcher(reader);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
QuerySearchResult results = context.queryResult();
assertThat(results.serviceTimeEWMA(), greaterThanOrEqualTo(0L));
assertThat(results.nodeQueueSize(), greaterThanOrEqualTo(0));
@ -245,15 +242,8 @@ public class QueryPhaseTests extends IndexShardTestCase {
w.addDocument(new Document());
}
w.close();
final AtomicBoolean collected = new AtomicBoolean();
IndexReader reader = DirectoryReader.open(dir);
IndexSearcher contextSearcher = new IndexSearcher(reader) {
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
collected.set(true);
super.search(leaves, weight, collector);
}
};
IndexSearcher contextSearcher = new IndexSearcher(reader);
TestSearchContext context = new TestSearchContext(null, indexShard);
context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
ScrollContext scrollContext = new ScrollContext();
@ -262,22 +252,22 @@ public class QueryPhaseTests extends IndexShardTestCase {
scrollContext.totalHits = -1;
context.scrollContext(scrollContext);
context.setTask(new SearchTask(123L, "", "", "", null));
context.setSize(10);
int size = randomIntBetween(2, 5);
context.setSize(size);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertThat(context.queryResult().topDocs().totalHits, equalTo((long) numDocs));
assertTrue(collected.get());
assertNull(context.queryResult().terminatedEarly());
assertThat(context.terminateAfter(), equalTo(0));
assertThat(context.queryResult().getTotalHits(), equalTo((long) numDocs));
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
contextSearcher = getAssertingEarlyTerminationSearcher(reader, size);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertThat(context.queryResult().topDocs().totalHits, equalTo((long) numDocs));
assertTrue(collected.get());
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.terminateAfter(), equalTo(10));
assertThat(context.terminateAfter(), equalTo(size));
assertThat(context.queryResult().getTotalHits(), equalTo((long) numDocs));
assertThat(context.queryResult().topDocs().scoreDocs[0].doc, greaterThanOrEqualTo(10));
assertThat(context.queryResult().topDocs().scoreDocs[0].doc, greaterThanOrEqualTo(size));
reader.close();
dir.close();
}
@ -304,26 +294,18 @@ public class QueryPhaseTests extends IndexShardTestCase {
context.parsedQuery(new ParsedQuery(new MatchAllDocsQuery()));
context.terminateAfter(1);
final AtomicBoolean collected = new AtomicBoolean();
final IndexReader reader = DirectoryReader.open(dir);
IndexSearcher contextSearcher = new IndexSearcher(reader) {
protected void search(List<LeafReaderContext> leaves, Weight weight, Collector collector) throws IOException {
collected.set(true);
super.search(leaves, weight, collector);
}
};
IndexSearcher contextSearcher = new IndexSearcher(reader);
{
context.setSize(1);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
context.setSize(0);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(0));
@ -331,8 +313,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
{
context.setSize(1);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
@ -344,40 +325,32 @@ public class QueryPhaseTests extends IndexShardTestCase {
.add(new TermQuery(new Term("foo", "baz")), Occur.SHOULD)
.build();
context.parsedQuery(new ParsedQuery(bq));
collected.set(false);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
context.setSize(0);
context.parsedQuery(new ParsedQuery(bq));
collected.set(false);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(0));
}
{
context.setSize(1);
collected.set(false);
TotalHitCountCollector collector = new TotalHitCountCollector();
context.queryCollectors().put(TotalHitCountCollector.class, collector);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
}
{
context.setSize(0);
collected.set(false);
TotalHitCountCollector collector = new TotalHitCountCollector();
context.queryCollectors().put(TotalHitCountCollector.class, collector);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null);
assertTrue(collected.get());
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertTrue(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(0));
@ -416,7 +389,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
final IndexReader reader = DirectoryReader.open(dir);
IndexSearcher contextSearcher = new IndexSearcher(reader);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertThat(context.queryResult().topDocs().totalHits, equalTo((long) numDocs));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
assertThat(context.queryResult().topDocs().scoreDocs[0], instanceOf(FieldDoc.class));
@ -425,7 +398,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
{
context.parsedPostFilter(new ParsedQuery(new MinDocQuery(1)));
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertNull(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo(numDocs - 1L));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
@ -435,7 +408,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
final TotalHitCountCollector totalHitCountCollector = new TotalHitCountCollector();
context.queryCollectors().put(TotalHitCountCollector.class, totalHitCountCollector);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertNull(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo((long) numDocs));
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
@ -448,13 +421,13 @@ public class QueryPhaseTests extends IndexShardTestCase {
{
contextSearcher = getAssertingEarlyTerminationSearcher(reader, 1);
context.trackTotalHits(false);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertNull(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
assertThat(context.queryResult().topDocs().scoreDocs[0], instanceOf(FieldDoc.class));
assertThat(fieldDoc.fields[0], anyOf(equalTo(1), equalTo(2)));
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertNull(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().scoreDocs.length, equalTo(1));
assertThat(context.queryResult().topDocs().scoreDocs[0], instanceOf(FieldDoc.class));
@ -502,7 +475,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
context.setSize(10);
context.sort(searchSortAndFormat);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, searchSortAndFormat.sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertThat(context.queryResult().topDocs().totalHits, equalTo((long) numDocs));
assertNull(context.queryResult().terminatedEarly());
assertThat(context.terminateAfter(), equalTo(0));
@ -511,7 +484,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
FieldDoc lastDoc = (FieldDoc) context.queryResult().topDocs().scoreDocs[sizeMinus1];
contextSearcher = getAssertingEarlyTerminationSearcher(reader, 10);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, searchSortAndFormat.sort);
QueryPhase.execute(context, contextSearcher, checkCancelled -> {});
assertNull(context.queryResult().terminatedEarly());
assertThat(context.queryResult().topDocs().totalHits, equalTo((long) numDocs));
assertThat(context.terminateAfter(), equalTo(0));