Skip to content

Commit 1b91933

Browse files
jed326Jay Deng
authored and
Jay Deng
committed
Make sure processPostCollection is called in case of early termination
Signed-off-by: Jay Deng <jayd0104@gmail.com>
1 parent 2f8cb07 commit 1b91933

File tree

2 files changed

+27
-0
lines changed

2 files changed

+27
-0
lines changed

server/src/internalClusterTest/java/org/opensearch/search/aggregations/AggregationsIntegrationIT.java

+24
Original file line numberDiff line numberDiff line change
@@ -187,4 +187,28 @@ public void testAggsOnEmptyShards() {
187187
// Validate non-global agg does not throw an exception
188188
assertSearchResponse(client().prepareSearch("idx").addAggregation(stats("value_stats").field("score")).get());
189189
}
190+
191+
public void testAggsWithTerminateAfter() throws InterruptedException {
192+
assertAcked(
193+
prepareCreate(
194+
"terminate_index",
195+
Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
196+
).setMapping("f", "type=keyword").get()
197+
);
198+
List<IndexRequestBuilder> docs = new ArrayList<>();
199+
for (int i = 0; i < randomIntBetween(5, 20); ++i) {
200+
docs.add(client().prepareIndex("terminate_index").setSource("f", Integer.toString(i / 3)));
201+
}
202+
indexRandom(true, docs);
203+
204+
SearchResponse response = client().prepareSearch("terminate_index")
205+
.setSize(2)
206+
.setTerminateAfter(1)
207+
.addAggregation(terms("f").field("f"))
208+
.get();
209+
assertSearchResponse(response);
210+
System.out.println(response);
211+
assertTrue(response.isTerminatedEarly());
212+
assertEquals(response.getHits().getHits().length, 1);
213+
}
190214
}

server/src/main/java/org/opensearch/search/query/QueryPhase.java

+3
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,9 @@ private static boolean searchWithCollector(
354354
try {
355355
searcher.search(query, queryCollector);
356356
} catch (EarlyTerminatingCollector.EarlyTerminationException e) {
357+
// EarlyTerminationException is not caught in ContextIndexSearcher to allow force termination of collection. Postcollection
358+
// still needs to be processed for Aggregations when early termination takes place.
359+
searchContext.bucketCollectorProcessor().processPostCollection(queryCollector);
357360
queryResult.terminatedEarly(true);
358361
}
359362
if (searchContext.isSearchTimedOut()) {

0 commit comments

Comments
 (0)