GlobalAggregation with profile option enabled returns incorrect result

Signed-off-by: Sorabh Hamirwasia <sohami.apache@gmail.com>
opensearch-project · Apr 11, 2023 · f4ada0c · f4ada0c
1 parent a15475a
commit f4ada0c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/GlobalIT.java b/server/src/internalClusterTest/java/org/opensearch/search/aggregations/bucket/GlobalIT.java
@@ -87,6 +87,7 @@ public void testWithStatsSubAggregator() throws Exception {
         SearchResponse response = client().prepareSearch("idx")
             .setQuery(QueryBuilders.termQuery("tag", "tag1"))
             .addAggregation(global("global").subAggregation(stats("value_stats").field("value")))
+            .setProfile(randomBoolean())
             .get();
 
         assertSearchResponse(response);

diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregationPhase.java b/server/src/main/java/org/opensearch/search/aggregations/AggregationPhase.java
@@ -38,6 +38,7 @@
 import org.opensearch.common.lucene.search.Queries;
 import org.opensearch.search.aggregations.bucket.global.GlobalAggregator;
 import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.profile.aggregation.ProfilingAggregator;
 import org.opensearch.search.profile.query.CollectorResult;
 import org.opensearch.search.profile.query.InternalProfileCollector;
 import org.opensearch.search.query.QueryPhaseExecutionException;
@@ -67,7 +68,7 @@ public void preProcess(SearchContext context) {
                 AggregatorFactories factories = context.aggregations().factories();
                 aggregators = factories.createTopLevelAggregators(context);
                 for (int i = 0; i < aggregators.length; i++) {
-                    if (aggregators[i] instanceof GlobalAggregator == false) {
+                    if (!isGlobalAggregator(context, aggregators[i])) {
                         collectors.add(aggregators[i]);
                     }
                 }
@@ -106,7 +107,7 @@ public void execute(SearchContext context) {
         Aggregator[] aggregators = context.aggregations().aggregators();
         List<Aggregator> globals = new ArrayList<>();
         for (int i = 0; i < aggregators.length; i++) {
-            if (aggregators[i] instanceof GlobalAggregator) {
+            if (isGlobalAggregator(context, aggregators[i])) {
                 globals.add(aggregators[i]);
             }
         }
@@ -168,4 +169,16 @@ private Collector createCollector(SearchContext context, List<Aggregator> collec
         }
         return collector;
     }
+
+    /**
+     * Checks if passed in aggregator is of type {@link GlobalAggregator}. This method takes care of Aggregator wrapped in
+     * {@link ProfilingAggregator} too
+     * @param context {@link SearchContext}
+     * @param aggregator input {@link Aggregator} instance to evaluate
+     * @return true input is {@link GlobalAggregator} instance or false otherwise
+     */
+    private boolean isGlobalAggregator(SearchContext context, Aggregator aggregator) {
+        return (aggregator instanceof GlobalAggregator
+            || (context.getProfilers() != null && ProfilingAggregator.unwrap(aggregator) instanceof GlobalAggregator));
+    }
 }