Make sure processPostCollection is called in case of early termination (

#14208) Signed-off-by: Jay Deng <[email protected]> (cherry picked from commit 5093901) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Jun 12, 2024 · 446f0e4 · 446f0e4
1 parent 5ad0f5d
commit 446f0e4
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 0 deletions.
diff --git a/...nternalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java b/...nternalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java
@@ -187,4 +187,27 @@ public void testAggsOnEmptyShards() {
         // Validate non-global agg does not throw an exception
         assertSearchResponse(client().prepareSearch("idx").addAggregation(stats("value_stats").field("score")).get());
     }
+
+    public void testAggsWithTerminateAfter() throws InterruptedException {
+        assertAcked(
+            prepareCreate(
+                "terminate_index",
+                Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+            ).setMapping("f", "type=keyword").get()
+        );
+        List<IndexRequestBuilder> docs = new ArrayList<>();
+        for (int i = 0; i < randomIntBetween(5, 20); ++i) {
+            docs.add(client().prepareIndex("terminate_index").setSource("f", Integer.toString(i / 3)));
+        }
+        indexRandom(true, docs);
+
+        SearchResponse response = client().prepareSearch("terminate_index")
+            .setSize(2)
+            .setTerminateAfter(1)
+            .addAggregation(terms("f").field("f"))
+            .get();
+        assertSearchResponse(response);
+        assertTrue(response.isTerminatedEarly());
+        assertEquals(response.getHits().getHits().length, 1);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/query/QueryPhase.java b/server/src/main/java/org/opensearch/search/query/QueryPhase.java
@@ -354,6 +354,9 @@ private static boolean searchWithCollector(
         try {
             searcher.search(query, queryCollector);
         } catch (EarlyTerminatingCollector.EarlyTerminationException e) {
+            // EarlyTerminationException is not caught in ContextIndexSearcher to allow force termination of collection. Postcollection
+            // still needs to be processed for Aggregations when early termination takes place.
+            searchContext.bucketCollectorProcessor().processPostCollection(queryCollector);
             queryResult.terminatedEarly(true);
         }
         if (searchContext.isSearchTimedOut()) {