Revert "Remove aggregation's postCollect phase (#68615)

This partially reverts #64016 and  and adds #67839 and adds
additional tests that would have caught issues with the changes
in #64016. It's mostly Nik's code, I am just cleaning things up
a bit.

Co-authored-by: Nik Everett <nik9000@gmail.com>

docs/reference/search/profile.asciidoc

@@ -799,7 +799,9 @@ This yields the following aggregation profile output:
               "collect": 45786,
               "collect_count": 4,
               "build_leaf_collector": 18211,
-              "build_leaf_collector_count": 1
+              "build_leaf_collector_count": 1,
+              "post_collection": 929,
+              "post_collection_count": 1
             },
             "debug": {
               "total_buckets": 1,
@@ -820,7 +822,9 @@ This yields the following aggregation profile output:
               "collect": 69401,
               "collect_count": 4,
               "build_leaf_collector": 8150,
-              "build_leaf_collector_count": 1
+              "build_leaf_collector_count": 1,
+              "post_collection": 1584,
+              "post_collection_count": 1
             },
             "children": [
               {
@@ -837,7 +841,9 @@ This yields the following aggregation profile output:
                   "collect": 61611,
                   "collect_count": 4,
                   "build_leaf_collector": 5564,
-                  "build_leaf_collector_count": 1
+                  "build_leaf_collector_count": 1,
+                  "post_collection": 471,
+                  "post_collection_count": 1
                 },
                 "debug": {
                   "total_buckets": 1,

modules/parent-join/src/main/java/org/elasticsearch/join/aggregations/ParentJoinAggregator.java

@@ -102,7 +102,12 @@ public abstract class ParentJoinAggregator extends BucketsAggregator implements
     }
 
     @Override
-    protected void prepareSubAggs(long[] bucketOrdsToCollect) throws IOException {
+    public void postCollection() throws IOException {
+        // Delaying until beforeBuildingBuckets
+    }
+
+    @Override
+    protected void prepareSubAggs(long[] ordsToCollect) throws IOException {
         IndexReader indexReader = searcher().getIndexReader();
         for (LeafReaderContext ctx : indexReader.leaves()) {
             Scorer childDocsScorer = outFilter.scorer(ctx);
@@ -144,13 +149,14 @@ public abstract class ParentJoinAggregator extends BucketsAggregator implements
                  * structure that maps a primitive long to a list of primitive
                  * longs.
                  */
-                for (long o: bucketOrdsToCollect) {
-                    if (collectionStrategy.exists(o, globalOrdinal)) {
-                        collectBucket(sub, docId, o);
+                for (long owningBucketOrd: ordsToCollect) {
+                    if (collectionStrategy.exists(owningBucketOrd, globalOrdinal)) {
+                        collectBucket(sub, docId, owningBucketOrd);
                     }
                 }
             }
         }
+        super.postCollection(); // Run post collection after collecting the sub-aggs
     }
 
     @Override

modules/parent-join/src/yamlRestTest/resources/rest-api-spec/test/50_order_by.yml

@@ -0,0 +1,238 @@
+---
+"order by sub agg containing join":
+  - skip:
+        reason: "https://github.com/elastic/elasticsearch/issues/66876"
+        version: "7.11.1 - "
+  - do:
+      indices.create:
+          index: test_1
+          body:
+            settings:
+              number_of_shards: 1
+              number_of_replicas: 0
+            mappings:
+              properties:
+                name:
+                  type: keyword
+                join:
+                  type: join
+                  relations:
+                    animal: feature
+
+  - do:
+      bulk:
+        index: test_1
+        refresh: true
+        body: |
+          { "index": {"_id": "sheep"} }
+          { "name": "sheep", "join": {"name": "animal"} }
+          { "index": {"_id": "cow"} }
+          { "name": "cow", "join": {"name": "animal"} }
+          { "index": {"_id": "pig"} }
+          { "name": "pig", "join": {"name": "animal"} }
+          { "index": {"routing": "sheep"} }
+          { "join": {"name": "feature", "parent": "sheep"}, "tag": "danger", "number": 1 }
+          { "index": {"routing": "sheep"} }
+          { "join": {"name": "feature", "parent": "sheep"}, "tag": "fluffiness", "number": 10 }
+          { "index": {"routing": "cow"} }
+          { "join": {"name": "feature", "parent": "cow"}, "tag": "danger", "number": 3 }
+          { "index": {"routing": "cow"} }
+          { "join": {"name": "feature", "parent": "cow"}, "tag": "fluffiness", "number": 1 }
+          { "index": {"routing": "pig"} }
+          { "join": {"name": "feature", "parent": "pig"}, "tag": "danger", "number": 100 }
+          { "index": {"routing": "pig"} }
+          { "join": {"name": "feature", "parent": "pig"}, "tag": "fluffiness", "number": 1 }
+
+  - do:
+      search:
+        index: test_1
+        body:
+          size: 0
+          aggs:
+            name:
+              terms:
+                size: 1
+                shard_size: 1
+                field: name
+                order:
+                  - "children>max_number": desc
+              aggs:
+                children:
+                  children:
+                    type: feature
+                  aggs:
+                    max_number:
+                      max:
+                        field: number
+  - length: { aggregations.name.buckets: 1 }
+  - match: { aggregations.name.buckets.0.key: pig }
+  - match: { aggregations.name.buckets.0.doc_count: 1 }
+  - match: { aggregations.name.buckets.0.children.max_number.value: 100.0 }
+
+---
+"order by sub agg containing join and nested":
+  - skip:
+      reason: "https://github.com/elastic/elasticsearch/issues/66876"
+      version: "7.11.1 - "
+  - do:
+      indices.create:
+          index: test_1
+          body:
+            settings:
+              number_of_shards: 1
+              number_of_replicas: 0
+            mappings:
+              properties:
+                name:
+                  type: keyword
+                join:
+                  type: join
+                  relations:
+                    animal: feature
+                nested:
+                  type: nested
+                  properties:
+                    number:
+                      type: long
+
+  - do:
+      bulk:
+        index: test_1
+        refresh: true
+        body: |
+          { "index": {"_id": "sheep"} }
+          { "name": "sheep", "join": {"name": "animal"} }
+          { "index": {"_id": "cow"} }
+          { "name": "cow", "join": {"name": "animal"} }
+          { "index": {"_id": "pig"} }
+          { "name": "pig", "join": {"name": "animal"} }
+          { "index": {"routing": "sheep"} }
+          { "join": {"name": "feature", "parent": "sheep"}, "tag": "danger", "nested": [{"number": 1}] }
+          { "index": {"routing": "sheep"} }
+          { "join": {"name": "feature", "parent": "sheep"}, "tag": "fluffiness", "nested": [{"number": 10}] }
+          { "index": {"routing": "cow"} }
+          { "join": {"name": "feature", "parent": "cow"}, "tag": "danger", "nested": [{"number": 3}] }
+          { "index": {"routing": "cow"} }
+          { "join": {"name": "feature", "parent": "cow"}, "tag": "fluffiness", "nested": [{"number": 1}] }
+          { "index": {"routing": "pig"} }
+          { "join": {"name": "feature", "parent": "pig"}, "tag": "danger", "nested": [{"number": 100}, {"number": 1}] }
+          { "index": {"routing": "pig"} }
+          { "join": {"name": "feature", "parent": "pig"}, "tag": "fluffiness", "nested": [{"number": 1}] }
+
+  - do:
+      search:
+        index: test_1
+        body:
+          size: 0
+          aggs:
+            name:
+              terms:
+                size: 1
+                shard_size: 1
+                field: name
+                order:
+                  - "children>nested>max_number": desc
+              aggs:
+                children:
+                  children:
+                    type: feature
+                  aggs:
+                    nested:
+                      nested:
+                        path: nested
+                      aggs:
+                        max_number:
+                          max:
+                            field: nested.number
+  - length: { aggregations.name.buckets: 1 }
+  - match: { aggregations.name.buckets.0.key: pig }
+  - match: { aggregations.name.buckets.0.doc_count: 1 }
+  - match: { aggregations.name.buckets.0.children.nested.max_number.value: 100.0 }
+
+---
+"order by sub agg containing join and nested and filter":
+  - skip:
+      reason: "https://github.com/elastic/elasticsearch/issues/66876"
+      version: "7.11.1 - "
+  - do:
+      indices.create:
+          index: test_1
+          body:
+            settings:
+              number_of_shards: 1
+              number_of_replicas: 0
+            mappings:
+              properties:
+                name:
+                  type: keyword
+                join:
+                  type: join
+                  relations:
+                    animal: feature
+                nested:
+                  type: nested
+                  properties:
+                    code:
+                      type:
+                        keyword
+                    number:
+                      type: long
+
+  - do:
+      bulk:
+        index: test_1
+        refresh: true
+        body: |
+          { "index": {"_id": "sheep"} }
+          { "name": "sheep", "join": {"name": "animal"} }
+          { "index": {"_id": "cow"} }
+          { "name": "cow", "join": {"name": "animal"} }
+          { "index": {"_id": "pig"} }
+          { "name": "pig", "join": {"name": "animal"} }
+          { "index": {"routing": "sheep"} }
+          { "join": {"name": "feature", "parent": "sheep"}, "tag": "danger", "nested": [{"code": "mighty", "number": 1}, {"code": "meek", "number": 100}] }
+          { "index": {"routing": "sheep"} }
+          { "join": {"name": "feature", "parent": "sheep"}, "tag": "fluffiness", "nested": [{"code": "mighty", "number": 10}, {"code": "meek", "number": 10}] }
+          { "index": {"routing": "cow"} }
+          { "join": {"name": "feature", "parent": "cow"}, "tag": "danger", "nested": [{"code": "mighty", "number": 3}, {"code": "meek", "number": 3}] }
+          { "index": {"routing": "cow"} }
+          { "join": {"name": "feature", "parent": "cow"}, "tag": "fluffiness", "nested": [{"code": "mighty", "number": 1}, {"code": "meek", "number": 1}] }
+          { "index": {"routing": "pig"} }
+          { "join": {"name": "feature", "parent": "pig"}, "tag": "danger", "nested": [{"code": "mighty", "number": 100}, {"code": "meek", "number": 1}] }
+          { "index": {"routing": "pig"} }
+          { "join": {"name": "feature", "parent": "pig"}, "tag": "fluffiness", "nested": [{"code": "mighty", "number": 1}, {"code": "meek", "number": 1}] }
+
+  - do:
+      search:
+        index: test_1
+        body:
+          size: 0
+          aggs:
+            name:
+              terms:
+                size: 1
+                shard_size: 1
+                field: name
+                order:
+                  - "children>nested>filter>max_number": desc
+              aggs:
+                children:
+                  children:
+                    type: feature
+                  aggs:
+                    nested:
+                      nested:
+                        path: nested
+                      aggs:
+                        filter:
+                          filter:
+                            match:
+                              nested.code: mighty
+                          aggs:
+                            max_number:
+                              max:
+                                field: nested.number
+  - length: { aggregations.name.buckets: 1 }
+  - match: { aggregations.name.buckets.0.key: pig }
+  - match: { aggregations.name.buckets.0.doc_count: 1 }
+  - match: { aggregations.name.buckets.0.children.nested.filter.max_number.value: 100.0 }

rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/170_cardinality_metric.yml

@@ -232,6 +232,7 @@ setup:
   - gt: { profile.shards.0.aggregations.0.breakdown.build_leaf_collector: 0 }
   - gt: { profile.shards.0.aggregations.0.breakdown.collect: 0 }
   - gt: { profile.shards.0.aggregations.0.breakdown.build_aggregation: 0 }
+  - gt: { profile.shards.0.aggregations.0.breakdown.post_collection: 0 }
   - match: { profile.shards.0.aggregations.0.debug.empty_collectors_used: 0 }
   - gt: { profile.shards.0.aggregations.0.debug.numeric_collectors_used: 0 }
   - match: { profile.shards.0.aggregations.0.debug.ordinals_collectors_used: 0 }
@@ -257,6 +258,7 @@ setup:
   - gt: { profile.shards.0.aggregations.0.breakdown.build_leaf_collector: 0 }
   - gt: { profile.shards.0.aggregations.0.breakdown.collect: 0 }
   - gt: { profile.shards.0.aggregations.0.breakdown.build_aggregation: 0 }
+  - gt: { profile.shards.0.aggregations.0.breakdown.post_collection: 0 }
   - match: { profile.shards.0.aggregations.0.debug.empty_collectors_used: 0 }
   - gt: { profile.shards.0.aggregations.0.debug.numeric_collectors_used: 0 }
   - match: { profile.shards.0.aggregations.0.debug.ordinals_collectors_used: 0 }
@@ -283,3 +285,4 @@ setup:
   - gt: { profile.shards.0.aggregations.0.breakdown.build_leaf_collector: 0 }
   - gt: { profile.shards.0.aggregations.0.breakdown.collect: 0 }
   - gt: { profile.shards.0.aggregations.0.breakdown.build_aggregation: 0 }
+  - gt: { profile.shards.0.aggregations.0.breakdown.post_collection: 0 }

rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/20_terms.yml

@@ -1069,3 +1069,154 @@ setup:
   - match: { aggregations.str_terms.buckets.2.key: pig }
   - match: { aggregations.str_terms.buckets.2.doc_count: 0 }
   - match: { aggregations.str_terms.buckets.2.max_number.value: null }
+
+---
+"order by sub agg":
+  - do:
+      bulk:
+        index: test_1
+        refresh: true
+        body: |
+          { "index": {} }
+          { "str": "sheep", "number": 1 }
+          { "index": {} }
+          { "str": "sheep", "number": 3 }
+          { "index": {} }
+          { "str": "cow", "number": 7 }
+          { "index": {} }
+          { "str": "pig", "number": 1 }
+
+  - do:
+      search:
+        index: test_1
+        body:
+          size: 0
+          aggs:
+            str_terms:
+              terms:
+                size: 1
+                shard_size: 1
+                field: str
+                order:
+                  - max_number: desc
+              aggs:
+                max_number:
+                  max:
+                    field: number
+  - length: { aggregations.str_terms.buckets: 1 }
+  - match: { aggregations.str_terms.buckets.0.key: cow }
+  - match: { aggregations.str_terms.buckets.0.doc_count: 1 }
+  - match: { aggregations.str_terms.buckets.0.max_number.value: 7.0 }
+
+---
+"order by sub agg containing nested":
+  - skip:
+        version: "7.99.99 - "
+        reason: "https://github.com/elastic/elasticsearch/issues/66876"
+  - do:
+      indices.put_mapping:
+          index: test_1
+          body:
+            properties:
+              nested:
+                type: nested
+                properties:
+                  tag:
+                    type: keyword
+                  number:
+                    type: long
+
+  - do:
+      bulk:
+        index: test_1
+        refresh: true
+        body: |
+          { "index": {} }
+          { "str": "sheep", "nested": [{"tag": "danger", "number": 1}, {"tag": "fluffiness", "number": 10}] }
+          { "index": {} }
+          { "str": "cow", "nested": [{"tag": "danger", "number": 3}, {"tag": "fluffiness", "number": 1}] }
+          { "index": {} }
+          { "str": "pig", "nested": [{"tag": "danger", "number": 100}, {"tag": "fluffiness", "number": 1}] }
+
+  - do:
+      search:
+        index: test_1
+        body:
+          size: 0
+          aggs:
+            str_terms:
+              terms:
+                size: 1
+                shard_size: 1
+                field: str
+                order:
+                  - "nested>max_number": desc
+              aggs:
+                nested:
+                  nested:
+                    path: nested
+                  aggs:
+                    max_number:
+                      max:
+                        field: nested.number
+  - length: { aggregations.str_terms.buckets: 1 }
+  - match: { aggregations.str_terms.buckets.0.key: pig }
+  - match: { aggregations.str_terms.buckets.0.doc_count: 1 }
+  - match: { aggregations.str_terms.buckets.0.nested.max_number.value: 100.0 }
+
+---
+"order by sub agg containing filter":
+  - do:
+      indices.put_mapping:
+          index: test_1
+          body:
+            properties:
+              nested:
+                type: nested
+                properties:
+                  tag:
+                    type: keyword
+                  number:
+                    type: long
+
+  - do:
+      bulk:
+        index: test_1
+        refresh: true
+        body: |
+          { "index": {} }
+          { "str": "sheep", "number": 1 }
+          { "index": {} }
+          { "str": "sheep", "number": 3 }
+          { "index": {} }
+          { "str": "cow", "number": 7 }
+          { "index": {} }
+          { "str": "pig", "number": 100 }
+
+  - do:
+      search:
+        index: test_1
+        body:
+          size: 0
+          aggs:
+            str_terms:
+              terms:
+                size: 1
+                shard_size: 1
+                field: str
+                order:
+                  - "filter>max_number": desc
+              aggs:
+                filter:
+                  filter:
+                    range:
+                      number:
+                        lt: 10
+                  aggs:
+                    max_number:
+                      max:
+                        field: number
+  - length: { aggregations.str_terms.buckets: 1 }
+  - match: { aggregations.str_terms.buckets.0.key: cow }
+  - match: { aggregations.str_terms.buckets.0.doc_count: 1 }
+  - match: { aggregations.str_terms.buckets.0.filter.max_number.value: 7.0 }

server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java

@@ -586,6 +586,9 @@ public class TransportSearchIT extends ESIntegTestCase {
         @Override
         public void preCollection() throws IOException {}
 
+        @Override
+        public void postCollection() throws IOException {}
+
         @Override
         public Aggregator[] subAggregators() {
             throw new UnsupportedOperationException();

server/src/internalClusterTest/java/org/elasticsearch/search/profile/aggregation/AggregationProfilerIT.java

@@ -47,6 +47,7 @@ import static org.hamcrest.Matchers.notNullValue;
 public class AggregationProfilerIT extends ESIntegTestCase {
     private static final String BUILD_LEAF_COLLECTOR = AggregationTimingType.BUILD_LEAF_COLLECTOR.toString();
     private static final String COLLECT = AggregationTimingType.COLLECT.toString();
+    private static final String POST_COLLECTION = AggregationTimingType.POST_COLLECTION.toString();
     private static final String INITIALIZE = AggregationTimingType.INITIALIZE.toString();
     private static final String BUILD_AGGREGATION = AggregationTimingType.BUILD_AGGREGATION.toString();
     private static final String REDUCE = AggregationTimingType.REDUCE.toString();
@@ -54,11 +55,13 @@ public class AggregationProfilerIT extends ESIntegTestCase {
         INITIALIZE,
         BUILD_LEAF_COLLECTOR,
         COLLECT,
+        POST_COLLECTION,
         BUILD_AGGREGATION,
         REDUCE,
         INITIALIZE + "_count",
         BUILD_LEAF_COLLECTOR + "_count",
         COLLECT + "_count",
+        POST_COLLECTION + "_count",
         BUILD_AGGREGATION + "_count",
         REDUCE + "_count"
     );
@@ -322,6 +325,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(diversifyBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(diversifyBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(diversifyBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(diversifyBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(diversifyBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(diversifyBreakdown.get(REDUCE), equalTo(0L));
             assertThat(diversifyAggResult.getDebugInfo(), equalTo(Map.of(DEFERRED, List.of("max"))));
@@ -338,6 +342,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(diversifyBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(diversifyBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(diversifyBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(diversifyBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(maxBreakdown.get(REDUCE), equalTo(0L));
             assertThat(maxAggResult.getDebugInfo(), equalTo(Map.of()));
@@ -381,6 +386,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(histoBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(histoBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(histoBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(histoBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(histoBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(histoBreakdown.get(REDUCE), equalTo(0L));
             Map<String, Object> histoDebugInfo = histoAggResult.getDebugInfo();
@@ -402,6 +408,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(tagsBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(tagsBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(tagsBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(tagsBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(tagsBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(tagsBreakdown.get(REDUCE), equalTo(0L));
             assertRemapTermsDebugInfo(tagsAggResult);
@@ -420,6 +427,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(avgBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(avgBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(avgBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(avgBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(avgBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(avgBreakdown.get(REDUCE), equalTo(0L));
             assertThat(avgAggResult.getDebugInfo(), equalTo(Map.of()));
@@ -435,6 +443,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(maxBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(maxBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(maxBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(maxBreakdown.get(REDUCE), equalTo(0L));
             assertThat(maxAggResult.getDebugInfo(), equalTo(Map.of()));
@@ -450,6 +459,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(stringsBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(stringsBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(stringsBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(stringsBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(stringsBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(stringsBreakdown.get(REDUCE), equalTo(0L));
             assertRemapTermsDebugInfo(stringsAggResult);
@@ -468,6 +478,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(avgBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(avgBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(avgBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(avgBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(avgBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(avgBreakdown.get(REDUCE), equalTo(0L));
             assertThat(avgAggResult.getDebugInfo(), equalTo(Map.of()));
@@ -483,6 +494,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(maxBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(maxBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(maxBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(maxBreakdown.get(REDUCE), equalTo(0L));
             assertThat(maxAggResult.getDebugInfo(), equalTo(Map.of()));
@@ -499,6 +511,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(tagsBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(tagsBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(tagsBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(tagsBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(tagsBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(tagsBreakdown.get(REDUCE), equalTo(0L));
             assertRemapTermsDebugInfo(tagsAggResult);
@@ -517,6 +530,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(avgBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(avgBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(avgBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(avgBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(avgBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(avgBreakdown.get(REDUCE), equalTo(0L));
             assertThat(avgAggResult.getDebugInfo(), equalTo(Map.of()));
@@ -532,6 +546,7 @@ public class AggregationProfilerIT extends ESIntegTestCase {
             assertThat(maxBreakdown.get(INITIALIZE), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_LEAF_COLLECTOR), greaterThan(0L));
             assertThat(maxBreakdown.get(COLLECT), greaterThan(0L));
+            assertThat(maxBreakdown.get(POST_COLLECTION), greaterThan(0L));
             assertThat(maxBreakdown.get(BUILD_AGGREGATION), greaterThan(0L));
             assertThat(maxBreakdown.get(REDUCE), equalTo(0L));
             assertThat(maxAggResult.getDebugInfo(), equalTo(Map.of()));

server/src/main/java/org/elasticsearch/search/aggregations/AdaptingAggregator.java

@@ -84,6 +84,11 @@ public abstract class AdaptingAggregator extends Aggregator {
         delegate.preCollection();
     }
 
+    @Override
+    public final void postCollection() throws IOException {
+        delegate.postCollection();
+    }
+
     @Override
     public final InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
         InternalAggregation[] delegateResults = delegate.buildAggregations(owningBucketOrds);

server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java

@@ -110,6 +110,7 @@ public class AggregationPhase {
         }
         for (Aggregator aggregator : context.aggregations().aggregators()) {
             try {
+                aggregator.postCollection();
                 aggregations.add(aggregator.buildTopLevel());
             } catch (IOException e) {
                 throw new AggregationExecutionException("Failed to build aggregation [" + aggregator.name() + "]", e);

server/src/main/java/org/elasticsearch/search/aggregations/Aggregator.java

@@ -125,7 +125,7 @@ public abstract class Aggregator extends BucketCollector implements Releasable {
 
     /**
      * Build the results of this aggregation.
-     * @param ordsToCollect the ordinals of the buckets that we want to
+      * @param ordsToCollect the ordinals of the buckets that we want to
      *        collect from this aggregation
      * @return the results for each ordinal, in the same order as the array
      *         of ordinals

server/src/main/java/org/elasticsearch/search/aggregations/AggregatorBase.java

@@ -80,6 +80,10 @@ public abstract class AggregatorBase extends Aggregator {
                 badState();
             }
 
+            @Override
+            public void postCollection() throws IOException {
+                badState();
+            }
             @Override
             public ScoreMode scoreMode() {
                 badState();
@@ -218,6 +222,19 @@ public abstract class AggregatorBase extends Aggregator {
         return subAggregatorbyName.get(aggName);
     }
 
+    /**
+     * Called after collection of all document is done.
+     * <p>
+     * Warning: this is not final only to allow the parent join aggregator
+     * to delay this until building buckets.
+     */
+    @Override
+    public void postCollection() throws IOException {
+        // post-collect this agg before subs to make it possible to buffer and then replay in postCollection()
+        doPostCollection();
+        collectableSubAggregators.postCollection();
+    }
+
     /** Called upon release of the aggregator. */
     @Override
     public void close() {
@@ -231,6 +248,12 @@ public abstract class AggregatorBase extends Aggregator {
     /** Release instance-specific data. */
     protected void doClose() {}
 
+    /**
+     * Can be overridden by aggregator implementation to be called back when the collection phase ends.
+     */
+    protected void doPostCollection() throws IOException {
+    }
+
     protected final InternalAggregations buildEmptySubAggregations() {
         List<InternalAggregation> aggs = new ArrayList<>();
         for (Aggregator aggregator : subAggregators) {

server/src/main/java/org/elasticsearch/search/aggregations/BucketCollector.java

@@ -31,6 +31,10 @@ public abstract class BucketCollector implements Collector {
             // no-op
         }
         @Override
+        public void postCollection() throws IOException {
+            // no-op
+        }
+        @Override
         public ScoreMode scoreMode() {
             return ScoreMode.COMPLETE_NO_SCORES;
         }
@@ -43,4 +47,10 @@ public abstract class BucketCollector implements Collector {
      * Pre collection callback.
      */
     public abstract void preCollection() throws IOException;
+
+    /**
+     * Post-collection callback.
+     */
+    public abstract void postCollection() throws IOException;
+
 }

server/src/main/java/org/elasticsearch/search/aggregations/MultiBucketCollector.java

@@ -115,6 +115,13 @@ public class MultiBucketCollector extends BucketCollector {
         }
     }
 
+    @Override
+    public void postCollection() throws IOException {
+        for (BucketCollector collector : collectors) {
+            collector.postCollection();
+        }
+    }
+
     @Override
     public String toString() {
         return Arrays.toString(collectors);

server/src/main/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollector.java

@@ -62,6 +62,7 @@ public class BestBucketsDeferringCollector extends DeferringBucketCollector {
     private PackedLongValues.Builder docDeltasBuilder;
     private PackedLongValues.Builder bucketsBuilder;
     private LongHash selectedBuckets;
+    private boolean finished = false;
 
     /**
      * Sole constructor.
@@ -134,12 +135,20 @@ public class BestBucketsDeferringCollector extends DeferringBucketCollector {
         collector.preCollection();
     }
 
+    @Override
+    public void postCollection() throws IOException {
+        finishLeaf();
+        finished = true;
+    }
+
     /**
      * Replay the wrapped collector, but only on a selection of buckets.
      */
     @Override
     public void prepareSelectedBuckets(long... selectedBuckets) throws IOException {
-        finishLeaf();
+        if (finished == false) {
+            throw new IllegalStateException("Cannot replay yet, collection is not finished: postCollect() has not been called");
+        }
         if (this.selectedBuckets != null) {
             throw new IllegalStateException("Already been replayed");
         }
@@ -191,6 +200,7 @@ public class BestBucketsDeferringCollector extends DeferringBucketCollector {
                 // continue with the following leaf
             }
         }
+        collector.postCollection();
     }
 
     /**

server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java

@@ -140,7 +140,7 @@ public abstract class BucketsAggregator extends AggregatorBase {
     /**
      * Hook to allow taking an action before building the sub agg results.
      */
-    protected void prepareSubAggs(long[] bucketOrdsToCollect) throws IOException {}
+    protected void prepareSubAggs(long[] ordsToCollect) throws IOException {}
 
     /**
      * Build the results of the sub-aggregations of the buckets at each of

server/src/main/java/org/elasticsearch/search/aggregations/bucket/DeferringBucketCollector.java

@@ -99,6 +99,12 @@ public abstract class DeferringBucketCollector extends BucketCollector {
                     "Deferred collectors cannot be collected directly. They must be collected through the recording wrapper.");
         }
 
+        @Override
+        public void postCollection() throws IOException {
+            throw new IllegalStateException(
+                    "Deferred collectors cannot be collected directly. They must be collected through the recording wrapper.");
+        }
+
         @Override
         public Aggregator resolveSortPath(PathElement next, Iterator<PathElement> path) {
             return in.resolveSortPath(next, path);

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

@@ -123,11 +123,15 @@ final class CompositeAggregator extends BucketsAggregator {
         collectableSubAggregators = BucketCollector.NO_OP_COLLECTOR;
     }
 
+    @Override
+    protected void doPostCollection() throws IOException {
+        finishLeaf();
+    }
+
     @Override
     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;
-        finishLeaf();
         if (deferredCollectors != NO_OP_COLLECTOR) {
             // Replay all documents that contain at least one top bucket (collected during the first pass).
             runDeferredCollections();
@@ -471,6 +475,7 @@ final class CompositeAggregator extends BucketsAggregator {
                 collector.collect(docID);
             }
         }
+        deferredCollectors.postCollection();
     }
 
     /**

server/src/main/java/org/elasticsearch/search/aggregations/bucket/nested/NestedAggregator.java

@@ -101,6 +101,11 @@ public class NestedAggregator extends BucketsAggregator implements SingleBucketA
         processBufferedDocs();
     }
 
+    @Override
+    protected void doPostCollection() throws IOException {
+        processBufferedDocs();
+    }
+
     private void processBufferedDocs() throws IOException {
         if (bufferingNestedLeafBucketCollector != null) {
             bufferingNestedLeafBucketCollector.processBufferedChildBuckets();
@@ -109,7 +114,6 @@ public class NestedAggregator extends BucketsAggregator implements SingleBucketA
 
     @Override
     public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
-        processBufferedDocs();
         return buildAggregationsForSingleBucket(owningBucketOrds, (owningBucketOrd, subAggregationResults) ->
             new InternalNested(name, bucketDocCount(owningBucketOrd), subAggregationResults, metadata()));
     }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/sampler/BestDocsDeferringCollector.java

@@ -113,9 +113,15 @@ public class BestDocsDeferringCollector extends DeferringBucketCollector impleme
         deferred.preCollection();
     }
 
+    @Override
+    public void postCollection() throws IOException {
+        runDeferredAggs();
+    }
+
+
     @Override
     public void prepareSelectedBuckets(long... selectedBuckets) throws IOException {
-        runDeferredAggs();  // TODO should we only prepare the selected buckets?!
+        // no-op - deferred aggs processed in postCollection call
     }
 
     private void runDeferredAggs() throws IOException {
@@ -149,6 +155,7 @@ public class BestDocsDeferringCollector extends DeferringBucketCollector impleme
             // done with allDocs now, reclaim some memory
             circuitBreakerConsumer.accept(-12L * shardSize);
         }
+        deferred.postCollection();
     }
 
     class PerParentBucketSamples {

server/src/main/java/org/elasticsearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -326,12 +326,11 @@ public class GlobalOrdinalsStringTermsAggregator extends AbstractStringTermsAggr
         }
 
         @Override
-        public InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        protected void doPostCollection() throws IOException {
             if (mapping != null) {
                 mapSegmentCountsToGlobalCounts(mapping);
                 mapping = null;
             }
-            return super.buildAggregations(owningBucketOrds);
         }
 
         @Override

server/src/main/java/org/elasticsearch/search/aggregations/metrics/CardinalityAggregator.java

@@ -8,7 +8,9 @@
 
 package org.elasticsearch.search.aggregations.metrics;
 
-import com.carrotsearch.hppc.BitMixer;
+import java.io.IOException;
+import java.util.Map;
+import java.util.function.BiConsumer;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
@@ -27,7 +29,6 @@ import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.ObjectArray;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
-import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.LeafBucketCollector;
@@ -35,9 +36,7 @@ import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
 
-import java.io.IOException;
-import java.util.Map;
-import java.util.function.BiConsumer;
+import com.carrotsearch.hppc.BitMixer;
 
 /**
  * An aggregator that computes approximate counts of unique values.
@@ -136,18 +135,12 @@ public class CardinalityAggregator extends NumericMetricsAggregator.SingleValue
     }
 
     @Override
-    protected void beforeBuildingResults(long[] ordsToCollect) throws IOException {
+    protected void doPostCollection() throws IOException {
         postCollectLastCollector();
     }
 
     @Override
     public double metric(long owningBucketOrd) {
-        try {
-            // Make sure all outstanding data has been synced down to the counts.
-            postCollectLastCollector();
-        } catch (IOException e) {
-            throw new AggregationExecutionException("error collecting data in last segment", e);
-        }
         return counts == null ? 0 : counts.cardinality(owningBucketOrd);
     }
 

server/src/main/java/org/elasticsearch/search/aggregations/metrics/GlobalOrdCardinalityAggregator.java

@@ -8,6 +8,8 @@
 
 package org.elasticsearch.search.aggregations.metrics;
 
+import java.io.IOException;
+import java.util.Map;
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedSetDocValues;
@@ -20,16 +22,12 @@ import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.BitArray;
 import org.elasticsearch.common.util.LongArray;
 import org.elasticsearch.common.util.ObjectArray;
-import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.LeafBucketCollector;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 
-import java.io.IOException;
-import java.util.Map;
-
 /**
  * An aggregator that computes approximate counts of unique values
  * using global ords.
@@ -91,15 +89,7 @@ public class GlobalOrdCardinalityAggregator extends NumericMetricsAggregator.Sin
         };
     }
 
-    @Override
-    protected void beforeBuildingResults(long[] ordsToCollect) throws IOException {
-        buildCountIfNeeded();
-    }
-
-    private void buildCountIfNeeded() throws IOException {
-        if (counts != null) {
-            return;
-        }
+    protected void doPostCollection() throws IOException {
         counts = new HyperLogLogPlusPlusSparse(precision, bigArrays, visitedOrds.size());
         try (LongArray hashes = bigArrays.newLongArray(maxOrd, false)) {
             try (BitArray allVisitedOrds = new BitArray(maxOrd, bigArrays)) {
@@ -138,18 +128,12 @@ public class GlobalOrdCardinalityAggregator extends NumericMetricsAggregator.Sin
 
     @Override
     public double metric(long owningBucketOrd) {
-        try {
-            // Make sure all outstanding data has been synced down to the counts.
-            buildCountIfNeeded();
-        } catch (IOException e) {
-            throw new AggregationExecutionException("error collecting data in last segment", e);
-        }
         return counts.cardinality(owningBucketOrd);
     }
 
     @Override
     public InternalAggregation buildAggregation(long owningBucketOrdinal) {
-        if (owningBucketOrdinal >= counts.maxOrd() || counts.cardinality(owningBucketOrdinal) == 0) {
+        if (counts == null || owningBucketOrdinal >= counts.maxOrd() || counts.cardinality(owningBucketOrdinal) == 0) {
             return buildEmptyAggregation();
         }
         // We need to build a copy because the returned Aggregation needs remain usable after

server/src/main/java/org/elasticsearch/search/aggregations/metrics/MetricsAggregator.java

@@ -28,24 +28,17 @@ public abstract class MetricsAggregator extends AggregatorBase {
          */
     }
 
-    /**
-     * Called once before any calls to {@link #buildAggregation(long)} so the
-     * Aggregator can finish up any work it has to do.
-     */
-    protected void beforeBuildingResults(long[] ordsToCollect) throws IOException {}
-
     /**
      * Build an aggregation for data that has been collected into
      * {@code owningBucketOrd}.
      */
-    public abstract InternalAggregation buildAggregation(long ordToCollect) throws IOException;
+    public abstract InternalAggregation buildAggregation(long owningBucketOrd) throws IOException;
 
     @Override
-    public final InternalAggregation[] buildAggregations(long[] ordsToCollect) throws IOException {
-        beforeBuildingResults(ordsToCollect);
-        InternalAggregation[] results = new InternalAggregation[ordsToCollect.length];
-        for (int ordIdx = 0; ordIdx < ordsToCollect.length; ordIdx++) {
-            results[ordIdx] = buildAggregation(ordsToCollect[ordIdx]);
+    public final InternalAggregation[] buildAggregations(long[] owningBucketOrds) throws IOException {
+        InternalAggregation[] results = new InternalAggregation[owningBucketOrds.length];
+        for (int ordIdx = 0; ordIdx < owningBucketOrds.length; ordIdx++) {
+            results[ordIdx] = buildAggregation(owningBucketOrds[ordIdx]);
         }
         return results;
     }

server/src/main/java/org/elasticsearch/search/profile/aggregation/AggregationTimingType.java

@@ -14,6 +14,7 @@ public enum AggregationTimingType {
     INITIALIZE,
     BUILD_LEAF_COLLECTOR,
     COLLECT,
+    POST_COLLECTION,
     BUILD_AGGREGATION,
     REDUCE;
 

server/src/main/java/org/elasticsearch/search/profile/aggregation/ProfilingAggregator.java

@@ -107,6 +107,17 @@ public class ProfilingAggregator extends Aggregator {
         profiler.pollLastElement();
     }
 
+    @Override
+    public void postCollection() throws IOException {
+        Timer timer = profileBreakdown.getTimer(AggregationTimingType.POST_COLLECTION);
+        timer.start();
+        try {
+            delegate.postCollection();
+        } finally {
+            timer.stop();
+        }
+    }
+
     @Override
     public String toString() {
         return delegate.toString();

server/src/test/java/org/elasticsearch/search/aggregations/MultiBucketCollectorTests.java

@@ -81,6 +81,9 @@ public class MultiBucketCollectorTests  extends ESTestCase {
 
         @Override
         public void preCollection() {}
+
+        @Override
+        public void postCollection() {}
     }
 
     private static class TotalHitCountBucketCollector extends BucketCollector {
@@ -108,6 +111,9 @@ public class MultiBucketCollectorTests  extends ESTestCase {
         @Override
         public void preCollection() {}
 
+        @Override
+        public void postCollection() {}
+
         int getTotalHits() {
             return count;
         }
@@ -141,6 +147,9 @@ public class MultiBucketCollectorTests  extends ESTestCase {
 
         @Override
         public void preCollection() {}
+
+        @Override
+        public void postCollection() {}
     }
 
     public void testCollectionTerminatedExceptionHandling() throws IOException {

server/src/test/java/org/elasticsearch/search/aggregations/bucket/BestBucketsDeferringCollectorTests.java

@@ -73,6 +73,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
         collector.setDeferredCollector(Collections.singleton(bla(deferredCollectedDocIds)));
         collector.preCollection();
         indexSearcher.search(termQuery, collector);
+        collector.postCollection();
         collector.prepareSelectedBuckets(0);
 
         assertEquals(topDocs.scoreDocs.length, deferredCollectedDocIds.size());
@@ -86,6 +87,7 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
         collector.setDeferredCollector(Collections.singleton(bla(deferredCollectedDocIds)));
         collector.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), collector);
+        collector.postCollection();
         collector.prepareSelectedBuckets(0);
 
         assertEquals(topDocs.scoreDocs.length, deferredCollectedDocIds.size());
@@ -113,6 +115,11 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
 
             }
 
+            @Override
+            public void postCollection() throws IOException {
+
+            }
+
             @Override
             public ScoreMode scoreMode() {
                 return ScoreMode.COMPLETE_NO_SCORES;
@@ -201,12 +208,16 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
                     @Override
                     public void preCollection() throws IOException {}
 
+                    @Override
+                    public void postCollection() throws IOException {}
+
                     @Override
                     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx) throws IOException {
                         LeafBucketCollector delegate = deferringCollector.getLeafCollector(ctx);
                         return leafCollector.apply(deferringCollector, delegate);
                     }
                 });
+                deferringCollector.postCollection();
                 verify.accept(deferringCollector, finalCollector);
             }
         }
@@ -232,5 +243,10 @@ public class BestBucketsDeferringCollectorTests extends AggregatorTestCase {
 
         @Override
         public void preCollection() throws IOException {}
+
+        @Override
+        public void postCollection() throws IOException {
+
+        }
     }
 }

server/src/test/java/org/elasticsearch/search/aggregations/bucket/GlobalAggregatorTests.java

@@ -69,6 +69,7 @@ public class GlobalAggregatorTests extends AggregatorTestCase {
         GlobalAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         InternalGlobal result = (InternalGlobal) aggregator.buildTopLevel();
         verify.accept(result, (InternalMin) result.getAggregations().asMap().get("in_global"));
 

server/src/test/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregatorTestCase.java

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

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

@@ -63,6 +63,7 @@ public class BestDocsDeferringCollectorTests extends AggregatorTestCase {
         collector.setDeferredCollector(Collections.singleton(testCollector(deferredCollectedDocIds)));
         collector.preCollection();
         indexSearcher.search(termQuery, collector);
+        collector.postCollection();
         collector.prepareSelectedBuckets(0);
 
         assertEquals(topDocs.scoreDocs.length, deferredCollectedDocIds.size());
@@ -91,6 +92,11 @@ public class BestDocsDeferringCollectorTests extends AggregatorTestCase {
 
             }
 
+            @Override
+            public void postCollection() throws IOException {
+
+            }
+
             @Override
             public ScoreMode scoreMode() {
                 return ScoreMode.COMPLETE_NO_SCORES;

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

@@ -261,6 +261,7 @@ public class RareTermsAggregatorTests extends AggregatorTestCase {
                         Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType1, fieldType2);
                         aggregator.preCollection();
                         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        aggregator.postCollection();
                         RareTerms result = (RareTerms) aggregator.buildTopLevel();
                         assertEquals("_name", result.getName());
                         assertEquals(0, result.getBuckets().size());

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

@@ -76,7 +76,9 @@ import org.elasticsearch.search.aggregations.bucket.nested.InternalNested;
 import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregatorTests;
 import org.elasticsearch.search.aggregations.metrics.CardinalityAggregationBuilder;
+import org.elasticsearch.search.aggregations.metrics.InternalMax;
 import org.elasticsearch.search.aggregations.metrics.InternalTopHits;
+import org.elasticsearch.search.aggregations.metrics.MaxAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.TopHitsAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.BucketScriptPipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
@@ -107,6 +109,7 @@ import static java.util.stream.Collectors.toList;
 import static org.elasticsearch.index.mapper.SeqNoFieldMapper.PRIMARY_TERM_NAME;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
 import static org.elasticsearch.search.aggregations.PipelineAggregatorBuilders.bucketScript;
+import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.instanceOf;
@@ -251,6 +254,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         TermsAggregator aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
                         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        aggregator.postCollection();
                         Terms result = reduce(aggregator, context.bigArrays());
                         assertEquals(5, result.getBuckets().size());
                         assertEquals("", result.getBuckets().get(0).getKeyAsString());
@@ -320,6 +324,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     TermsAggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals(10, result.getBuckets().size());
                     assertEquals("val000", result.getBuckets().get(0).getKeyAsString());
@@ -355,6 +360,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(5, result.getBuckets().size());
                     assertEquals("val001", result.getBuckets().get(0).getKeyAsString());
@@ -379,6 +385,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(8, result.getBuckets().size());
                     assertEquals("val002", result.getBuckets().get(0).getKeyAsString());
@@ -408,6 +415,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals("val010", result.getBuckets().get(0).getKeyAsString());
@@ -425,6 +433,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals("val000", result.getBuckets().get(0).getKeyAsString());
@@ -443,6 +452,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals("val000", result.getBuckets().get(0).getKeyAsString());
@@ -461,6 +471,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals("val000", result.getBuckets().get(0).getKeyAsString());
@@ -479,6 +490,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals("val001", result.getBuckets().get(0).getKeyAsString());
@@ -534,6 +546,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     TermsAggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals(0L, result.getBuckets().get(0).getKey());
@@ -551,6 +564,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(4, result.getBuckets().size());
                     assertEquals(1L, result.getBuckets().get(0).getKey());
@@ -574,6 +588,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(2, result.getBuckets().size());
                     assertEquals(0.0, result.getBuckets().get(0).getKey());
@@ -591,6 +606,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals(4, result.getBuckets().size());
                     assertEquals(1.0, result.getBuckets().get(0).getKey());
@@ -757,6 +773,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals(size, result.getBuckets().size());
                     for (int i = 0; i < size; i++) {
@@ -783,6 +800,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
                         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        aggregator.postCollection();
                         result = ((Filter) reduce(aggregator, context.bigArrays())).getAggregations().get("_name2");
                         int expectedFilteredCounts = 0;
                         for (Integer count : filteredCounts.values()) {
@@ -856,6 +874,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals(size, result.getBuckets().size());
                     for (int i = 0; i < size; i++) {
@@ -885,6 +904,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
                     assertEquals(0, result.getBuckets().size());
@@ -894,6 +914,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
                     assertEquals(0, result.getBuckets().size());
@@ -903,6 +924,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     result = reduce(aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
                     assertEquals(0, result.getBuckets().size());
@@ -926,6 +948,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         Aggregator aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
                         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        aggregator.postCollection();
                         Terms result = reduce(aggregator, context.bigArrays());
                         assertEquals("_name", result.getName());
                         assertEquals(0, result.getBuckets().size());
@@ -962,6 +985,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                         Aggregator aggregator = createAggregator(aggregationBuilder, context);
                         aggregator.preCollection();
                         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                        aggregator.postCollection();
                         Terms result = reduce(aggregator, context.bigArrays());
                         assertEquals("_name", result.getName());
                         assertEquals(1, result.getBuckets().size());
@@ -1034,6 +1058,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals("_name", result.getName());
                     assertEquals(1, result.getBuckets().size());
@@ -1082,6 +1107,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
                     Aggregator aggregator = createAggregator(aggregationBuilder, context);
                     aggregator.preCollection();
                     indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+                    aggregator.postCollection();
                     Terms result = reduce(aggregator, context.bigArrays());
                     assertEquals(3, result.getBuckets().size());
                     assertEquals("a", result.getBuckets().get(0).getKeyAsString());
@@ -1261,6 +1287,74 @@ public class TermsAggregatorTests extends AggregatorTestCase {
         }
     }
 
+    public void testHeisenpig() throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
+                String[] tags = new String[] {"danger", "fluffiness"};
+                indexWriter.addDocuments(generateAnimalDocsWithNested("1", "sheep", tags, new int[] {1, 10}));
+                indexWriter.addDocuments(generateAnimalDocsWithNested("2", "cow", tags, new int[] {3, 1}));
+                indexWriter.addDocuments(generateAnimalDocsWithNested("3", "pig", tags, new int[] {100, 1}));
+                indexWriter.commit();
+                NestedAggregationBuilder nested = new NestedAggregationBuilder("nested", "nested_object")
+                    .subAggregation(
+                        new MaxAggregationBuilder("max_number").field("number")
+                    );
+                TermsAggregationBuilder terms = new TermsAggregationBuilder("str_terms")
+                    .field("str")
+                    .subAggregation(nested)
+                    .shardSize(10)
+                    .size(10)
+                    .order(BucketOrder.aggregation("nested>max_number", false));
+                MappedFieldType nestedFieldType = new NumberFieldMapper.NumberFieldType("number", NumberFieldMapper.NumberType.LONG);
+                MappedFieldType fieldType = new KeywordFieldMapper.KeywordFieldType("str");
+                try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) {
+                    StringTerms result = searchAndReduce(newSearcher(indexReader, false, true),
+                        // match root document only
+                        new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), terms, fieldType, nestedFieldType);
+                    assertThat(result.getBuckets().get(0).getKeyAsString(), equalTo("pig"));
+                    assertThat(result.getBuckets().get(0).docCount, equalTo(1L));
+                    assertThat(((InternalMax) (((InternalNested)result.getBuckets().get(0).getAggregations().get("nested"))
+                        .getAggregations().get("max_number"))).getValue(), closeTo(100.0, 0.00001));
+                }
+            }
+        }
+    }
+
+    public void testSortingWithNestedAggregations() throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
+                for (int i = 0; i < 12; i++) {
+                    int[] nestedValues = new int[i];
+                    for (int j = 0; j < i; j++) {
+                        nestedValues[j] = j;
+                    }
+                    indexWriter.addDocuments(generateDocsWithNested(Integer.toString(i), i % 4, nestedValues));
+                }
+                indexWriter.commit();
+                NestedAggregationBuilder nested = new NestedAggregationBuilder("nested", "nested_object")
+                    .subAggregation(
+                        new MaxAggregationBuilder("max_val").field("nested_value")
+                    );
+                TermsAggregationBuilder terms = new TermsAggregationBuilder("terms")
+                    .field("value")
+                    .subAggregation(nested)
+                    .shardSize(1)
+                    .size(1)
+                    .order(BucketOrder.aggregation("nested>max_val", false));
+                MappedFieldType nestedFieldType = new NumberFieldMapper.NumberFieldType("nested_value", NumberFieldMapper.NumberType.LONG);
+                MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType("value", NumberFieldMapper.NumberType.LONG);
+                try (IndexReader indexReader = wrapInMockESDirectoryReader(DirectoryReader.open(directory))) {
+                    LongTerms result = searchAndReduce(newSearcher(indexReader, false, true),
+                        // match root document only
+                        new DocValuesFieldExistsQuery(PRIMARY_TERM_NAME), terms, fieldType, nestedFieldType);
+                    assertThat(result.getBuckets().get(0).term, equalTo(3L));
+                    assertThat(((InternalMax) (((InternalNested)result.getBuckets().get(0).getAggregations().get("nested"))
+                        .getAggregations().get("max_val"))).getValue(), closeTo(10.0, 0.00001));
+                }
+            }
+        }
+    }
+
     public void testNumberToStringValueScript() throws IOException {
         MappedFieldType fieldType
             = new NumberFieldMapper.NumberFieldType("number", NumberFieldMapper.NumberType.INTEGER);
@@ -1531,6 +1625,28 @@ public class TermsAggregatorTests extends AggregatorTestCase {
         return documents;
     }
 
+    private List<Document> generateAnimalDocsWithNested(String id, String animal, String[] tags, int[] nestedValues) {
+        List<Document> documents = new ArrayList<>();
+
+        for (int i = 0; i < tags.length; i++) {
+            Document document = new Document();
+            document.add(new Field(IdFieldMapper.NAME, Uid.encodeId(id), IdFieldMapper.Defaults.NESTED_FIELD_TYPE));
+
+            document.add(new Field(NestedPathFieldMapper.NAME, "nested_object", NestedPathFieldMapper.Defaults.FIELD_TYPE));
+            document.add(new SortedDocValuesField("tag", new BytesRef(tags[i])));
+            document.add(new SortedNumericDocValuesField("number", nestedValues[i]));
+            documents.add(document);
+        }
+
+        Document document = new Document();
+        document.add(new Field(IdFieldMapper.NAME, Uid.encodeId(id), IdFieldMapper.Defaults.FIELD_TYPE));
+        document.add(new SortedDocValuesField("str", new BytesRef(animal)));
+        document.add(new Field(NestedPathFieldMapper.NAME, "docs", NestedPathFieldMapper.Defaults.FIELD_TYPE));
+        document.add(sequenceIDFields.primaryTerm);
+        documents.add(document);
+
+        return documents;
+    }
 
     private IndexReader createIndexWithLongs() throws IOException {
         Directory directory = newDirectory();
@@ -1575,6 +1691,7 @@ public class TermsAggregatorTests extends AggregatorTestCase {
         TermsAggregator aggregator = createAggregator(builder, searcher, fieldType);
         aggregator.preCollection();
         searcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         return aggregator.buildTopLevel();
     }
 

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

@@ -276,6 +276,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         AvgAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         InternalAvg avg = (InternalAvg) aggregator.buildAggregation(0L);
 
@@ -520,6 +521,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         TermsAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         Terms terms = (Terms) aggregator.buildTopLevel();
         assertNotNull(terms);
@@ -593,6 +595,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         AvgAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         InternalAvg avg = (InternalAvg) aggregator.buildAggregation(0L);
 
@@ -639,6 +642,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         AvgAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         InternalAvg avg = (InternalAvg) aggregator.buildAggregation(0L);
 
@@ -657,6 +661,7 @@ public class AvgAggregatorTests extends AggregatorTestCase {
         aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         avg = (InternalAvg) aggregator.buildAggregation(0L);
 

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

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

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

@@ -450,6 +450,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         GlobalAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         Global global = (Global) aggregator.buildTopLevel();
         assertNotNull(global);
@@ -495,6 +496,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         MaxAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         InternalMax max = (InternalMax) aggregator.buildAggregation(0L);
 
@@ -703,6 +705,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         GlobalAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         Global global = (Global) aggregator.buildTopLevel();
         assertNotNull(global);
@@ -743,6 +746,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         TermsAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         Terms terms = (Terms) aggregator.buildTopLevel();
         assertNotNull(terms);
@@ -796,6 +800,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         BucketCollector bucketCollector = MultiBucketCollector.wrap(maxAggregator, countAggregator);
         bucketCollector.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), bucketCollector);
+        bucketCollector.postCollection();
 
         InternalMax max = (InternalMax) maxAggregator.buildAggregation(0L);
         assertNotNull(max);
@@ -847,6 +852,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
             BucketCollector bucketCollector = MultiBucketCollector.wrap(maxAggregator, countAggregator, termsAggregator);
             bucketCollector.preCollection();
             indexSearcher.search(new MatchAllDocsQuery(), bucketCollector);
+            bucketCollector.postCollection();
 
             InternalMax max = (InternalMax) maxAggregator.buildTopLevel();
             assertNotNull(max);
@@ -904,6 +910,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         MaxAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         InternalMax max = (InternalMax) aggregator.buildAggregation(0L);
 
@@ -950,6 +957,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         MaxAggregator aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         InternalMax max = (InternalMax) aggregator.buildAggregation(0L);
 
@@ -967,6 +975,7 @@ public class MaxAggregatorTests extends AggregatorTestCase {
         aggregator = createAggregator(aggregationBuilder, context);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
 
         max = (InternalMax) aggregator.buildAggregation(0L);
 

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

@@ -184,6 +184,7 @@ public class TDigestPercentilesAggregatorTests extends AggregatorTestCase {
                 TDigestPercentilesAggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
                 indexSearcher.search(query, aggregator);
+                aggregator.postCollection();
                 verify.accept((InternalTDigestPercentiles) aggregator.buildAggregation(0L));
             }
         }

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

@@ -464,6 +464,7 @@ public class WeightedAvgAggregatorTests extends AggregatorTestCase {
             WeightedAvgAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType, fieldType2);
             aggregator.preCollection();
             indexSearcher.search(query, aggregator);
+            aggregator.postCollection();
             verify.accept((InternalWeightedAvg) aggregator.buildAggregation(0L));
         } finally {
             indexReader.close();

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

@@ -466,11 +466,13 @@ public abstract class AggregatorTestCase extends ESTestCase {
                 a.preCollection();
                 Weight weight = subSearcher.createWeight(rewritten, ScoreMode.COMPLETE, 1f);
                 subSearcher.search(weight, a);
+                a.postCollection();
                 aggs.add(a.buildTopLevel());
             }
         } else {
             root.preCollection();
             searcher.search(rewritten, root);
+            root.postCollection();
             aggs.add(root.buildTopLevel());
         }
 

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

@@ -153,6 +153,7 @@ public class HDRPreAggregatedPercentilesAggregatorTests extends AggregatorTestCa
                 Aggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
                 indexSearcher.search(query, aggregator);
+                aggregator.postCollection();
                 verify.accept((InternalHDRPercentiles) aggregator.buildTopLevel());
 
             }

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

@@ -127,6 +127,7 @@ public class TDigestPreAggregatedPercentilesAggregatorTests extends AggregatorTe
                 Aggregator aggregator = createAggregator(builder, indexSearcher, fieldType);
                 aggregator.preCollection();
                 indexSearcher.search(query, aggregator);
+                aggregator.postCollection();
                 verify.accept((InternalTDigestPercentiles) aggregator.buildTopLevel());
 
             }

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

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

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

@@ -1212,6 +1212,7 @@ public class RollupResponseTranslationTests extends AggregatorTestCase {
         try {
             aggregator.preCollection();
             indexSearcher.search(query, aggregator);
+            aggregator.postCollection();
             return aggregator.buildTopLevel();
         } finally {
             indexReader.close();

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

@@ -107,6 +107,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
         directory.close();
@@ -167,6 +168,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
         directory.close();
@@ -219,6 +221,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, valueFieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
         directory.close();
@@ -278,6 +281,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
         directory.close();
@@ -461,6 +465,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, valueFieldType, metricFieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
         directory.close();
@@ -531,6 +536,7 @@ public class IndexerUtilsTests extends AggregatorTestCase {
         Aggregator aggregator = createAggregator(compositeBuilder, indexSearcher, timestampFieldType, valueFieldType);
         aggregator.preCollection();
         indexSearcher.search(new MatchAllDocsQuery(), aggregator);
+        aggregator.postCollection();
         CompositeAggregation composite = (CompositeAggregation) aggregator.buildTopLevel();
         indexReader.close();
         directory.close();

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

@@ -298,6 +298,8 @@ public abstract class GeoShapeGeoGridTestCase<T extends InternalGeoGridBucket<T>
         Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
         aggregator.preCollection();
         indexSearcher.search(query, aggregator);
+        aggregator.postCollection();
+
         verify.accept((InternalGeoGrid<T>) aggregator.buildTopLevel());
 
         indexReader.close();