From c56c693af4eca97b16efceee1ef81516557acac3 Mon Sep 17 00:00:00 2001 From: Jay Deng Date: Tue, 11 Jun 2024 18:36:08 -0700 Subject: [PATCH] Make sure processPostCollection is called in case of early termination Signed-off-by: Jay Deng --- .../AggregationsIntegrationIT.java | 23 +++++++++++++++++++ .../opensearch/search/query/QueryPhase.java | 3 +++ 2 files changed, 26 insertions(+) diff --git a/server/src/internalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java b/server/src/internalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java index 4a8b00ea45738..4051bee3e4e5c 100644 --- a/server/src/internalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java +++ b/server/src/internalClusterTest/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 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 index 608649ad22b23..55b7c0bc5178d 100644 --- 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()) {