Fix the regression of terms agg optimization

Signed-off-by: bowenlan-amzn <bowenlan23@gmail.com>

Author: bowenlan-amzn

server/src/main/java/org/opensearch/common/settings/ClusterSettings.java

@@ -594,6 +594,7 @@ public void apply(Settings value, Settings current, Settings previous) {
                 SearchService.SEARCH_MAX_QUERY_STRING_LENGTH,
                 SearchService.SEARCH_MAX_QUERY_STRING_LENGTH_MONITOR_ONLY,
                 SearchService.CARDINALITY_AGGREGATION_PRUNING_THRESHOLD,
+                SearchService.TERMS_AGGREGATION_MAX_PRECOMPUTE_CARDINALITY,
                 CardinalityAggregator.CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_ENABLED,
                 CardinalityAggregator.CARDINALITY_AGGREGATION_HYBRID_COLLECTOR_MEMORY_THRESHOLD,
                 SearchService.KEYWORD_INDEX_OR_DOC_VALUES_ENABLED,

server/src/main/java/org/opensearch/search/DefaultSearchContext.java

@@ -230,6 +230,7 @@ final class DefaultSearchContext extends SearchContext {
     private final int maxAggRewriteFilters;
     private final int filterRewriteSegmentThreshold;
     private final int cardinalityAggregationPruningThreshold;
+    private final long termsAggregationMaxPrecomputeCardinality;
     private final CardinalityAggregationContext cardinalityAggregationContext;
     private final int bucketSelectionStrategyFactor;
     private final boolean keywordIndexOrDocValuesEnabled;
@@ -298,6 +299,7 @@ final class DefaultSearchContext extends SearchContext {
         this.maxAggRewriteFilters = evaluateFilterRewriteSetting();
         this.filterRewriteSegmentThreshold = evaluateAggRewriteFilterSegThreshold();
         this.cardinalityAggregationPruningThreshold = evaluateCardinalityAggregationPruningThreshold();
+        this.termsAggregationMaxPrecomputeCardinality = evaluateTermsAggregationMaxPrecomputeCardinality();
         this.cardinalityAggregationContext = evaluateCardinalityAggregationContext();
         this.bucketSelectionStrategyFactor = evaluateBucketSelectionStrategyFactor();
         this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
@@ -1260,6 +1262,11 @@ public int cardinalityAggregationPruningThreshold() {
         return cardinalityAggregationPruningThreshold;
     }
 
+    @Override
+    public long termsAggregationMaxPrecomputeCardinality() {
+        return termsAggregationMaxPrecomputeCardinality;
+    }
+
     @Override
     public CardinalityAggregationContext cardinalityAggregationContext() {
         return cardinalityAggregationContext;
@@ -1282,6 +1289,13 @@ private int evaluateCardinalityAggregationPruningThreshold() {
         return 0;
     }
 
+    private long evaluateTermsAggregationMaxPrecomputeCardinality() {
+        if (clusterService != null) {
+            return clusterService.getClusterSettings().get(SearchService.TERMS_AGGREGATION_MAX_PRECOMPUTE_CARDINALITY);
+        }
+        return 30_000L;
+    }
+
     private CardinalityAggregationContext evaluateCardinalityAggregationContext() {
         if (clusterService != null) {
             boolean hybridCollectorEnabled = clusterService.getClusterSettings()

server/src/main/java/org/opensearch/search/SearchService.java

@@ -442,6 +442,15 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         Property.NodeScope
     );
 
+    @ExperimentalApi
+    public static final Setting<Long> TERMS_AGGREGATION_MAX_PRECOMPUTE_CARDINALITY = Setting.longSetting(
+        "search.aggregations.terms.max_precompute_cardinality",
+        30_000L,
+        0L,
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     public static final int DEFAULT_BUCKET_SELECTION_STRATEGY_FACTOR = 5;
     public static final Setting<Integer> BUCKET_SELECTION_STRATEGY_FACTOR_SETTING = Setting.intSetting(
         "search.aggregation.bucket_selection_strategy_factor",

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java

@@ -196,6 +196,11 @@ boolean tryCollectFromTermFrequencies(LeafReaderContext ctx, BiConsumer<Long, In
             return false;
         }
 
+        long termCount = segmentTerms.size();
+        if (termCount == -1 || termCount > context.termsAggregationMaxPrecomputeCardinality()) {
+            return false;
+        }
+
         NumericDocValues docCountValues = DocValues.getNumeric(ctx.reader(), DocCountFieldMapper.NAME);
         if (docCountValues.nextDoc() != NO_MORE_DOCS) {
             // This segment has at least one document with the _doc_count field.

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/TermsAggregatorFactory.java

@@ -776,8 +776,9 @@ private boolean shouldDisableStreamingForOrdinals(SearchContext searchContext, W
         }
 
         // Check 2: Match-all query with the majority of docs in clean segments
-        // Traditional aggregator can use term frequency optimization for these segments
-        if (isMatchAllQuery(searchContext.query())) {
+        // and cardinality within the precompute threshold.
+        // Traditional aggregator can use term frequency optimization for these segments.
+        if (isMatchAllQuery(searchContext.query()) && maxCardinality <= searchContext.termsAggregationMaxPrecomputeCardinality()) {
             double cleanRatio = totalDocs > 0 ? (double) docsInCleanSegments / totalDocs : 0;
             return cleanRatio > 0.8;
         }

server/src/main/java/org/opensearch/search/internal/SearchContext.java

@@ -560,6 +560,10 @@ public int cardinalityAggregationPruningThreshold() {
         return 0;
     }
 
+    public long termsAggregationMaxPrecomputeCardinality() {
+        return 30_000L;
+    }
+
     public CardinalityAggregationContext cardinalityAggregationContext() {
         return new CardinalityAggregationContext(false, Runtime.getRuntime().maxMemory() / 100);
     }

server/src/test/java/org/opensearch/search/aggregations/FactoryStreamingCostEstimationTests.java

@@ -667,6 +667,37 @@ public void testTermsFactoryStreamableNonMatchAllQuery() throws IOException {
         }
     }
 
+    /**
+     * Test that match-all query with clean segments but cardinality above the precompute threshold IS streamable.
+     * When cardinality exceeds the threshold, Check 2 is skipped because tryCollectFromTermFrequencies
+     * would bail out on high cardinality anyway.
+     */
+    public void testTermsFactoryStreamableMatchAllHighCardinality() throws IOException {
+        try (Directory directory = newDirectory()) {
+            try (IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) {
+                // Create index with cardinality above the default threshold (30,000)
+                for (int i = 0; i < 50000; i++) {
+                    Document doc = new Document();
+                    doc.add(new SortedSetDocValuesField("category", new BytesRef("cat_" + i)));
+                    writer.addDocument(doc);
+                }
+
+                try (IndexReader reader = DirectoryReader.open(writer)) {
+                    IndexSearcher searcher = newIndexSearcher(reader);
+                    MappedFieldType fieldType = new KeywordFieldMapper.KeywordFieldType("category");
+
+                    TermsAggregationBuilder termsBuilder = new TermsAggregationBuilder("terms").field("category").size(10);
+
+                    // Match-all with clean segments but high cardinality - should be streamable
+                    FactoryAndContext result = createAggregatorFactoryWithQuery(termsBuilder, searcher, new MatchAllDocsQuery(), fieldType);
+                    StreamingCostMetrics metrics = ((StreamingCostEstimable) result.factory).estimateStreamingCost(result.searchContext);
+
+                    assertTrue("Match-all with high cardinality should be streamable", metrics.streamable());
+                }
+            }
+        }
+    }
+
     // ========================================
     // Helper methods
     // ========================================

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

@@ -550,6 +550,7 @@ public boolean shouldCache(Query query) {
         fieldNameToType.putAll(getFieldAliases(fieldTypes));
 
         when(searchContext.maxAggRewriteFilters()).thenReturn(10_000);
+        when(searchContext.termsAggregationMaxPrecomputeCardinality()).thenReturn(30_000L);
         when(searchContext.cardinalityAggregationContext()).thenReturn(
             new org.opensearch.search.aggregations.metrics.CardinalityAggregationContext(false, Runtime.getRuntime().maxMemory() / 100)
         );