[Bug Fix] [Star Tree] Fix sub-aggregator casting for search with profile=true (#19652)

* fix sub-aggregator casting when used with profile=true

Signed-off-by: Sandesh Kumar <[email protected]>

---------

Signed-off-by: Sandesh Kumar <[email protected]>

Author: sandeshkr419

CHANGELOG.md

@@ -25,6 +25,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix flaky test FieldDataLoadingIT.testIndicesFieldDataCacheSizeSetting ([#19571](https://github.com/opensearch-project/OpenSearch/pull/19571))
 - Avoid primary shard failure caused by merged segment warmer exceptions ([#19436](https://github.com/opensearch-project/OpenSearch/pull/19436))
 - Fix pull-based ingestion out-of-bounds offset scenarios and remove persisted offsets ([#19607](https://github.com/opensearch-project/OpenSearch/pull/19607))
+- [Star Tree] Fix sub-aggregator casting for search with profile=true ([19652](https://github.com/opensearch-project/OpenSearch/pull/19652))
 
 ### Dependencies
 - Update to Gradle 9.1 ([#19575](https://github.com/opensearch-project/OpenSearch/pull/19575))

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

@@ -170,6 +170,15 @@ public BucketComparator bucketComparator(String key, SortOrder order) {
         );
     }
 
+    /**
+     * Returns the underlying Aggregator responsible for creating the bucket collector.
+     * For most aggregators, this is the aggregator itself.
+     * For wrappers like ProfilingAggregator, it's the delegate.
+     */
+    public Aggregator unwrapAggregator() {
+        return this;
+    }
+
     /**
      * Compare two buckets by their ordinal.
      *

server/src/main/java/org/opensearch/search/aggregations/bucket/histogram/DateHistogramAggregator.java

@@ -335,7 +335,9 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             @Override
             public void setSubCollectors() throws IOException {
                 for (Aggregator aggregator : subAggregators) {
-                    this.subCollectors.add(((StarTreePreComputeCollector) aggregator).getStarTreeBucketCollector(ctx, starTree, this));
+                    this.subCollectors.add(
+                        ((StarTreePreComputeCollector) aggregator.unwrapAggregator()).getStarTreeBucketCollector(ctx, starTree, this)
+                    );
                 }
             }
 

server/src/main/java/org/opensearch/search/aggregations/bucket/range/RangeAggregator.java

@@ -401,7 +401,9 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             @Override
             public void setSubCollectors() throws IOException {
                 for (Aggregator aggregator : subAggregators) {
-                    this.subCollectors.add(((StarTreePreComputeCollector) aggregator).getStarTreeBucketCollector(ctx, starTree, this));
+                    this.subCollectors.add(
+                        ((StarTreePreComputeCollector) aggregator.unwrapAggregator()).getStarTreeBucketCollector(ctx, starTree, this)
+                    );
                 }
             }
 

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

@@ -361,7 +361,9 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             @Override
             public void setSubCollectors() throws IOException {
                 for (Aggregator aggregator : subAggregators) {
-                    this.subCollectors.add(((StarTreePreComputeCollector) aggregator).getStarTreeBucketCollector(ctx, starTree, this));
+                    this.subCollectors.add(
+                        ((StarTreePreComputeCollector) aggregator.unwrapAggregator()).getStarTreeBucketCollector(ctx, starTree, this)
+                    );
                 }
             }
 

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

@@ -311,9 +311,9 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             @Override
             public void setSubCollectors() throws IOException {
                 for (Aggregator aggregator : subAggregators) {
-                    if (aggregator instanceof StarTreePreComputeCollector collector) {
-                        this.subCollectors.add(collector.getStarTreeBucketCollector(ctx, starTree, this));
-                    }
+                    this.subCollectors.add(
+                        ((StarTreePreComputeCollector) aggregator.unwrapAggregator()).getStarTreeBucketCollector(ctx, starTree, this)
+                    );
                 }
             }
 

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

@@ -191,7 +191,9 @@ public StarTreeBucketCollector getStarTreeBucketCollector(
             @Override
             public void setSubCollectors() throws IOException {
                 for (Aggregator aggregator : subAggregators) {
-                    this.subCollectors.add(((StarTreePreComputeCollector) aggregator).getStarTreeBucketCollector(ctx, starTree, this));
+                    this.subCollectors.add(
+                        ((StarTreePreComputeCollector) aggregator.unwrapAggregator()).getStarTreeBucketCollector(ctx, starTree, this)
+                    );
                 }
             }
 

server/src/main/java/org/opensearch/search/aggregations/support/AggregationPath.java

@@ -40,7 +40,6 @@
 import org.opensearch.search.aggregations.InternalAggregations;
 import org.opensearch.search.aggregations.bucket.SingleBucketAggregator;
 import org.opensearch.search.aggregations.metrics.NumericMetricsAggregator;
-import org.opensearch.search.profile.aggregation.ProfilingAggregator;
 import org.opensearch.search.sort.SortOrder;
 
 import java.util.ArrayList;
@@ -236,7 +235,7 @@ public Aggregator resolveAggregator(Aggregator root) {
     public Aggregator resolveTopmostAggregator(Aggregator root) {
         AggregationPath.PathElement token = pathElements.get(0);
         // TODO both unwrap and subAggregator are only used here!
-        Aggregator aggregator = ProfilingAggregator.unwrap(root.subAggregator(token.name));
+        Aggregator aggregator = root.subAggregator(token.name).unwrapAggregator();
         assert (aggregator instanceof SingleBucketAggregator) || (aggregator instanceof NumericMetricsAggregator)
             : "this should be picked up before aggregation execution - on validate";
         return aggregator;

server/src/main/java/org/opensearch/search/profile/aggregation/ProfilingAggregator.java

@@ -56,7 +56,7 @@ public class ProfilingAggregator extends Aggregator implements Streamable {
     private final AggregationProfiler profiler;
     private AggregationProfileBreakdown profileBreakdown;
 
-    public ProfilingAggregator(Aggregator delegate, AggregationProfiler profiler) throws IOException {
+    public ProfilingAggregator(Aggregator delegate, AggregationProfiler profiler) {
         this.profiler = profiler;
         this.delegate = delegate;
     }
@@ -164,19 +164,13 @@ public String toString() {
         return delegate.toString();
     }
 
-    public static Aggregator unwrap(Aggregator agg) {
-        if (agg instanceof ProfilingAggregator) {
-            return ((ProfilingAggregator) agg).delegate;
-        }
-        return agg;
-    }
-
-    public Aggregator getDelegate() {
-        return delegate;
-    }
-
     @Override
     public StreamingCostMetrics getStreamingCostMetrics() {
         return delegate instanceof Streamable ? ((Streamable) delegate).getStreamingCostMetrics() : StreamingCostMetrics.nonStreamable();
     }
+
+    @Override
+    public Aggregator unwrapAggregator() {
+        return delegate;
+    }
 }

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

@@ -139,19 +139,13 @@ private static StreamingCostMetrics collectMetrics(Collector collector) {
     }
 
     private static Collector[] getChildren(Collector collector) {
-        if (collector instanceof AggregatorBase) {
-            return ((AggregatorBase) collector).subAggregators();
-        }
-        if (collector instanceof MultiCollector) {
-            return ((MultiCollector) collector).getCollectors();
-        }
-        if (collector instanceof MultiBucketCollector) {
-            return ((MultiBucketCollector) collector).getCollectors();
-        }
-        if (collector instanceof ProfilingAggregator) {
-            return getChildren(((ProfilingAggregator) collector).getDelegate());
-        }
-        return new Collector[0];
+        return switch (collector) {
+            case AggregatorBase aggregatorBase -> aggregatorBase.subAggregators();
+            case MultiCollector multiCollector -> multiCollector.getCollectors();
+            case MultiBucketCollector multiBucketCollector -> multiBucketCollector.getCollectors();
+            case ProfilingAggregator profilingAggregator -> getChildren(profilingAggregator.unwrapAggregator());
+            default -> new Collector[0];
+        };
     }
 
     /**