tidy up post collection.

martijnvg · Aug 10, 2023 · dc80bd8 · dc80bd8
1 parent 4023454
commit dc80bd8
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 13 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java
@@ -40,7 +40,6 @@
 import org.elasticsearch.common.util.concurrent.FutureUtils;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.lucene.util.CombinedBitSet;
-import org.elasticsearch.search.aggregations.BucketCollector;
 import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.profile.Timer;
 import org.elasticsearch.search.profile.query.InternalProfileCollector;
@@ -471,11 +470,8 @@ public void search(List<LeafReaderContext> leaves, Weight weight, Collector coll
 
     private void doAggregationPostCollection(Collector collector) throws IOException {
         if (collector instanceof QueryPhaseCollector queryPhaseCollector) {
-            if (queryPhaseCollector.getAggsCollector() instanceof BucketCollector.BucketCollectorWrapper aggsCollector) {
-                aggsCollector.bucketCollector().postCollection();
-            }
-        }
-        if (collector instanceof InternalProfileCollector profilerCollector) {
+            queryPhaseCollector.doPostCollection();
+        } else if (collector instanceof InternalProfileCollector profilerCollector) {
             profilerCollector.doPostCollection();
         }
     }

diff --git a/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java b/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java
@@ -82,11 +82,7 @@ public void doPostCollection() throws IOException {
         if (wrappedCollector instanceof InternalProfileCollector profileCollector) {
             profileCollector.doPostCollection();
         } else if (wrappedCollector instanceof QueryPhaseCollector queryPhaseCollector) {
-            if (queryPhaseCollector.getAggsCollector() instanceof BucketCollector.BucketCollectorWrapper aggsCollector) {
-                aggsCollector.bucketCollector().postCollection();
-            } else if (queryPhaseCollector.getAggsCollector() instanceof InternalProfileCollector profileCollector) {
-                profileCollector.doPostCollection();
-            }
+            queryPhaseCollector.doPostCollection();
         } else if (wrappedCollector instanceof BucketCollector.BucketCollectorWrapper aggsCollector) {
             aggsCollector.bucketCollector().postCollection();
         }

diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhaseCollector.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhaseCollector.java
@@ -22,6 +22,8 @@
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.search.aggregations.BucketCollector;
+import org.elasticsearch.search.profile.query.InternalProfileCollector;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -458,7 +460,11 @@ boolean isTerminatedAfter() {
         }
     }
 
-    public Collector getAggsCollector() {
-        return aggsCollector;
+    public void doPostCollection() throws IOException {
+        if (aggsCollector instanceof BucketCollector.BucketCollectorWrapper bucketCollectorWrapper) {
+            bucketCollectorWrapper.bucketCollector().postCollection();
+        } else if (aggsCollector instanceof InternalProfileCollector profileCollector) {
+            profileCollector.doPostCollection();
+        }
     }
 }