Shortcut query phase using the results of other shards

modules/reindex/src/main/java/org/elasticsearch/index/reindex/BulkByScrollParallelizationHelper.java

@@ -203,7 +203,7 @@ static SearchRequest[] sliceIntoSubRequests(SearchRequest request, String field,
                 if (request.source().slice() != null) {
                     throw new IllegalStateException("Can't slice a request that already has a slice configuration");
                 }
-                slicedSource = request.source().copyWithNewSlice(sliceBuilder);
+                slicedSource = request.source().shallowCopy().slice(sliceBuilder);
             }
             SearchRequest searchRequest = new SearchRequest(request);
             searchRequest.source(slicedSource);

rest-api-spec/src/main/resources/rest-api-spec/test/search/140_pre_filter_search_shards.yml

@@ -154,10 +154,9 @@ setup:
   - match: { hits.total: 2 }
   - length: { aggregations.idx_terms.buckets: 2 }
 
-    # check that empty responses are correctly handled when rewriting to match_no_docs
+  # check that empty responses are correctly handled when rewriting to match_no_docs
   - do:
       search:
-        rest_total_hits_as_int: true
         # ensure that one shard can return empty response
         max_concurrent_shard_requests: 1
         body: { "size" : 0, "query" : { "range" : { "created_at" : { "gte" : "2016-02-01", "lt": "2018-02-01"}}}, "aggs" : { "idx_terms" :  { "terms" : { "field" : "_index" } } } }
@@ -166,12 +165,11 @@ setup:
   - match: { _shards.successful: 3 }
   - match: { _shards.skipped : 0 }
   - match: { _shards.failed: 0 }
-  - match: { hits.total: 2 }
+  - match: { hits.total.value: 2 }
   - length: { aggregations.idx_terms.buckets: 2 }
 
   - do:
       search:
-        rest_total_hits_as_int: true
         # ensure that one shard can return empty response
         max_concurrent_shard_requests: 2
         body: { "size" : 0, "query" : { "range" : { "created_at" : { "gte" : "2019-02-01"}}}, "aggs" : { "idx_terms" :  { "terms" : { "field" : "_index" } } } }
@@ -180,5 +178,47 @@ setup:
   - match: { _shards.successful: 3 }
   - match: { _shards.skipped : 0 }
   - match: { _shards.failed: 0 }
-  - match: { hits.total: 0 }
+  - match: { hits.total.value: 0 }
   - length: { aggregations.idx_terms.buckets: 0 }
+
+  # check field sort is correct when skipping query phase
+  - do:
+      search:
+        # ensure that one shard can return empty response
+        max_concurrent_shard_requests: 1
+        pre_filter_shard_size: 1
+        body:
+          "size": 1
+          "track_total_hits": 1
+          "sort": [{ "created_at": "desc"}]
+
+  - match: { _shards.total: 3 }
+  - match: { _shards.successful: 3 }
+  - match: { _shards.skipped: 0 }
+  - match: { _shards.failed: 0 }
+  - match: { hits.total.value: 1 }
+  - match: { hits.total.relation: "gte" }
+  - length: { hits.hits: 1 }
+  - match: { hits.hits.0._id: "3" }
+
+  # same with aggs
+  - do:
+      search:
+        # ensure that one shard can return empty response
+        max_concurrent_shard_requests: 1
+        pre_filter_shard_size: 1
+        body:
+          "size": 1
+          "track_total_hits": 1
+          "sort": [{ "created_at": "desc"}]
+          "aggs" : { "idx_terms" :  { "terms" : { "field" : "_index" } } }
+
+  - match: { _shards.total: 3 }
+  - match: { _shards.successful: 3 }
+  - match: { _shards.skipped: 0 }
+  - match: { _shards.failed: 0 }
+  - match: { hits.total.value: 1 }
+  - match: { hits.total.relation: "gte" }
+  - length: { hits.hits: 1 }
+  - match: {hits.hits.0._id: "3" }
+  - length: { aggregations.idx_terms.buckets: 3 }

server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java

@@ -76,7 +76,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
      **/
     private final BiFunction<String, String, Transport.Connection> nodeIdToConnection;
     private final SearchTask task;
-    private final SearchPhaseResults<Result> results;
+    protected final SearchPhaseResults<Result> results;
     private final long clusterStateVersion;
     private final Map<String, AliasFilter> aliasFilter;
     private final Map<String, Float> concreteIndexBoosts;

server/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java

@@ -611,9 +611,9 @@ static final class QueryPhaseResultConsumer extends ArraySearchPhaseResults<Sear
          * @param bufferSize the size of the reduce buffer. if the buffer size is smaller than the number of expected results
          *                   the buffer is used to incrementally reduce aggregation results before all shards responded.
          */
-        private QueryPhaseResultConsumer(SearchProgressListener progressListener, SearchPhaseController controller,
-                                         int expectedResultSize, int bufferSize, boolean hasTopDocs, boolean hasAggs,
-                                         int trackTotalHitsUpTo, int topNSize, boolean performFinalReduce) {
+        QueryPhaseResultConsumer(SearchProgressListener progressListener, SearchPhaseController controller,
+                                 int expectedResultSize, int bufferSize, boolean hasTopDocs, boolean hasAggs,
+                                 int trackTotalHitsUpTo, int topNSize, boolean performFinalReduce) {
             super(expectedResultSize);
             if (expectedResultSize != 1 && bufferSize < 2) {
                 throw new IllegalArgumentException("buffer size must be >= 2 if there is more than one expected result");
@@ -683,6 +683,15 @@ private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
             processedShards[querySearchResult.getShardIndex()] = querySearchResult.getSearchShardTarget();
         }
 
+        public TopDocs getBufferTopDocs() {
+            if (hasTopDocs) {
+                synchronized (this) {
+                    return topDocsBuffer[0];
+                }
+            }
+            return null;
+        }
+
         private synchronized List<InternalAggregations> getRemainingAggs() {
             return hasAggs ? Arrays.asList(aggsBuffer).subList(0, index) : null;
         }

server/src/main/java/org/elasticsearch/action/search/SearchQueryThenFetchAsyncAction.java

@@ -20,22 +20,33 @@
 package org.elasticsearch.action.search;
 
 import org.apache.logging.log4j.Logger;
+import org.apache.lucene.search.FieldDoc;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TopFieldDocs;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.routing.GroupShardsIterator;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.internal.AliasFilter;
+import org.elasticsearch.search.internal.ShardSearchRequest;
+import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.transport.Transport;
 
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 
-final class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPhaseResult> {
+import static org.elasticsearch.action.search.SearchPhaseController.getTopDocsSize;
+import static org.elasticsearch.search.internal.SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO;
+
+class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPhaseResult> {
 
     private final SearchPhaseController searchPhaseController;
+    private final Supplier<TopDocs> topDocsSupplier;
+    private final int topDocsSize;
     private final SearchProgressListener progressListener;
 
     SearchQueryThenFetchAsyncAction(final Logger logger, final SearchTransportService searchTransportService,
@@ -49,18 +60,20 @@ final class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<Se
                 executor, request, listener, shardsIts, timeProvider, clusterStateVersion, task,
                 searchPhaseController.newSearchPhaseResults(task.getProgressListener(), request, shardsIts.size()),
                 request.getMaxConcurrentShardRequests(), clusters);
+        this.topDocsSize = getTopDocsSize(request);
+        this.topDocsSupplier = getBufferTopDocsSupplier(request, results);
         this.searchPhaseController = searchPhaseController;
         this.progressListener = task.getProgressListener();
-        final SearchProgressListener progressListener = task.getProgressListener();
         final SearchSourceBuilder sourceBuilder = request.source();
         progressListener.notifyListShards(progressListener.searchShards(this.shardsIts),
             sourceBuilder == null || sourceBuilder.size() != 0);
     }
 
     protected void executePhaseOnShard(final SearchShardIterator shardIt, final ShardRouting shard,
                                        final SearchActionListener<SearchPhaseResult> listener) {
+        ShardSearchRequest request = rewriteShardRequest(buildShardSearchRequest(shardIt));
         getSearchTransport().sendExecuteQuery(getConnection(shardIt.getClusterAlias(), shard.currentNodeId()),
-            buildShardSearchRequest(shardIt), getTask(), listener);
+            request, getTask(), listener);
     }
 
     @Override
@@ -72,4 +85,38 @@ protected void onShardGroupFailure(int shardIndex, Exception exc) {
     protected SearchPhase getNextPhase(final SearchPhaseResults<SearchPhaseResult> results, final SearchPhaseContext context) {
         return new FetchSearchPhase(results, searchPhaseController, context);
     }
+
+    ShardSearchRequest rewriteShardRequest(ShardSearchRequest request) {
+        TopDocs topDocs = topDocsSupplier.get();
+        if (topDocs != null && topDocs instanceof TopFieldDocs) {
+            SearchSourceBuilder source = request.source();
+            int trackTotalHits = source.trackTotalHitsUpTo() == null ? DEFAULT_TRACK_TOTAL_HITS_UP_TO :
+                source.trackTotalHitsUpTo();
+            if (topDocs.totalHits.value >= trackTotalHits) {
+                SearchSourceBuilder newSource = source.shallowCopy();
+                newSource.trackTotalHits(false);
+                if (topDocs.scoreDocs.length >= topDocsSize) {
+                    FieldDoc bottomDoc = (FieldDoc) topDocs.scoreDocs[topDocs.scoreDocs.length-1];
+                    request.setRawBottomSortValues(bottomDoc.fields);
+                }
+                request.source(newSource);
+            }
+        }
+        return request;
+    }
+
+    private Supplier<TopDocs> getBufferTopDocsSupplier(SearchRequest request,
+                                                       SearchPhaseResults<SearchPhaseResult> searchPhaseResults) {
+        if (searchPhaseResults instanceof SearchPhaseController.QueryPhaseResultConsumer == false) {
+            return () -> null;
+        }
+        int size = getTopDocsSize(request);
+        FieldSortBuilder fieldSort = FieldSortBuilder.getPrimaryFieldSortOrNull(request.source());
+        if (size == 0
+                || fieldSort == null
+                || fieldSort.canRewriteToMatchNone() == false) {
+            return () -> null;
+        }
+        return ((SearchPhaseController.QueryPhaseResultConsumer) searchPhaseResults)::getBufferTopDocs;
+    }
 }

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -959,13 +959,13 @@ public SearchSourceBuilder rewrite(QueryRewriteContext context) throws IOExcepti
     /**
      * Create a shallow copy of this builder with a new slice configuration.
      */
-    public SearchSourceBuilder copyWithNewSlice(SliceBuilder slice) {
-        return shallowCopy(queryBuilder, postQueryBuilder, aggregations, slice, sorts, rescoreBuilders, highlightBuilder);
+    public SearchSourceBuilder shallowCopy() {
+        return shallowCopy(queryBuilder, postQueryBuilder, aggregations, sliceBuilder, sorts, rescoreBuilders, highlightBuilder);
     }
 
     /**
      * Create a shallow copy of this source replaced {@link #queryBuilder}, {@link #postQueryBuilder}, and {@link #sliceBuilder}. Used by
-     * {@link #rewrite(QueryRewriteContext)} and {@link #copyWithNewSlice(SliceBuilder)}.
+     * {@link #rewrite(QueryRewriteContext)}}.
      */
     @SuppressWarnings("rawtypes")
     private SearchSourceBuilder shallowCopy(QueryBuilder queryBuilder, QueryBuilder postQueryBuilder,

server/src/main/java/org/elasticsearch/search/internal/ShardSearchRequest.java

@@ -37,17 +37,21 @@
 import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.indices.AliasFilterParsingException;
 import org.elasticsearch.indices.InvalidAliasNameException;
 import org.elasticsearch.search.Scroll;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.query.QuerySearchResult;
+import org.elasticsearch.search.sort.FieldSortBuilder;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskId;
 import org.elasticsearch.transport.TransportRequest;
@@ -57,6 +61,8 @@
 import java.util.Map;
 import java.util.function.Function;
 
+import static org.elasticsearch.search.internal.SearchContext.TRACK_TOTAL_HITS_DISABLED;
+
 /**
  * Shard level request that represents a search.
  * It provides all the methods that the {@link SearchContext} needs.
@@ -77,6 +83,7 @@ public class ShardSearchRequest extends TransportRequest implements IndicesReque
     private final OriginalIndices originalIndices;
 
     private boolean canReturnNullResponseIfMatchNoDocs;
+    private Object[] rawBottomSortValues;
 
     //these are the only mutable fields, as they are subject to rewriting
     private AliasFilter aliasFilter;
@@ -172,8 +179,10 @@ public ShardSearchRequest(StreamInput in) throws IOException {
         preference = in.readOptionalString();
         if (in.getVersion().onOrAfter(Version.V_8_0_0)) {
             canReturnNullResponseIfMatchNoDocs = in.readBoolean();
+            rawBottomSortValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new);
         } else {
             canReturnNullResponseIfMatchNoDocs = false;
+            rawBottomSortValues = null;
         }
         originalIndices = OriginalIndices.readOriginalIndices(in);
     }
@@ -211,6 +220,7 @@ protected final void innerWriteTo(StreamOutput out, boolean asKey) throws IOExce
         }
         if (out.getVersion().onOrAfter(Version.V_8_0_0)) {
             out.writeBoolean(canReturnNullResponseIfMatchNoDocs);
+            out.writeOptionalArray(Lucene::writeSortValue, rawBottomSortValues);
         }
     }
 
@@ -286,6 +296,20 @@ public String preference() {
         return preference;
     }
 
+    /**
+     * Sets the bottom sort values that can be used by the searcher to filter documents
+     * that are after it. This value is computed by coordinating nodes that throttles the
+     * query phase. After a partial merge of successful shards the sort values of the
+     * bottom top document are passed as an hint on subsequent shard requests.
+     */
+    public void setRawBottomSortValues(Object[] values) {
+        this.rawBottomSortValues = values;
+    }
+
+    public Object[] getRawBottomSortValues() {
+        return  rawBottomSortValues;
+    }
+
     /**
      * Returns true if the caller can handle null response {@link QuerySearchResult#nullInstance()}.
      * Defaults to false since the coordinator node needs at least one shard response to build the global
@@ -343,6 +367,25 @@ static class RequestRewritable implements Rewriteable<Rewriteable> {
         public Rewriteable rewrite(QueryRewriteContext ctx) throws IOException {
             SearchSourceBuilder newSource = request.source() == null ? null : Rewriteable.rewrite(request.source(), ctx);
             AliasFilter newAliasFilter = Rewriteable.rewrite(request.getAliasFilter(), ctx);
+
+            QueryShardContext shardContext = ctx.convertToShardContext();
+            FieldSortBuilder primarySort = FieldSortBuilder.getPrimaryFieldSortOrNull(newSource);
+            // checks if the bottom sort values are guaranteed to be more competitive than all the documents
+            // contained in the shard
+            if (shardContext != null
+                    && primarySort != null
+                    && primarySort.isBottomSortWithinShard(shardContext, request.getRawBottomSortValues()) == false) {
+                newSource = newSource.shallowCopy();
+                if (newSource.trackTotalHitsUpTo() == TRACK_TOTAL_HITS_DISABLED
+                        && newSource.aggregations() == null) {
+                    newSource.query(new MatchNoneQueryBuilder());
+                } else {
+                    newSource.size(0);
+                }
+                request.source(newSource);
+                request.setRawBottomSortValues(null);
+            }
+
             if (newSource == request.source() && newAliasFilter == request.getAliasFilter()) {
                 return this;
             } else {

server/src/main/java/org/elasticsearch/search/sort/FieldSortBuilder.java

@@ -60,6 +60,7 @@
 import java.util.Objects;
 import java.util.function.Function;
 
+import static org.elasticsearch.index.mapper.DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER;
 import static org.elasticsearch.index.mapper.DateFieldMapper.Resolution.MILLISECONDS;
 import static org.elasticsearch.index.mapper.DateFieldMapper.Resolution.NANOSECONDS;
 import static org.elasticsearch.index.search.NestedHelper.parentObject;
@@ -375,6 +376,42 @@ public SortFieldAndFormat build(QueryShardContext context) throws IOException {
         return new SortFieldAndFormat(field, fieldType.docValueFormat(null, null));
     }
 
+    public boolean canRewriteToMatchNone() {
+        return nestedSort == null && (missing == null || "_last".equals(missing));
+    }
+
+    /**
+     * Returns whether all values of the given {@link QueryShardContext#getIndexReader()} are within the
+     * primary sort value provided in the <code>rawBottomSortValues</code>.
+     */
+    public boolean isBottomSortWithinShard(QueryShardContext context,
+                                           Object[] rawBottomSortValues) {
+        if (rawBottomSortValues == null || rawBottomSortValues.length == 0) {
+            return true;
+        }
+
+        if (canRewriteToMatchNone() == false) {
+            return true;
+        }
+        MappedFieldType fieldType = context.fieldMapper(fieldName);
+        if (fieldType == null) {
+            // unmapped
+            return false;
+        }
+        if (fieldType.indexOptions() == IndexOptions.NONE) {
+            return true;
+        }
+        Object minValue = order() == SortOrder.DESC ? rawBottomSortValues[0] : null;
+        Object maxValue = order() == SortOrder.DESC ? null : rawBottomSortValues[0];
+        try {
+            MappedFieldType.Relation relation = fieldType.isFieldWithinQuery(context.getIndexReader(), minValue, maxValue,
+                true, true, null, DEFAULT_DATE_TIME_FORMATTER.toDateMathParser(), context);
+            return relation != MappedFieldType.Relation.DISJOINT;
+        } catch (Exception exc) {
+            return true;
+        }
+    }
+
     /**
      * Return true if the primary sort in the provided <code>source</code>
      * is an instance of {@link FieldSortBuilder}.