Added check to not assert early termination in case of invalid field names

Signed-off-by: expani <anijainc@amazon.com>

Author: expani

server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java

@@ -41,7 +41,6 @@
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.NumberFieldMapper;
-import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.TermQueryBuilder;
@@ -285,8 +284,7 @@ public void testStarTreeDocValues() throws IOException {
         );
 
         // Test that feature parity is maintained for unmapped field names.
-        sumAggregationBuilder = new SumAggregationBuilder("sumaggs").field("hello");
-        queryBuilder = new MatchAllQueryBuilder();
+        sumAggregationBuilder = sum("sumaggs").field("hello");
         testCase(
             indexSearcher,
             query,
@@ -295,7 +293,8 @@ public void testStarTreeDocValues() throws IOException {
             starTree,
             supportedDimensions,
             verifyAggregation(InternalSum::getValue),
-            sumAggregationBuilder.build(queryShardContext, null)
+            sumAggregationBuilder.build(queryShardContext, null),
+            false // Invalid fields will return null Star Query Context which will not cause early termination for leaf collector
         );
 
         ir.close();
@@ -319,7 +318,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
         List<Dimension> supportedDimensions,
         BiConsumer<V, V> verify
     ) throws IOException {
-        testCase(searcher, query, queryBuilder, aggBuilder, starTree, supportedDimensions, verify, null);
+        testCase(searcher, query, queryBuilder, aggBuilder, starTree, supportedDimensions, verify, null, true);
     }
 
     private <T extends AggregationBuilder, V extends InternalAggregation> void testCase(
@@ -330,7 +329,8 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
         CompositeIndexFieldInfo starTree,
         List<Dimension> supportedDimensions,
         BiConsumer<V, V> verify,
-        AggregatorFactory aggregatorFactory
+        AggregatorFactory aggregatorFactory,
+        boolean assertCollectorEarlyTermination
     ) throws IOException {
         V starTreeAggregation = searchAndReduceStarTree(
             createIndexSettings(),
@@ -343,6 +343,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
             DEFAULT_MAX_BUCKETS,
             false,
             aggregatorFactory,
+            assertCollectorEarlyTermination,
             DEFAULT_MAPPED_FIELD
         );
         V expectedAggregation = searchAndReduceStarTree(
@@ -356,6 +357,7 @@ private <T extends AggregationBuilder, V extends InternalAggregation> void testC
             DEFAULT_MAX_BUCKETS,
             false,
             aggregatorFactory,
+            assertCollectorEarlyTermination,
             DEFAULT_MAPPED_FIELD
         );
         verify.accept(expectedAggregation, starTreeAggregation);

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

@@ -751,6 +751,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
         int maxBucket,
         boolean hasNested,
         AggregatorFactory aggregatorFactory,
+        boolean assertCollectorEarlyTermination,
         MappedFieldType... fieldTypes
     ) throws IOException {
         query = query.rewrite(searcher);
@@ -782,7 +783,7 @@ protected <A extends InternalAggregation, C extends Aggregator> A searchAndReduc
         searcher.search(query, countingAggregator);
         countingAggregator.postCollection();
         aggs.add(countingAggregator.buildTopLevel());
-        if (compositeIndexFieldInfo != null) {
+        if (compositeIndexFieldInfo != null && assertCollectorEarlyTermination) {
             assertEquals(0, countingAggregator.collectCounter.get());
         }