diff --git a/core/src/main/java/org/apache/lucene/queries/SearchAfterSortedDocQuery.java b/core/src/main/java/org/apache/lucene/queries/SearchAfterSortedDocQuery.java index 75fdeee2719b2..5da0e618752e2 100644 --- a/core/src/main/java/org/apache/lucene/queries/SearchAfterSortedDocQuery.java +++ b/core/src/main/java/org/apache/lucene/queries/SearchAfterSortedDocQuery.java @@ -79,7 +79,7 @@ public Scorer scorer(LeafReaderContext context) throws IOException { 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 @@ static TopComparator getTopComparator(FieldComparator[] fieldComparators, } } - if (topDoc <= doc) { + if (doc <= topDoc) { return false; } return true; diff --git a/core/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java b/core/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java index acb679b2180f9..2ed806a32ae14 100644 --- a/core/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java +++ b/core/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java @@ -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.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 @@ protected InternalProfileCollector createWithProfiler(InternalProfileCollector i return new InternalProfileCollector(collector, profilerName, in != null ? Collections.singletonList(in) : Collections.emptyList()); } - /** - * A value of false 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 result 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 collectors @@ -175,11 +158,6 @@ static QueryCollectorContext createCancellableCollectorContext(BooleanSupplier c Collector create(Collector in) throws IOException { return new CancellableCollector(cancelled, in); } - - @Override - boolean shouldCollect() { - return false; - } }; } @@ -198,52 +176,11 @@ Collector create(Collector in) throws IOException { } @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 numHits per segment. - * The total hit count matching the query is also computed if trackTotalHits 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()); - } - } - }; - } } diff --git a/core/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/core/src/main/java/org/elasticsearch/search/query/QueryPhase.java index e966d6ef2b827..f028f6014adb9 100644 --- a/core/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/core/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -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.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.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 void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep // 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 void execute(SearchContext searchContext) throws QueryPhaseExecutionExcep * wire everything (mapperService, etc.) * @return whether the rescoring phase should be executed */ - static boolean execute(SearchContext searchContext, final IndexSearcher searcher, - Consumer checkCancellationSetter, @Nullable Sort indexSort) throws QueryPhaseExecutionException { + static boolean execute(SearchContext searchContext, + final IndexSearcher searcher, + Consumer 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 @@ static boolean execute(SearchContext searchContext, final IndexSearcher searcher // ... 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 @@ static boolean execute(SearchContext searchContext, final IndexSearcher searcher } final LinkedList 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 @@ static boolean execute(SearchContext searchContext, final IndexSearcher searcher 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 @@ static boolean execute(SearchContext searchContext, final IndexSearcher searcher // 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 @@ static boolean execute(SearchContext searchContext, final IndexSearcher searcher } 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 @@ static boolean execute(SearchContext searchContext, final IndexSearcher searcher 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 @@ static boolean returnsDocsInOrder(Query query, SortAndFormats sf) { } /** - * Returns true if the provided searchContext can early terminate based on indexSort - * @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 reader can be early-terminated if it sorts + * with sortAndFormats. + **/ + 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 {} diff --git a/core/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java b/core/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java index a5eff58584701..18e351e34a79b 100644 --- a/core/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java +++ b/core/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java @@ -27,6 +27,7 @@ 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 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 @@ boolean shouldRescore() { 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 @@ Collector create(Collector in) { } @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 @@ Collector create(Collector in) throws IOException { } @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 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 @@ static class ScrollingTopDocsCollectorContext extends SimpleTopDocsCollectorCont 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 @@ static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOExcep } /** - * Creates a {@link TopDocsCollectorContext} from the provided searchContext + * Creates a {@link TopDocsCollectorContext} from the provided searchContext. + * @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 @@ static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searc 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; diff --git a/core/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java b/core/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java index 425f5a56f8ea3..d651c92cd611e 100644 --- a/core/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java +++ b/core/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java @@ -27,7 +27,6 @@ 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 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 @@ private void countTestCase(Query query, IndexReader reader, boolean shouldCollec 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 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 void testCountWithDeletions() throws Exception { } 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 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 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()); - - context.minimumScore(1); - QueryPhase.execute(context, contextSearcher, checkCancelled -> {}, null); + contextSearcher = new IndexSearcher(reader); + context.minimumScore(100); + QueryPhase.execute(context, contextSearcher, checkCancelled -> {}); assertEquals(0, context.queryResult().topDocs().totalHits); - assertTrue(collected.get()); + reader.close(); + dir.close(); } public void testQueryCapturesThreadPoolStats() throws Exception { @@ -226,7 +223,7 @@ public void testQueryCapturesThreadPoolStats() throws Exception { 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 void testInOrderScrollOptimization() throws Exception { 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 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 @@ protected void search(List leaves, Weight weight, Collector c 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 void testTerminateAfterEarlyTermination() throws Exception { 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 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 @@ protected void search(List leaves, Weight weight, Collector c { 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 @@ protected void search(List leaves, Weight weight, Collector c .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 void testIndexSortingEarlyTermination() throws Exception { 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 void testIndexSortingEarlyTermination() throws Exception { { 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 void testIndexSortingEarlyTermination() throws Exception { 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 void testIndexSortingEarlyTermination() throws Exception { { 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 void testIndexSortScrollOptimization() throws Exception { 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 void testIndexSortScrollOptimization() throws Exception { 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));