Remove QueryCollectorContext abstraction (#95383)

The QueryCollectorContext abstraction was introduced by #24864 based on the requirement that the top docs collector creation needed to be delayed until after all the other collectors had been created. At the same time, collectors get wrapped depending on the search features enabled by the request, but the top score / total hit count collector is the root collector where the wrapping starts, which is why its corresponding context gets added at position 0 in the list of collector contexts.

Requirements have changed since #27666 , which means that we can go back to a simpler way of creating collectors and wrapping them. We no longer need a QueryCollectorContext abstraction, and we can instead create collectors straight-away, and wrap them as needed. This is much easier to follow compared to the very generic create(Collector) method that the context exposes.

TopDocsCollectorContext adds some value in that it incorporates all the logic around creating the top docs collector, yet it can be further simplified as well by making the postProcess method more specific.

server/src/main/java/org/elasticsearch/search/profile/query/CollectorResult.java

@@ -67,7 +67,7 @@ public class CollectorResult implements ToXContentObject, Writeable {
     /**
      * A list of children collectors "embedded" inside this collector
      */
-    private List<CollectorResult> children;
+    private final List<CollectorResult> children;
 
     public CollectorResult(String collectorName, String reason, long time, List<CollectorResult> children) {
         this.collectorName = collectorName;

server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java

@@ -47,15 +47,15 @@ public class InternalProfileCollector implements Collector {
      */
     private final InternalProfileCollector[] children;
 
-    public InternalProfileCollector(Collector collector, String reason, InternalProfileCollector... children) {
+    public InternalProfileCollector(Collector collector, String reason, Collector... children) {
         this.collector = new ProfileCollector(collector);
         this.reason = reason;
         this.collectorName = deriveCollectorName(collector);
         Objects.requireNonNull(children, "children collectors cannot be null");
-        for (InternalProfileCollector child : children) {
-            Objects.requireNonNull(child, "child collector cannot be null");
+        this.children = new InternalProfileCollector[children.length];
+        for (int i = 0; i < children.length; i++) {
+            this.children[i] = (InternalProfileCollector) Objects.requireNonNull(children[i], "child collector cannot be null");
         }
-        this.children = children;
     }
 
     /**

server/src/main/java/org/elasticsearch/search/query/QueryCollectorContext.java

@@ -1,169 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0 and the Server Side Public License, v 1; you may not use this file except
- * in compliance with, at your election, the Elastic License 2.0 or the Server
- * Side Public License, v 1.
- */
-
-package org.elasticsearch.search.query;
-
-import org.apache.lucene.search.Collector;
-import org.apache.lucene.search.IndexSearcher;
-import org.apache.lucene.search.MultiCollector;
-import org.apache.lucene.search.Query;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.SimpleCollector;
-import org.apache.lucene.search.Weight;
-import org.elasticsearch.common.lucene.MinimumScoreCollector;
-import org.elasticsearch.common.lucene.search.FilteredCollector;
-import org.elasticsearch.search.profile.query.InternalProfileCollector;
-
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
-
-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;
-
-abstract class QueryCollectorContext {
-    private static final Collector EMPTY_COLLECTOR = new SimpleCollector() {
-        @Override
-        public void collect(int doc) {}
-
-        @Override
-        public ScoreMode scoreMode() {
-            return ScoreMode.COMPLETE_NO_SCORES;
-        }
-    };
-
-    private final String profilerName;
-
-    QueryCollectorContext(String profilerName) {
-        this.profilerName = profilerName;
-    }
-
-    /**
-     * Creates a collector that delegates documents to the provided <code>in</code> collector.
-     * @param in The delegate collector
-     */
-    abstract Collector create(Collector in) throws IOException;
-
-    /**
-     * Wraps this collector with a profiler
-     */
-    protected InternalProfileCollector createWithProfiler(InternalProfileCollector in) throws IOException {
-        final Collector collector = create(in);
-        if (in == null) {
-            return new InternalProfileCollector(collector, profilerName);
-        }
-        return new InternalProfileCollector(collector, profilerName, in);
-    }
-
-    /**
-     * Post-process <code>result</code> after search execution.
-     *
-     * @param result The query search result to populate
-     */
-    void postProcess(QuerySearchResult result) throws IOException {}
-
-    /**
-     * Creates the collector tree from the provided <code>collectors</code>
-     * @param collectors Ordered list of collector context
-     */
-    static Collector createQueryCollector(List<QueryCollectorContext> collectors) throws IOException {
-        Collector collector = null;
-        for (QueryCollectorContext ctx : collectors) {
-            collector = ctx.create(collector);
-        }
-        return collector;
-    }
-
-    /**
-     * Creates the collector tree from the provided <code>collectors</code> and wraps each collector with a profiler
-     * @param collectors Ordered list of collector context
-     */
-    static InternalProfileCollector createQueryCollectorWithProfiler(List<QueryCollectorContext> collectors) throws IOException {
-        InternalProfileCollector collector = null;
-        for (QueryCollectorContext ctx : collectors) {
-            collector = ctx.createWithProfiler(collector);
-        }
-        return collector;
-    }
-
-    /**
-     * Filters documents with a query score greater than <code>minScore</code>
-     * @param minScore The minimum score filter
-     */
-    static QueryCollectorContext createMinScoreCollectorContext(float minScore) {
-        return new QueryCollectorContext(REASON_SEARCH_MIN_SCORE) {
-            @Override
-            Collector create(Collector in) {
-                return new MinimumScoreCollector(in, minScore);
-            }
-        };
-    }
-
-    /**
-     * Filters documents based on the provided <code>query</code>
-     */
-    static QueryCollectorContext createFilteredCollectorContext(IndexSearcher searcher, Query query) {
-        return new QueryCollectorContext(REASON_SEARCH_POST_FILTER) {
-            @Override
-            Collector create(Collector in) throws IOException {
-                final Weight filterWeight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f);
-                return new FilteredCollector(in, filterWeight);
-            }
-        };
-    }
-
-    /**
-     * Creates a multi collector from the provided sub-collector
-     */
-    static QueryCollectorContext createAggsCollectorContext(Collector subCollector) {
-        return new QueryCollectorContext(REASON_SEARCH_MULTI) {
-            @Override
-            Collector create(Collector in) {
-                List<Collector> subCollectors = new ArrayList<>();
-                subCollectors.add(in);
-                subCollectors.add(subCollector);
-                return MultiCollector.wrap(subCollectors);
-            }
-
-            @Override
-            protected InternalProfileCollector createWithProfiler(InternalProfileCollector in) {
-                if (subCollector instanceof InternalProfileCollector == false) {
-                    throw new IllegalArgumentException("non-profiling collector");
-                }
-                final Collector collector = MultiCollector.wrap(in, subCollector);
-                return new InternalProfileCollector(collector, REASON_SEARCH_MULTI, in, (InternalProfileCollector) subCollector);
-            }
-        };
-    }
-
-    /**
-     * Creates collector limiting the collection to the first <code>numHits</code> documents
-     */
-    static QueryCollectorContext createEarlyTerminationCollectorContext(int numHits) {
-        return new QueryCollectorContext(REASON_SEARCH_TERMINATE_AFTER_COUNT) {
-            private Collector collector;
-
-            /**
-             * Creates a {@link MultiCollector} to ensure that the {@link EarlyTerminatingCollector}
-             * can terminate the collection independently of the provided <code>in</code> {@link Collector}.
-             */
-            @Override
-            Collector create(Collector in) {
-                assert collector == null;
-
-                List<Collector> subCollectors = new ArrayList<>();
-                subCollectors.add(new EarlyTerminatingCollector(EMPTY_COLLECTOR, numHits, true));
-                subCollectors.add(in);
-                this.collector = MultiCollector.wrap(subCollectors);
-                return collector;
-            }
-        };
-    }
-}

server/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -16,12 +16,18 @@ import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.FieldDoc;
+import org.apache.lucene.search.MultiCollector;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.SimpleCollector;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.TopDocs;
 import org.apache.lucene.search.TotalHits;
+import org.apache.lucene.search.Weight;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.common.lucene.MinimumScoreCollector;
+import org.elasticsearch.common.lucene.search.FilteredCollector;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.common.util.concurrent.EWMATrackingEsThreadPoolExecutor;
 import org.elasticsearch.common.util.concurrent.EsThreadPoolExecutor;
@@ -33,6 +39,7 @@ import org.elasticsearch.search.aggregations.AggregationPhase;
 import org.elasticsearch.search.internal.ContextIndexSearcher;
 import org.elasticsearch.search.internal.ScrollContext;
 import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.search.profile.Profilers;
 import org.elasticsearch.search.profile.query.InternalProfileCollector;
 import org.elasticsearch.search.rescore.RescorePhase;
 import org.elasticsearch.search.sort.SortAndFormats;
@@ -40,13 +47,12 @@ import org.elasticsearch.search.suggest.SuggestPhase;
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.io.IOException;
-import java.util.LinkedList;
 import java.util.concurrent.ExecutorService;
 
-import static org.elasticsearch.search.query.QueryCollectorContext.createAggsCollectorContext;
-import static org.elasticsearch.search.query.QueryCollectorContext.createEarlyTerminationCollectorContext;
-import static org.elasticsearch.search.query.QueryCollectorContext.createFilteredCollectorContext;
-import static org.elasticsearch.search.query.QueryCollectorContext.createMinScoreCollectorContext;
+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.createTopDocsCollectorContext;
 
 /**
@@ -125,30 +131,65 @@ public class QueryPhase {
                 }
             }
 
-            final LinkedList<QueryCollectorContext> collectors = new LinkedList<>();
-            // whether the chain contains a collector that filters documents
-            boolean hasFilterCollector = false;
+            // create the top docs collector last when the other collectors are known
+            final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(
+                searchContext,
+                searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null
+            );
+
+            Collector collector = wrapWithProfilerCollectorIfNeeded(
+                searchContext.getProfilers(),
+                topDocsFactory.collector(),
+                topDocsFactory.profilerName
+            );
+
             if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER) {
                 // add terminate_after before the filter collectors
                 // it will only be applied on documents accepted by these filter collectors
-                collectors.add(createEarlyTerminationCollectorContext(searchContext.terminateAfter()));
+                EarlyTerminatingCollector earlyTerminatingCollector = new EarlyTerminatingCollector(
+                    EMPTY_COLLECTOR,
+                    searchContext.terminateAfter(),
+                    true
+                );
+                collector = wrapWithProfilerCollectorIfNeeded(
+                    searchContext.getProfilers(),
+                    MultiCollector.wrap(earlyTerminatingCollector, collector),
+                    REASON_SEARCH_TERMINATE_AFTER_COUNT,
+                    collector
+                );
             }
             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;
+                final Weight filterWeight = searcher.createWeight(
+                    searcher.rewrite(searchContext.parsedPostFilter().query()),
+                    ScoreMode.COMPLETE_NO_SCORES,
+                    1f
+                );
+                collector = wrapWithProfilerCollectorIfNeeded(
+                    searchContext.getProfilers(),
+                    new FilteredCollector(collector, filterWeight),
+                    REASON_SEARCH_POST_FILTER,
+                    collector
+                );
             }
             if (searchContext.getAggsCollector() != null) {
-                // plug in additional collectors, like aggregations
-                collectors.add(createAggsCollectorContext(searchContext.getAggsCollector()));
+                collector = wrapWithProfilerCollectorIfNeeded(
+                    searchContext.getProfilers(),
+                    MultiCollector.wrap(collector, searchContext.getAggsCollector()),
+                    REASON_SEARCH_MULTI,
+                    collector,
+                    searchContext.getAggsCollector()
+                );
             }
             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;
+                collector = wrapWithProfilerCollectorIfNeeded(
+                    searchContext.getProfilers(),
+                    new MinimumScoreCollector(collector, searchContext.minimumScore()),
+                    REASON_SEARCH_MIN_SCORE,
+                    collector
+                );
             }
 
             boolean timeoutSet = scrollContext == null
@@ -171,7 +212,8 @@ public class QueryPhase {
             }
 
             try {
-                searchWithCollector(searchContext, searcher, query, collectors, hasFilterCollector, timeoutSet);
+                searchWithCollector(searchContext, searcher, query, collector, timeoutSet);
+                queryResult.topDocs(topDocsFactory.topDocsAndMaxScore(), topDocsFactory.sortValueFormats);
                 ExecutorService executor = searchContext.indexShard().getThreadPool().executor(ThreadPool.Names.SEARCH);
                 assert executor instanceof EWMATrackingEsThreadPoolExecutor
                     || (executor instanceof EsThreadPoolExecutor == false /* in case thread pool is mocked out in tests */)
@@ -192,30 +234,31 @@ public class QueryPhase {
         }
     }
 
+    private static Collector wrapWithProfilerCollectorIfNeeded(
+        Profilers profilers,
+        Collector collector,
+        String profilerName,
+        Collector... children
+    ) {
+        if (profilers == null) {
+            return collector;
+        }
+        return new InternalProfileCollector(collector, profilerName, children);
+    }
+
     private static void searchWithCollector(
         SearchContext searchContext,
         ContextIndexSearcher searcher,
         Query query,
-        LinkedList<QueryCollectorContext> collectors,
-        boolean hasFilterCollector,
+        Collector collector,
         boolean timeoutSet
     ) throws IOException {
-        // create the top docs collector last when the other collectors are known
-        final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
-        // add the top docs collector, the first collector context in the chain
-        collectors.addFirst(topDocsFactory);
-
-        final Collector queryCollector;
         if (searchContext.getProfilers() != null) {
-            InternalProfileCollector profileCollector = QueryCollectorContext.createQueryCollectorWithProfiler(collectors);
-            searchContext.getProfilers().getCurrentQueryProfiler().setCollector(profileCollector);
-            queryCollector = profileCollector;
-        } else {
-            queryCollector = QueryCollectorContext.createQueryCollector(collectors);
+            searchContext.getProfilers().getCurrentQueryProfiler().setCollector((InternalProfileCollector) collector);
         }
         QuerySearchResult queryResult = searchContext.queryResult();
         try {
-            searcher.search(query, queryCollector);
+            searcher.search(query, collector);
         } catch (EarlyTerminatingCollector.EarlyTerminationException e) {
             queryResult.terminatedEarly(true);
         } catch (TimeExceededException e) {
@@ -229,9 +272,6 @@ public class QueryPhase {
         if (searchContext.terminateAfter() != SearchContext.DEFAULT_TERMINATE_AFTER && queryResult.terminatedEarly() == null) {
             queryResult.terminatedEarly(false);
         }
-        for (QueryCollectorContext ctx : collectors) {
-            ctx.postProcess(queryResult);
-        }
     }
 
     /**
@@ -253,4 +293,14 @@ public class QueryPhase {
     }
 
     public static class TimeExceededException extends RuntimeException {}
+
+    private static final Collector EMPTY_COLLECTOR = new SimpleCollector() {
+        @Override
+        public void collect(int doc) {}
+
+        @Override
+        public ScoreMode scoreMode() {
+            return ScoreMode.COMPLETE_NO_SCORES;
+        }
+    };
 }

server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorContext.java

@@ -64,16 +64,21 @@ import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEAR
 import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TOP_HITS;
 
 /**
- * A {@link QueryCollectorContext} that creates top docs collector
+ * Creates and holds the main {@link Collector} that will be used to search and collect top hits.
  */
-abstract class TopDocsCollectorContext extends QueryCollectorContext {
-    protected final int numHits;
+abstract class TopDocsCollectorContext {
+    final String profilerName;
+    final DocValueFormat[] sortValueFormats;
 
-    TopDocsCollectorContext(String profilerName, int numHits) {
-        super(profilerName);
-        this.numHits = numHits;
+    TopDocsCollectorContext(String profilerName, DocValueFormat[] sortValueFormats) {
+        this.profilerName = profilerName;
+        this.sortValueFormats = sortValueFormats;
     }
 
+    abstract Collector collector();
+
+    abstract TopDocsAndMaxScore topDocsAndMaxScore() throws IOException;
+
     static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext {
         private final Sort sort;
         private final Collector collector;
@@ -85,7 +90,7 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
          * @param trackTotalHitsUpTo The threshold up to which total hit count needs to be tracked
          */
         private EmptyTopDocsCollectorContext(@Nullable SortAndFormats sortAndFormats, int trackTotalHitsUpTo) {
-            super(REASON_SEARCH_COUNT, 0);
+            super(REASON_SEARCH_COUNT, null);
             this.sort = sortAndFormats == null ? null : sortAndFormats.sort;
             if (trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_DISABLED) {
                 this.collector = new EarlyTerminatingCollector(new TotalHitCountCollector(), 0, false);
@@ -108,13 +113,12 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
         }
 
         @Override
-        Collector create(Collector in) {
-            assert in == null;
+        Collector collector() {
             return collector;
         }
 
         @Override
-        void postProcess(QuerySearchResult result) {
+        TopDocsAndMaxScore topDocsAndMaxScore() {
             final TotalHits totalHitCount = hitCountSupplier.get();
             final TopDocs topDocs;
             if (sort != null) {
@@ -122,12 +126,11 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
             } else {
                 topDocs = new TopDocs(totalHitCount, Lucene.EMPTY_SCORE_DOCS);
             }
-            result.topDocs(new TopDocsAndMaxScore(topDocs, Float.NaN), null);
+            return new TopDocsAndMaxScore(topDocs, Float.NaN);
         }
     }
 
     static class CollapsingTopDocsCollectorContext extends TopDocsCollectorContext {
-        private final DocValueFormat[] sortFmt;
         private final SinglePassGroupingCollector<?> topDocsCollector;
         private final Supplier<Float> maxScoreSupplier;
 
@@ -145,11 +148,10 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
             boolean trackMaxScore,
             @Nullable FieldDoc after
         ) {
-            super(REASON_SEARCH_TOP_HITS, numHits);
+            super(REASON_SEARCH_TOP_HITS, sortAndFormats == null ? new DocValueFormat[] { DocValueFormat.RAW } : sortAndFormats.formats);
             assert numHits > 0;
             assert collapseContext != null;
             Sort sort = sortAndFormats == null ? Sort.RELEVANCE : sortAndFormats.sort;
-            this.sortFmt = sortAndFormats == null ? new DocValueFormat[] { DocValueFormat.RAW } : sortAndFormats.formats;
             this.topDocsCollector = collapseContext.createTopDocs(sort, numHits, after);
 
             MaxScoreCollector maxScoreCollector;
@@ -162,15 +164,14 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
         }
 
         @Override
-        Collector create(Collector in) throws IOException {
-            assert in == null;
+        Collector collector() {
             return topDocsCollector;
         }
 
         @Override
-        void postProcess(QuerySearchResult result) throws IOException {
+        TopDocsAndMaxScore topDocsAndMaxScore() throws IOException {
             TopFieldGroups topDocs = topDocsCollector.getTopGroups(0);
-            result.topDocs(new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get()), sortFmt);
+            return new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get());
         }
     }
 
@@ -189,7 +190,6 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
             }
         }
 
-        protected final @Nullable SortAndFormats sortAndFormats;
         private final Collector collector;
         private final Supplier<TotalHits> totalHitsSupplier;
         private final Supplier<TopDocs> topDocsSupplier;
@@ -197,12 +197,13 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
 
         /**
          * Ctr
-         * @param reader The index reader
-         * @param query The Lucene query
-         * @param sortAndFormats The sort clause if provided
-         * @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 reader             The index reader
+         * @param query              The Lucene query
+         * @param sortAndFormats     The sort clause if provided
+         * @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 trackTotalHitsUpTo Threshold up to which total hit count should be tracked
          * @param hasFilterCollector True if the collector chain contains at least one collector that can filter documents out
          */
@@ -216,11 +217,9 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
             int trackTotalHitsUpTo,
             boolean hasFilterCollector
         ) throws IOException {
-            super(REASON_SEARCH_TOP_HITS, numHits);
-            this.sortAndFormats = sortAndFormats;
+            super(REASON_SEARCH_TOP_HITS, sortAndFormats == null ? null : sortAndFormats.formats);
 
             final TopDocsCollector<?> topDocsCollector;
-
             if ((sortAndFormats == null || SortField.FIELD_SCORE.equals(sortAndFormats.sort.getSort()[0])) && hasInfMaxScore(query)) {
                 // disable max score optimization since we have a mandatory clause
                 // that doesn't track the maximum score
@@ -268,8 +267,7 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
         }
 
         @Override
-        Collector create(Collector in) {
-            assert in == null;
+        Collector collector() {
             return collector;
         }
 
@@ -286,9 +284,8 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
         }
 
         @Override
-        void postProcess(QuerySearchResult result) throws IOException {
-            final TopDocsAndMaxScore topDocs = newTopDocs();
-            result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats);
+        TopDocsAndMaxScore topDocsAndMaxScore() {
+            return newTopDocs();
         }
     }
 
@@ -322,7 +319,7 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
         }
 
         @Override
-        void postProcess(QuerySearchResult result) throws IOException {
+        TopDocsAndMaxScore topDocsAndMaxScore() {
             final TopDocsAndMaxScore topDocs = newTopDocs();
             if (scrollContext.totalHits == null) {
                 // first round
@@ -341,7 +338,7 @@ abstract class TopDocsCollectorContext extends QueryCollectorContext {
                     scrollContext.lastEmittedDoc = topDocs.topDocs.scoreDocs[topDocs.topDocs.scoreDocs.length - 1];
                 }
             }
-            result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats);
+            return topDocs;
         }
     }
 

server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java

@@ -680,7 +680,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
         context.setSize(3);
         context.trackTotalHitsUpTo(3);
         TopDocsCollectorContext topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false);
-        assertEquals(topDocsContext.create(null).scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE);
+        assertEquals(topDocsContext.collector().scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE);
         QueryPhase.executeInternal(context);
         assertEquals(5, context.queryResult().topDocs().topDocs.totalHits.value);
         assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.EQUAL_TO);
@@ -688,7 +688,7 @@ public class QueryPhaseTests extends IndexShardTestCase {
 
         context.sort(new SortAndFormats(new Sort(new SortField("other", SortField.Type.INT)), new DocValueFormat[] { DocValueFormat.RAW }));
         topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false);
-        assertEquals(topDocsContext.create(null).scoreMode(), org.apache.lucene.search.ScoreMode.TOP_DOCS);
+        assertEquals(topDocsContext.collector().scoreMode(), org.apache.lucene.search.ScoreMode.TOP_DOCS);
         QueryPhase.executeInternal(context);
         assertEquals(5, context.queryResult().topDocs().topDocs.totalHits.value);
         assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3));