Fix profiling of ordered terms aggs (#31814)

This change fixes the profiling of `terms` aggregation ordered
by a sub aggregation.

Closes #22123

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

@@ -28,6 +28,7 @@ import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregation;
 import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregator;
 import org.elasticsearch.search.aggregations.metrics.InternalNumericMetricsAggregation;
 import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregator;
+import org.elasticsearch.search.profile.aggregation.ProfilingAggregator;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -256,7 +257,7 @@ public class AggregationPath {
         Aggregator aggregator = root;
         for (int i = 0; i < pathElements.size(); i++) {
             AggregationPath.PathElement token = pathElements.get(i);
-            aggregator = aggregator.subAggregator(token.name);
+            aggregator = ProfilingAggregator.unwrap(aggregator.subAggregator(token.name));
             assert (aggregator instanceof SingleBucketAggregator && i <= pathElements.size() - 1)
                     || (aggregator instanceof NumericMetricsAggregator && i == pathElements.size() - 1) :
                     "this should be picked up before aggregation execution - on validate";
@@ -272,7 +273,7 @@ public class AggregationPath {
      */
     public Aggregator resolveTopmostAggregator(Aggregator root) {
         AggregationPath.PathElement token = pathElements.get(0);
-        Aggregator aggregator = root.subAggregator(token.name);
+        Aggregator aggregator = ProfilingAggregator.unwrap(root.subAggregator(token.name));
         assert (aggregator instanceof SingleBucketAggregator )
                 || (aggregator instanceof NumericMetricsAggregator) : "this should be picked up before aggregation execution - on validate";
         return aggregator;
@@ -287,7 +288,7 @@ public class AggregationPath {
     public void validate(Aggregator root) throws AggregationExecutionException {
         Aggregator aggregator = root;
         for (int i = 0; i < pathElements.size(); i++) {
-            aggregator = aggregator.subAggregator(pathElements.get(i).name);
+            aggregator = ProfilingAggregator.unwrap(aggregator.subAggregator(pathElements.get(i).name));
             if (aggregator == null) {
                 throw new AggregationExecutionException("Invalid aggregator order path [" + this + "]. Unknown aggregation ["
                         + pathElements.get(i).name + "]");

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

@@ -114,4 +114,11 @@ public class ProfilingAggregator extends Aggregator {
     public String toString() {
         return delegate.toString();
     }
+
+    public static Aggregator unwrap(Aggregator agg) {
+        if (agg instanceof ProfilingAggregator) {
+            return ((ProfilingAggregator) agg).delegate;
+        }
+        return agg;
+    }
 }

server/src/test/java/org/elasticsearch/search/profile/aggregation/AggregationProfilerIT.java

@@ -22,6 +22,7 @@ package org.elasticsearch.search.profile.aggregation;
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.search.aggregations.Aggregator.SubAggCollectionMode;
+import org.elasticsearch.search.aggregations.BucketOrder;
 import org.elasticsearch.search.aggregations.bucket.sampler.DiversifiedOrdinalsSamplerAggregator;
 import org.elasticsearch.search.aggregations.bucket.terms.GlobalOrdinalsStringTermsAggregator;
 import org.elasticsearch.search.aggregations.metrics.avg.AvgAggregator;
@@ -120,9 +121,17 @@ public class AggregationProfilerIT extends ESIntegTestCase {
 
     public void testMultiLevelProfile() {
         SearchResponse response = client().prepareSearch("idx").setProfile(true)
-                .addAggregation(histogram("histo").field(NUMBER_FIELD).interval(1L)
-                        .subAggregation(terms("terms").field(TAG_FIELD)
-                                .subAggregation(avg("avg").field(NUMBER_FIELD)))).get();
+                .addAggregation(
+                    histogram("histo")
+                        .field(NUMBER_FIELD)
+                        .interval(1L)
+                        .subAggregation(
+                            terms("terms")
+                                .field(TAG_FIELD)
+                                .order(BucketOrder.aggregation("avg", false))
+                                .subAggregation(avg("avg").field(NUMBER_FIELD))
+                        )
+                ).get();
         assertSearchResponse(response);
         Map<String, ProfileShardResult> profileResults = response.getProfileResults();
         assertThat(profileResults, notNullValue());