Browse Source

Remove Collector implementation from BucketCollector (#88444)

BucketCollector has now a method called #asCollector that returns the current BucketCollector wrapped as a 
Lucene Collector.
Ignacio Vera 3 years ago
parent
commit
04bdefd58c
26 changed files with 115 additions and 85 deletions
  1. 1 1
      docs/reference/search/profile.asciidoc
  2. 2 2
      server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java
  3. 39 7
      server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java
  4. 5 5
      server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java
  5. 3 3
      server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregator.java
  6. 2 2
      server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java
  7. 9 11
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java
  8. 3 3
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/filter/FiltersAggregatorTests.java
  9. 1 1
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramAggregatorTests.java
  10. 1 1
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java
  11. 1 1
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/RareTermsAggregatorTests.java
  12. 14 14
      server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/TermsAggregatorTests.java
  13. 5 5
      server/src/test/java/org/elasticsearch/search/aggregations/metrics/AvgAggregatorTests.java
  14. 2 2
      server/src/test/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorTests.java
  15. 1 1
      server/src/test/java/org/elasticsearch/search/aggregations/metrics/HDRPercentilesAggregatorTests.java
  16. 9 9
      server/src/test/java/org/elasticsearch/search/aggregations/metrics/MaxAggregatorTests.java
  17. 1 1
      server/src/test/java/org/elasticsearch/search/aggregations/metrics/TDigestPercentilesAggregatorTests.java
  18. 1 1
      server/src/test/java/org/elasticsearch/search/aggregations/metrics/WeightedAvgAggregatorTests.java
  19. 3 3
      test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java
  20. 1 1
      test/framework/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregatorTestCase.java
  21. 1 1
      x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/aggregations/metrics/HDRPreAggregatedPercentilesAggregatorTests.java
  22. 1 1
      x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/aggregations/metrics/TDigestPreAggregatedPercentilesAggregatorTests.java
  23. 1 1
      x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorTests.java
  24. 1 1
      x-pack/plugin/rollup/src/test/java/org/elasticsearch/xpack/rollup/RollupResponseTranslationTests.java
  25. 6 6
      x-pack/plugin/rollup/src/test/java/org/elasticsearch/xpack/rollup/job/IndexerUtilsTests.java
  26. 1 1
      x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/search/aggregations/bucket/geogrid/GeoShapeGeoGridTestCase.java

+ 1 - 1
docs/reference/search/profile.asciidoc

@@ -696,7 +696,7 @@ The API returns the following result:
                     ]
                   },
                   {
-                    "name": "MultiBucketCollector: [[my_scoped_agg, my_global_agg]]",
+                    "name": "BucketCollectorWrapper: [BucketCollectorWrapper[bucketCollector=[my_scoped_agg, my_global_agg]]]",
                     "reason": "aggregation",
                     "time_in_nanos": 867617
                   }

+ 2 - 2
server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java

@@ -52,8 +52,8 @@ public class AggregationPhase {
             context.queryCollectors().put(AggregationPhase.class, BucketCollector.NO_OP_COLLECTOR);
         } else {
             Collector collector = context.getProfilers() == null
-                ? bucketCollector
-                : new InternalProfileCollector(bucketCollector, CollectorResult.REASON_AGGREGATION, List.of());
+                ? bucketCollector.asCollector()
+                : new InternalProfileCollector(bucketCollector.asCollector(), CollectorResult.REASON_AGGREGATION, List.of());
             context.queryCollectors().put(AggregationPhase.class, collector);
         }
     }

+ 39 - 7
server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java

@@ -10,6 +10,7 @@ package org.elasticsearch.search.aggregations;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.ScoreMode;
 
 import java.io.IOException;
@@ -17,9 +18,9 @@ import java.io.IOException;
 /**
  * A Collector that can collect data in separate buckets.
  */
-public abstract class BucketCollector implements Collector {
+public abstract class BucketCollector {
 
-    public static final BucketCollector NO_OP_COLLECTOR = new BucketCollector() {
+    public static final BucketCollector NO_OP_BUCKET_COLLECTOR = new BucketCollector() {
 
         @Override
         public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) {
@@ -41,11 +42,18 @@ public abstract class BucketCollector implements Collector {
         }
     };
 
-    // TODO: will remove it in a follow up PR
-    @Override
-    public final LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException {
-        return getLeafCollector(new AggregationExecutionContext(ctx, null, null));
-    }
+    public static final Collector NO_OP_COLLECTOR = new Collector() {
+
+        @Override
+        public LeafCollector getLeafCollector(LeafReaderContext context) {
+            return LeafBucketCollector.NO_OP_COLLECTOR;
+        }
+
+        @Override
+        public ScoreMode scoreMode() {
+            return ScoreMode.COMPLETE_NO_SCORES;
+        }
+    };
 
     public abstract LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException;
 
@@ -59,4 +67,28 @@ public abstract class BucketCollector implements Collector {
      */
     public abstract void postCollection() throws IOException;
 
+    /**
+     *  Indicates what features are required from the scorer.
+     */
+    public abstract ScoreMode scoreMode();
+
+    /**
+     * Return this BucketCollector wrapped as a {@link Collector}
+     */
+    public final Collector asCollector() {
+        return new BucketCollectorWrapper(this);
+    }
+
+    private record BucketCollectorWrapper(BucketCollector bucketCollector) implements Collector {
+
+        @Override
+        public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
+            return bucketCollector.getLeafCollector(new AggregationExecutionContext(context, null, null));
+        }
+
+        @Override
+        public ScoreMode scoreMode() {
+            return bucketCollector.scoreMode();
+        }
+    }
 }

+ 5 - 5
server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java

@@ -23,7 +23,7 @@ import java.util.List;
 /**
  * A {@link BucketCollector} which allows running a bucket collection with several
  * {@link BucketCollector}s. It is similar to the {@link MultiCollector} except that the
- * {@link #wrap} method filters out the {@link BucketCollector#NO_OP_COLLECTOR}s and not
+ * {@link #wrap} method filters out the {@link BucketCollector#NO_OP_BUCKET_COLLECTOR}s and not
  * the null ones.
  */
 public class MultiBucketCollector extends BucketCollector {
@@ -31,12 +31,12 @@ public class MultiBucketCollector extends BucketCollector {
      * Wraps a list of {@link BucketCollector}s with a {@link MultiBucketCollector}. This
      * method works as follows:
      * <ul>
-     * <li>Filters out the {@link BucketCollector#NO_OP_COLLECTOR}s collectors, so they are not used
+     * <li>Filters out the {@link BucketCollector#NO_OP_BUCKET_COLLECTOR}s collectors, so they are not used
      * during search time.
      * <li>If the input contains 1 real collector we wrap it in a collector that takes
      * {@code terminateIfNoop} into account.
      * <li>Otherwise the method returns a {@link MultiBucketCollector} which wraps the
-     * non-{@link BucketCollector#NO_OP_COLLECTOR} collectors.
+     * non-{@link BucketCollector#NO_OP_BUCKET_COLLECTOR} collectors.
      * </ul>
      * @param terminateIfNoop Pass true if {@link #getLeafCollector} should throw
      * {@link CollectionTerminatedException} if all leaf collectors are noop. Pass
@@ -52,13 +52,13 @@ public class MultiBucketCollector extends BucketCollector {
         // and dropped from the array we save for actual collection time.
         int n = 0;
         for (BucketCollector c : collectors) {
-            if (c != NO_OP_COLLECTOR) {
+            if (c != NO_OP_BUCKET_COLLECTOR) {
                 n++;
             }
         }
 
         if (n == 0) {
-            return NO_OP_COLLECTOR;
+            return NO_OP_BUCKET_COLLECTOR;
         } else if (n == 1) {
             // only 1 Collector - return it.
             BucketCollector col = null;

+ 3 - 3
server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -158,7 +158,7 @@ public final class CompositeAggregator extends BucketsAggregator implements Size
     @Override
     protected void doPreCollection() throws IOException {
         deferredCollectors = MultiBucketCollector.wrap(false, Arrays.asList(subAggregators));
-        collectableSubAggregators = BucketCollector.NO_OP_COLLECTOR;
+        collectableSubAggregators = BucketCollector.NO_OP_BUCKET_COLLECTOR;
     }
 
     @Override
@@ -170,7 +170,7 @@ public final class CompositeAggregator extends BucketsAggregator implements Size
     public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
         // Composite aggregator must be at the top of the aggregation tree
         assert owningBucketOrds.length == 1 && owningBucketOrds[0] == 0L;
-        if (deferredCollectors != NO_OP_COLLECTOR) {
+        if (deferredCollectors != NO_OP_BUCKET_COLLECTOR) {
             // Replay all documents that contain at least one top bucket (collected during the first pass).
             runDeferredCollections();
         }
@@ -440,7 +440,7 @@ public final class CompositeAggregator extends BucketsAggregator implements Size
     protected LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx, LeafBucketCollector sub) throws IOException {
         finishLeaf();
 
-        boolean fillDocIdSet = deferredCollectors != NO_OP_COLLECTOR;
+        boolean fillDocIdSet = deferredCollectors != NO_OP_BUCKET_COLLECTOR;
 
         Sort indexSortPrefix = buildIndexSortPrefix(aggCtx.getLeafReaderContext());
         int sortPrefixLen = computeSortPrefixLen(indexSortPrefix);

+ 2 - 2
server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java

@@ -176,7 +176,7 @@ public class MultiBucketCollectorTests extends ESTestCase {
                 expectedCounts.put(collector, expectedCount);
                 collectors.add(new TerminateAfterBucketCollector(collector, terminateAfter));
             }
-            searcher.search(new MatchAllDocsQuery(), MultiBucketCollector.wrap(true, collectors));
+            searcher.search(new MatchAllDocsQuery(), MultiBucketCollector.wrap(true, collectors).asCollector());
             for (Map.Entry<TotalHitCountBucketCollector, Integer> expectedCount : expectedCounts.entrySet()) {
                 assertEquals(expectedCount.getValue().intValue(), expectedCount.getKey().getTotalHits());
             }
@@ -252,7 +252,7 @@ public class MultiBucketCollectorTests extends ESTestCase {
         Collections.shuffle(collectors, random());
         BucketCollector collector = MultiBucketCollector.wrap(true, collectors);
 
-        LeafBucketCollector leafCollector = collector.getLeafCollector((LeafReaderContext) null);
+        LeafBucketCollector leafCollector = collector.getLeafCollector(null);
         leafCollector.setScorer(scorer);
         assertTrue(setScorerCalled1.get());
         assertTrue(setScorerCalled2.get());

+ 9 - 11
server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java

@@ -14,7 +14,9 @@ import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
+import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
@@ -72,7 +74,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
         Set<Integer> deferredCollectedDocIds = new HashSet<>();
         collector.setDeferredCollector(Collections.singleton(bla(deferredCollectedDocIds)));
         collector.preCollection();
-        indexSearcher.search(termQuery, collector);
+        indexSearcher.search(termQuery, collector.asCollector());
         collector.postCollection();
         collector.prepareSelectedBuckets(0);
 
@@ -86,7 +88,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
         deferredCollectedDocIds = new HashSet<>();
         collector.setDeferredCollector(Collections.singleton(bla(deferredCollectedDocIds)));
         collector.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), collector);
+        indexSearcher.search(new MatchAllDocsQuery(), collector.asCollector());
         collector.postCollection();
         collector.prepareSelectedBuckets(0);
 
@@ -199,21 +201,17 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
                 CollectingBucketCollector finalCollector = new CollectingBucketCollector();
                 deferringCollector.setDeferredCollector(Collections.singleton(finalCollector));
                 deferringCollector.preCollection();
-                indexSearcher.search(query, new BucketCollector() {
+                indexSearcher.search(query, new Collector() {
                     @Override
                     public ScoreMode scoreMode() {
                         return ScoreMode.COMPLETE_NO_SCORES;
                     }
 
                     @Override
-                    public void preCollection() throws IOException {}
-
-                    @Override
-                    public void postCollection() throws IOException {}
-
-                    @Override
-                    public LeafBucketCollector getLeafCollector(AggregationExecutionContext aggCtx) throws IOException {
-                        LeafBucketCollector delegate = deferringCollector.getLeafCollector(aggCtx);
+                    public LeafBucketCollector getLeafCollector(LeafReaderContext context) throws IOException {
+                        LeafBucketCollector delegate = deferringCollector.getLeafCollector(
+                            new AggregationExecutionContext(context, null, null)
+                        );
                         return leafCollector.apply(deferringCollector, delegate);
                     }
                 });

+ 3 - 3
server/src/test/java/org/elasticsearch/search/aggregations/bucket/filter/FiltersAggregatorTests.java

@@ -669,7 +669,7 @@ public class FiltersAggregatorTests extends AggregatorTestCase {
                 AggregationContext context = createAggregationContext(searcher, new MatchAllDocsQuery());
                 FilterByFilterAggregator aggregator = createAggregator(builder, context);
                 aggregator.preCollection();
-                searcher.search(context.query(), aggregator);
+                searcher.search(context.query(), aggregator.asCollector());
                 aggregator.postCollection();
 
                 InternalAggregation result = aggregator.buildTopLevel();
@@ -746,7 +746,7 @@ public class FiltersAggregatorTests extends AggregatorTestCase {
                 AggregationContext context = createAggregationContext(searcher, new MatchAllDocsQuery(), ft);
                 FilterByFilterAggregator aggregator = createAggregator(builder, context);
                 aggregator.preCollection();
-                searcher.search(context.query(), aggregator);
+                searcher.search(context.query(), aggregator.asCollector());
                 aggregator.postCollection();
 
                 InternalAggregation result = aggregator.buildTopLevel();
@@ -812,7 +812,7 @@ public class FiltersAggregatorTests extends AggregatorTestCase {
                 AggregationContext context = createAggregationContext(searcher, new MatchAllDocsQuery(), ft);
                 FilterByFilterAggregator aggregator = createAggregator(builder, context);
                 aggregator.preCollection();
-                searcher.search(context.query(), aggregator);
+                searcher.search(context.query(), aggregator.asCollector());
                 aggregator.postCollection();
 
                 InternalAggregation result = aggregator.buildTopLevel();

+ 1 - 1
server/src/test/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramAggregatorTests.java

@@ -1038,7 +1038,7 @@ public class DateHistogramAggregatorTests extends DateHistogramAggregatorTestCas
                 }
                 assertThat(agg, matcher);
                 agg.preCollection();
-                context.searcher().search(context.query(), agg);
+                context.searcher().search(context.query(), agg.asCollector());
                 InternalDateHistogram result = (InternalDateHistogram) agg.buildTopLevel();
                 result = (InternalDateHistogram) result.reduce(
                     List.of(result),

+ 1 - 1
server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollectorTests.java

@@ -65,7 +65,7 @@ public class BestDocsDeferringCollectorTests extends AggregatorTestCase {
         Set<Integer> deferredCollectedDocIds = new HashSet<>();
         collector.setDeferredCollector(Collections.singleton(testCollector(deferredCollectedDocIds)));
         collector.preCollection();
-        indexSearcher.search(termQuery, collector);
+        indexSearcher.search(termQuery, collector.asCollector());
         collector.postCollection();
         collector.prepareSelectedBuckets(0);
 

+ 1 - 1
server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/RareTermsAggregatorTests.java

@@ -250,7 +250,7 @@ public class RareTermsAggregatorTests extends AggregatorTestCase {
                         RareTermsAggregationBuilder aggregationBuilder = new RareTermsAggregationBuilder("_name").field(fieldNames[i]);
                         Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType1, fieldType2);
                         aggregator.preCollection();
-                        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                         aggregator.postCollection();
                         RareTerms result = (RareTerms) aggregator.buildTopLevel();
                         assertEquals("_name", result.getName());

+ 14 - 14
server/src/test/java/org/elasticsearch/search/aggregations/bucket/terms/TermsAggregatorTests.java

@@ -692,7 +692,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
                     TermsAggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
@@ -710,7 +710,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     context = createAggregationContext(indexSearcher, null, fieldType);
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(4, result.getBuckets().size());
@@ -735,7 +735,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     context = createAggregationContext(indexSearcher, null, fieldType);
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
@@ -755,7 +755,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     context = createAggregationContext(indexSearcher, null, fieldType);
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(4, result.getBuckets().size());
@@ -930,7 +930,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     AggregationContext context = createAggregationContext(indexSearcher, new MatchAllDocsQuery(), fieldType);
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(size, result.getBuckets().size());
@@ -958,7 +958,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         context = createAggregationContext(indexSearcher, new MatchAllDocsQuery(), fieldType, filterFieldType);
                         aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
-                        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                         aggregator.postCollection();
                         result = ((Filter) reduce(aggregationBuilder, aggregator, context.bigArrays())).getAggregations().get("_name2");
                         int expectedFilteredCounts = 0;
@@ -1038,7 +1038,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     AggregationContext context = createAggregationContext(indexSearcher, new MatchAllDocsQuery(), fieldType, fieldType2);
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(size, result.getBuckets().size());
@@ -1065,7 +1065,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     AggregationContext context = createAggregationContext(indexSearcher, null, fieldType1);
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
@@ -1075,7 +1075,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     context = createAggregationContext(indexSearcher, null, fieldType2);
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
@@ -1085,7 +1085,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     context = createAggregationContext(indexSearcher, null, fieldType3);
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
@@ -1108,7 +1108,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         AggregationContext context = createAggregationContext(indexSearcher, null);
                         Aggregator aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
-                        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                         aggregator.postCollection();
                         Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                         assertEquals("_name", result.getName());
@@ -1144,7 +1144,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         AggregationContext context = createAggregationContext(indexSearcher, new MatchAllDocsQuery(), fieldType1);
                         Aggregator aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
-                        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                         aggregator.postCollection();
                         Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                         assertEquals("_name", result.getName());
@@ -1260,7 +1260,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     AggregationContext context = createAggregationContext(indexSearcher, new MatchAllDocsQuery(), fieldType1, fieldType2);
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
-                    indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
                     aggregator.postCollection();
                     Terms result = reduce(aggregationBuilder, aggregator, context.bigArrays());
                     assertEquals(3, result.getBuckets().size());
@@ -2325,7 +2325,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
         throws IOException {
         TermsAggregator aggregator = createAggregator(builder, searcher, fieldType);
         aggregator.preCollection();
-        searcher.search(new MatchAllDocsQuery(), aggregator);
+        searcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         return aggregator.buildTopLevel();
     }

+ 5 - 5
server/src/test/java/org/elasticsearch/search/aggregations/metrics/AvgAggregatorTests.java

@@ -283,7 +283,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
 
         AvgAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         InternalAvg avg = (InternalAvg) aggregator.buildAggregation(0L);
@@ -544,7 +544,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
 
         TermsAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Terms terms = (Terms) aggregator.buildTopLevel();
@@ -616,7 +616,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
         AvgAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         InternalAvg avg = (InternalAvg) aggregator.buildAggregation(0L);
@@ -662,7 +662,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
         AvgAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         InternalAvg avg = (InternalAvg) aggregator.buildAggregation(0L);
@@ -680,7 +680,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         context = createAggregationContext(indexSearcher, null, fieldType);
         aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         avg = (InternalAvg) aggregator.buildAggregation(0L);

+ 2 - 2
server/src/test/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregatorTests.java

@@ -432,7 +432,7 @@ public class CardinalityAggregatorTests extends AggregatorTestCase {
 
         final CardinalityAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         final InternalCardinality cardinality = (InternalCardinality) aggregator.buildAggregation(0L);
@@ -642,7 +642,7 @@ public class CardinalityAggregatorTests extends AggregatorTestCase {
         final AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
         final CardinalityAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         final InternalCardinality cardinality = (InternalCardinality) aggregator.buildAggregation(0L);

+ 1 - 1
server/src/test/java/org/elasticsearch/search/aggregations/metrics/HDRPercentilesAggregatorTests.java

@@ -207,7 +207,7 @@ public class HDRPercentilesAggregatorTests extends AggregatorTestCase {
 
                 HDRPercentilesAggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
-                indexSearcher.search(query, aggregator);
+                indexSearcher.search(query, aggregator.asCollector());
                 aggregator.postCollection();
                 verify.accept((InternalHDRPercentiles) aggregator.buildAggregation(0L));
 

+ 9 - 9
server/src/test/java/org/elasticsearch/search/aggregations/metrics/MaxAggregatorTests.java

@@ -441,7 +441,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
 
         GlobalAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Global global = (Global) aggregator.buildTopLevel();
@@ -486,7 +486,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
 
         MaxAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Max max = (Max) aggregator.buildAggregation(0L);
@@ -695,7 +695,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
 
         GlobalAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Global global = (Global) aggregator.buildTopLevel();
@@ -737,7 +737,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
 
         TermsAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Terms terms = (Terms) aggregator.buildTopLevel();
@@ -790,7 +790,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
 
         BucketCollector bucketCollector = MultiBucketCollector.wrap(true, List.of(maxAggregator, countAggregator));
         bucketCollector.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), bucketCollector);
+        indexSearcher.search(new MatchAllDocsQuery(), bucketCollector.asCollector());
         bucketCollector.postCollection();
 
         Max max = (Max) maxAggregator.buildAggregation(0L);
@@ -840,7 +840,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
 
             BucketCollector bucketCollector = MultiBucketCollector.wrap(true, List.of(maxAggregator, countAggregator, termsAggregator));
             bucketCollector.preCollection();
-            indexSearcher.search(new MatchAllDocsQuery(), bucketCollector);
+            indexSearcher.search(new MatchAllDocsQuery(), bucketCollector.asCollector());
             bucketCollector.postCollection();
 
             Max max = (Max) maxAggregator.buildTopLevel();
@@ -896,7 +896,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
         MaxAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Max max = (Max) aggregator.buildAggregation(0L);
@@ -942,7 +942,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         AggregationContext context = createAggregationContext(indexSearcher, null, fieldType);
         MaxAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Max max = (Max) aggregator.buildAggregation(0L);
@@ -959,7 +959,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         context = createAggregationContext(indexSearcher, null, fieldType);
         aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         max = (Max) aggregator.buildAggregation(0L);

+ 1 - 1
server/src/test/java/org/elasticsearch/search/aggregations/metrics/TDigestPercentilesAggregatorTests.java

@@ -182,7 +182,7 @@ public class TDigestPercentilesAggregatorTests extends AggregatorTestCase {
                 MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType("number", NumberFieldMapper.NumberType.LONG);
                 TDigestPercentilesAggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
-                indexSearcher.search(query, aggregator);
+                indexSearcher.search(query, aggregator.asCollector());
                 aggregator.postCollection();
                 verify.accept((InternalTDigestPercentiles) aggregator.buildAggregation(0L));
             }

+ 1 - 1
server/src/test/java/org/elasticsearch/search/aggregations/metrics/WeightedAvgAggregatorTests.java

@@ -531,7 +531,7 @@ public class WeightedAvgAggregatorTests extends AggregatorTestCase {
             MappedFieldType fieldType2 = new NumberFieldMapper.NumberFieldType("weight_field", fieldNumberType);
             WeightedAvgAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType, fieldType2);
             aggregator.preCollection();
-            indexSearcher.search(query, aggregator);
+            indexSearcher.search(query, aggregator.asCollector());
             aggregator.postCollection();
             verify.accept((InternalWeightedAvg) aggregator.buildAggregation(0L));
         } finally {

+ 3 - 3
test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java

@@ -574,7 +574,7 @@ public abstract class AggregatorTestCase extends ESTestCase {
                     new TimeSeriesIndexSearcher(subSearcher, List.of()).search(rewritten, a);
                 } else {
                     Weight weight = subSearcher.createWeight(rewritten, ScoreMode.COMPLETE, 1f);
-                    subSearcher.search(weight, a);
+                    subSearcher.search(weight, a.asCollector());
                 }
                 a.postCollection();
                 aggs.add(a.buildTopLevel());
@@ -584,7 +584,7 @@ public abstract class AggregatorTestCase extends ESTestCase {
             if (context.isInSortOrderExecutionRequired()) {
                 new TimeSeriesIndexSearcher(searcher, List.of()).search(rewritten, MultiBucketCollector.wrap(true, List.of(root)));
             } else {
-                searcher.search(rewritten, MultiBucketCollector.wrap(true, List.of(root)));
+                searcher.search(rewritten, MultiBucketCollector.wrap(true, List.of(root)).asCollector());
             }
             root.postCollection();
             aggs.add(root.buildTopLevel());
@@ -753,7 +753,7 @@ public abstract class AggregatorTestCase extends ESTestCase {
         );
         Aggregator aggregator = createAggregator(builder, context);
         aggregator.preCollection();
-        searcher.search(context.query(), aggregator);
+        searcher.search(context.query(), aggregator.asCollector());
         aggregator.postCollection();
         InternalAggregation r = aggregator.buildTopLevel();
         r = r.reduce(

+ 1 - 1
test/framework/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregatorTestCase.java

@@ -342,7 +342,7 @@ public abstract class GeoGridAggregatorTestCase<T extends InternalGeoGridBucket>
 
         Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(query, aggregator);
+        indexSearcher.search(query, aggregator.asCollector());
         aggregator.postCollection();
         @SuppressWarnings("unchecked")
         InternalGeoGrid<T> topLevel = (InternalGeoGrid<T>) aggregator.buildTopLevel();

+ 1 - 1
x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/aggregations/metrics/HDRPreAggregatedPercentilesAggregatorTests.java

@@ -152,7 +152,7 @@ public class HDRPreAggregatedPercentilesAggregatorTests extends AggregatorTestCa
                 MappedFieldType fieldType = new HistogramFieldMapper.HistogramFieldType("number", Collections.emptyMap(), null);
                 Aggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
-                indexSearcher.search(query, aggregator);
+                indexSearcher.search(query, aggregator.asCollector());
                 aggregator.postCollection();
                 verify.accept((InternalHDRPercentiles) aggregator.buildTopLevel());
 

+ 1 - 1
x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/aggregations/metrics/TDigestPreAggregatedPercentilesAggregatorTests.java

@@ -133,7 +133,7 @@ public class TDigestPreAggregatedPercentilesAggregatorTests extends AggregatorTe
                 MappedFieldType fieldType = new HistogramFieldMapper.HistogramFieldType("number", Collections.emptyMap(), null);
                 Aggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
-                indexSearcher.search(query, aggregator);
+                indexSearcher.search(query, aggregator.asCollector());
                 aggregator.postCollection();
                 verify.accept((InternalTDigestPercentiles) aggregator.buildTopLevel());
 

+ 1 - 1
x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/stringstats/StringStatsAggregatorTests.java

@@ -241,7 +241,7 @@ public class StringStatsAggregatorTests extends AggregatorTestCase {
 
         TermsAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, numericFieldType, textFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
 
         Terms terms = (Terms) aggregator.buildTopLevel();

+ 1 - 1
x-pack/plugin/rollup/src/test/java/org/elasticsearch/xpack/rollup/RollupResponseTranslationTests.java

@@ -1316,7 +1316,7 @@ public class RollupResponseTranslationTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(aggBuilder, indexSearcher, fieldType);
         try {
             aggregator.preCollection();
-            indexSearcher.search(query, aggregator);
+            indexSearcher.search(query, aggregator.asCollector());
             aggregator.postCollection();
             return aggregator.buildTopLevel();
         } finally {

+ 6 - 6
x-pack/plugin/rollup/src/test/java/org/elasticsearch/xpack/rollup/job/IndexerUtilsTests.java

@@ -112,7 +112,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
 
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
@@ -174,7 +174,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
 
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
@@ -228,7 +228,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
 
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, valueFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
@@ -289,7 +289,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
 
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
@@ -477,7 +477,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
 
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, valueFieldType, metricFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
@@ -551,7 +551,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
 
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        indexSearcher.search(new MatchAllDocsQuery(), aggregator.asCollector());
         aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();

+ 1 - 1
x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/search/aggregations/bucket/geogrid/GeoShapeGeoGridTestCase.java

@@ -287,7 +287,7 @@ public abstract class GeoShapeGeoGridTestCase<T extends InternalGeoGridBucket> e
 
         Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
-        indexSearcher.search(query, aggregator);
+        indexSearcher.search(query, aggregator.asCollector());
         aggregator.postCollection();
 
         verify.accept((InternalGeoGrid<T>) aggregator.buildTopLevel());