Refactor null out from estimation recursive loop

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

Author: bowenlan-amzn

server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java

@@ -307,16 +307,12 @@ public List<Aggregator> createTopLevelNonGlobalAggregators(SearchContext searchC
 
     private List<Aggregator> createTopLevelAggregators(SearchContext searchContext, Predicate<AggregatorFactory> factoryFilter)
         throws IOException {
-        // Estimate streaming cost from factories BEFORE creating any aggregators.
-        // This allows the correct aggregator type to be created on the first try,
-        // avoiding double-creation when streaming is not beneficial.
         if (searchContext.isStreamSearch() && searchContext.getFlushMode() == null) {
-            StreamingCostMetrics metrics = estimateStreamingCostFromFactories(factories, searchContext);
             FlushMode decision;
-            if (metrics == null) {
-                // No factories provided streaming metrics - default to PER_SHARD
+            if (factories.length == 0) {
                 decision = FlushMode.PER_SHARD;
             } else {
+                StreamingCostMetrics metrics = estimateStreamingCostFromFactories(factories, searchContext);
                 long maxBucket = searchContext.getStreamingMaxEstimatedBucketCount();
                 double minRatio = searchContext.getStreamingMinCardinalityRatio();
                 long minBucket = searchContext.getStreamingMinEstimatedBucketCount();
@@ -348,25 +344,19 @@ private List<Aggregator> createTopLevelAggregators(SearchContext searchContext,
     /**
      * Recursively estimates streaming cost from the factory tree.
      *
-     * <p>Traverses the aggregator factory tree, collecting streaming cost metrics from factories
-     * that implement {@link StreamingCostEstimable}. Combines metrics from sibling factories
-     * and nested sub-aggregations to produce a combined estimate.
-     *
-     * @param factories Array of aggregator factories to estimate
+     * @param factories Array of aggregator factories to estimate (must be non-empty)
      * @param searchContext Search context providing access to index metadata
-     * @return Combined streaming cost metrics, null if no factories provide metrics,
-     *         or non-streamable if any factory explicitly returns non-streamable
+     * @return Combined streaming cost metrics, or non-streamable if any factory cannot be streamed
      */
     private static StreamingCostMetrics estimateStreamingCostFromFactories(AggregatorFactory[] factories, SearchContext searchContext) {
+        assert factories.length > 0 : "factories array must be non-empty";
         StreamingCostMetrics combined = null;
         for (AggregatorFactory factory : factories) {
             StreamingCostMetrics metrics = estimateFromFactory(factory, searchContext);
-            if (metrics != null && !metrics.streamable()) {
+            if (!metrics.streamable()) {
                 return StreamingCostMetrics.nonStreamable();
             }
-            if (metrics != null) {
-                combined = (combined == null) ? metrics : combined.combineWithSibling(metrics);
-            }
+            combined = (combined == null) ? metrics : combined.combineWithSibling(metrics);
         }
         return combined;
     }
@@ -380,30 +370,26 @@ private static StreamingCostMetrics estimateStreamingCostFromFactories(Aggregato
      *
      * @param factory The aggregator factory to estimate
      * @param searchContext Search context providing access to index metadata
-     * @return Streaming cost metrics for this factory and its sub-aggregations
+     * @return Streaming cost metrics for this factory and its sub-aggregations (never null)
      */
     private static StreamingCostMetrics estimateFromFactory(AggregatorFactory factory, SearchContext searchContext) {
-        StreamingCostMetrics factoryMetrics;
-        if (factory instanceof StreamingCostEstimable estimable) {
-            factoryMetrics = estimable.estimateStreamingCost(searchContext);
-            if (!factoryMetrics.streamable()) {
-                return StreamingCostMetrics.nonStreamable();
-            }
-        } else {
-            // Factory doesn't implement StreamingCostEstimable - not streaming compatible
+        if (!(factory instanceof StreamingCostEstimable estimable)) {
+            return StreamingCostMetrics.nonStreamable();
+        }
+
+        StreamingCostMetrics factoryMetrics = estimable.estimateStreamingCost(searchContext);
+        if (!factoryMetrics.streamable()) {
             return StreamingCostMetrics.nonStreamable();
         }
 
         // Recursively estimate sub-factories
         AggregatorFactory[] subFactories = factory.getSubFactories().getFactories();
         if (subFactories.length > 0) {
             StreamingCostMetrics subMetrics = estimateStreamingCostFromFactories(subFactories, searchContext);
-            if (subMetrics != null && !subMetrics.streamable()) {
+            if (!subMetrics.streamable()) {
                 return StreamingCostMetrics.nonStreamable();
             }
-            if (subMetrics != null) {
-                factoryMetrics = factoryMetrics.combineWithSubAggregation(subMetrics);
-            }
+            factoryMetrics = factoryMetrics.combineWithSubAggregation(subMetrics);
         }
 
         return factoryMetrics;

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

@@ -679,16 +679,6 @@ static Aggregator createStreamNumericTermsAggregator(
         );
     }
 
-    /**
-     * Estimates streaming cost metrics before aggregator creation.
-     *
-     * <p>This method enables factory-level streaming decision making, allowing
-     * the correct aggregator type to be selected BEFORE creation, avoiding
-     * double-creation overhead.
-     *
-     * @param searchContext The search context providing access to index reader
-     * @return Streaming cost metrics for this terms aggregation
-     */
     @Override
     public StreamingCostMetrics estimateStreamingCost(SearchContext searchContext) {
         ValuesSource valuesSource = config.getValuesSource();
@@ -709,8 +699,7 @@ public StreamingCostMetrics estimateStreamingCost(SearchContext searchContext) {
             return StreamingCostEstimator.estimateStringTerms(searchContext.searcher().getIndexReader(), ordinalsVS, effectiveShardSize);
         }
 
-        // Numeric terms - estimation is less reliable, return non-streamable
-        // to let AggregatorTreeEvaluator handle it with collector-level metrics
+        // Numeric terms - estimation is less reliable
         if (valuesSource instanceof ValuesSource.Numeric) {
             return StreamingCostEstimator.estimateNumericTerms(
                 searchContext.searcher().getIndexReader(),

server/src/main/java/org/opensearch/search/streaming/FlushModeResolver.java

@@ -71,10 +71,6 @@ private FlushModeResolver() {}
     /**
      * Evaluates cost metrics to determine if streaming is beneficial.
      *
-     * <p>This method is used by factories to make streaming decisions before
-     * creating aggregators, eliminating the need for double-creation when streaming
-     * is not beneficial.
-     *
      * @param metrics combined cost metrics from the factory tree
      * @param defaultMode fallback mode when streaming is not beneficial
      * @param maxBucketCount maximum bucket count threshold

server/src/main/java/org/opensearch/search/streaming/StreamingCostEstimable.java

@@ -14,13 +14,8 @@
 /**
  * Interface for aggregator factories that can estimate streaming cost without creating aggregators.
  *
- * <p>This interface enables factory-level streaming cost estimation, allowing the streaming decision
- * to be made BEFORE any aggregators are created. This eliminates the double-creation problem where
- * streaming aggregators are created speculatively, metrics collected, and then recreated as traditional
- * aggregators if streaming is not beneficial.
- *
  * <p>Implementing classes should estimate the streaming cost based on field metadata (ordinals,
- * cardinality) without creating the actual aggregator instance.
+ * cardinality)
  *
  * @opensearch.experimental
  */
@@ -30,12 +25,8 @@ public interface StreamingCostEstimable {
     /**
      * Estimates streaming cost metrics before aggregator creation.
      *
-     * <p>Called only when streaming search is enabled and flushMode has not yet been determined.
-     * The returned metrics represent this factory only (excluding sub-factories, which are
-     * handled separately by the caller).
-     *
      * @param searchContext The search context providing access to index reader and configuration
-     * @return StreamingCostMetrics for this factory, or {@link StreamingCostMetrics#nonStreamable()}
+     * @return StreamingCostMetrics for this factory excluding sub-factories, or {@link StreamingCostMetrics#nonStreamable()}
      *         if this factory cannot support streaming
      */
     StreamingCostMetrics estimateStreamingCost(SearchContext searchContext);

server/src/main/java/org/opensearch/search/streaming/StreamingCostMetrics.java

@@ -70,7 +70,8 @@ public static StreamingCostMetrics neutral() {
      * @return combined metrics reflecting the nested relationship, or non-streamable if either input is non-streamable
      */
     public StreamingCostMetrics combineWithSubAggregation(StreamingCostMetrics subAggMetrics) {
-        if (!this.streamable || subAggMetrics == null || !subAggMetrics.streamable) {
+        assert subAggMetrics != null : "subAggMetrics must not be null";
+        if (!this.streamable || !subAggMetrics.streamable) {
             return nonStreamable();
         }
 
@@ -108,7 +109,8 @@ public StreamingCostMetrics combineWithSubAggregation(StreamingCostMetrics subAg
      * @return combined metrics reflecting the parallel relationship, or non-streamable if either input is non-streamable
      */
     public StreamingCostMetrics combineWithSibling(StreamingCostMetrics siblingMetrics) {
-        if (!this.streamable || siblingMetrics == null || !siblingMetrics.streamable) {
+        assert siblingMetrics != null : "siblingMetrics must not be null";
+        if (!this.streamable || !siblingMetrics.streamable) {
             return nonStreamable();
         }