Fix profiled global agg

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.aggregation/380_global.yml

@@ -0,0 +1,64 @@
+setup:
+  - do:
+     bulk:
+        refresh: true
+        index: test
+        body:
+          - '{"index": {}}'
+          - '{"name": "one"}'
+          - '{"index": {}}'
+          - '{"name": "two"}'
+          - '{"index": {}}'
+          - '{"name": "two"}'
+
+---
+simple:
+  - do:
+      search:
+        index: test
+        body:
+          size: 0
+          query:
+            match:
+              name: two
+          aggs:
+            g:
+              global: {}
+              aggs:
+                t:
+                  terms:
+                    field: name.keyword
+
+  - match: { aggregations.g.doc_count: 3 }
+  - length: { aggregations.g.t.buckets: 2 }
+  - match: { aggregations.g.t.buckets.0.key: two }
+  - match: { aggregations.g.t.buckets.1.key: one }
+
+---
+profile:
+  - skip:
+      version: " - 7.99.99"
+      reason:  fixed in 8.0.0 (to be backported to 7.13.0)
+
+  - do:
+      search:
+        index: test
+        body:
+          profile: true
+          size: 0
+          query:
+            match:
+              name: two
+          aggs:
+            g:
+              global: {}
+              aggs:
+                t:
+                  terms:
+                    field: name.keyword
+
+  - match: { aggregations.g.doc_count: 3 }
+  - length: { aggregations.g.t.buckets: 2 }
+  - match: { aggregations.g.t.buckets.0.key: two }
+  - match: { aggregations.g.t.buckets.1.key: one }
+  - match: { profile.shards.0.aggregations.0.description: g }

server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java

@@ -8,18 +8,13 @@
 package org.elasticsearch.search.aggregations;
 
 import org.apache.lucene.search.Collector;
-import org.apache.lucene.search.Query;
 import org.elasticsearch.common.inject.Inject;
-import org.elasticsearch.common.lucene.search.Queries;
-import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregator;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.profile.query.CollectorResult;
 import org.elasticsearch.search.profile.query.InternalProfileCollector;
-import org.elasticsearch.search.query.QueryPhaseExecutionException;
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 
 /**
@@ -32,31 +27,20 @@ public AggregationPhase() {
     }
 
     public void preProcess(SearchContext context) {
-        if (context.aggregations() != null) {
-            List<Aggregator> collectors = new ArrayList<>();
-            Aggregator[] aggregators;
-            try {
-                aggregators = context.aggregations().factories().createTopLevelAggregators();
-                for (int i = 0; i < aggregators.length; i++) {
-                    if (aggregators[i] instanceof GlobalAggregator == false) {
-                        collectors.add(aggregators[i]);
-                    }
-                }
-                context.aggregations().aggregators(aggregators);
-                if (collectors.isEmpty() == false) {
-                    Collector collector = MultiBucketCollector.wrap(true, collectors);
-                    ((BucketCollector)collector).preCollection();
-                    if (context.getProfilers() != null) {
-                        collector = new InternalProfileCollector(collector, CollectorResult.REASON_AGGREGATION,
-                                // TODO: report on child aggs as well
-                                Collections.emptyList());
-                    }
-                    context.queryCollectors().put(AggregationPhase.class, collector);
-                }
-            } catch (IOException e) {
-                throw new AggregationInitializationException("Could not initialize aggregators", e);
-            }
+        if (context.aggregations() == null) {
+            return;
+        }
+        BucketCollector bucketCollector;
+        try {
+            context.aggregations().aggregators(context.aggregations().factories().createTopLevelAggregators());
+            bucketCollector = MultiBucketCollector.wrap(true, List.of(context.aggregations().aggregators()));
+            bucketCollector.preCollection();
+        } catch (IOException e) {
+            throw new AggregationInitializationException("Could not initialize aggregators", e);
         }
+        Collector collector = context.getProfilers() == null ?
+            bucketCollector : new InternalProfileCollector(bucketCollector, CollectorResult.REASON_AGGREGATION, List.of());
+        context.queryCollectors().put(AggregationPhase.class, collector);
     }
 
     public void execute(SearchContext context) {
@@ -71,37 +55,6 @@ 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) {
-                globals.add(aggregators[i]);
-            }
-        }
-
-        // optimize the global collector based execution
-        if (globals.isEmpty() == false) {
-            BucketCollector globalsCollector = MultiBucketCollector.wrap(false, globals);
-            Query query = context.buildFilteredQuery(Queries.newMatchAllQuery());
-
-            try {
-                final Collector collector;
-                if (context.getProfilers() == null) {
-                    collector = globalsCollector;
-                } else {
-                    InternalProfileCollector profileCollector = new InternalProfileCollector(
-                            globalsCollector, CollectorResult.REASON_AGGREGATION_GLOBAL,
-                            // TODO: report on sub collectors
-                            Collections.emptyList());
-                    collector = profileCollector;
-                    // start a new profile with this collector
-                    context.getProfilers().addQueryProfiler().setCollector(profileCollector);
-                }
-                globalsCollector.preCollection();
-                context.searcher().search(query, collector);
-            } catch (Exception e) {
-                throw new QueryPhaseExecutionException(context.shardTarget(), "Failed to execute global aggregators", e);
-            }
-        }
 
         List<InternalAggregation> aggregations = new ArrayList<>(aggregators.length);
         if (context.aggregations().factories().context() != null) {

server/src/main/java/org/elasticsearch/search/aggregations/bucket/global/GlobalAggregator.java

@@ -8,11 +8,14 @@
 package org.elasticsearch.search.aggregations.bucket.global;
 
 import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.LeafCollector;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.Scorable;
+import org.apache.lucene.search.Weight;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.CardinalityUpperBound;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.LeafBucketCollector;
-import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
 import org.elasticsearch.search.aggregations.bucket.BucketsAggregator;
 import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregator;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
@@ -21,22 +24,30 @@
 import java.util.Map;
 
 public class GlobalAggregator extends BucketsAggregator implements SingleBucketAggregator {
+    private final Weight weight;
 
     public GlobalAggregator(String name, AggregatorFactories subFactories, AggregationContext context, Map<String, Object> metadata)
         throws IOException {
+
         super(name, subFactories, context, null, CardinalityUpperBound.ONE, metadata);
+        weight = new MatchAllDocsQuery().createWeight(context.searcher(), scoreMode(), 1.0f);
     }
 
     @Override
-    public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
-            final LeafBucketCollector sub) throws IOException {
-        return new LeafBucketCollectorBase(sub, null) {
+    public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, LeafBucketCollector sub) throws IOException {
+        // Run sub-aggregations on child documents
+        weight.bulkScorer(ctx).score(new LeafCollector() {
+            @Override
+            public void collect(int doc) throws IOException {
+                collectBucket(sub, doc, 0);
+            }
+
             @Override
-            public void collect(int doc, long bucket) throws IOException {
-                assert bucket == 0 : "global aggregator can only be a top level aggregator";
-                collectBucket(sub, doc, bucket);
+            public void setScorer(Scorable scorer) throws IOException {
+                sub.setScorer(scorer);
             }
-        };
+        }, ctx.reader().getLiveDocs());
+        return LeafBucketCollector.NO_OP_COLLECTOR;
     }
 
     @Override

server/src/test/java/org/elasticsearch/search/aggregations/bucket/DocCountProviderTests.java

@@ -18,9 +18,11 @@
 import org.elasticsearch.index.mapper.DocCountFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
+import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.elasticsearch.search.aggregations.AggregatorTestCase;
-import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
-import org.elasticsearch.search.aggregations.bucket.global.InternalGlobal;
+import org.elasticsearch.search.aggregations.bucket.filter.FilterAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.filter.InternalFilter;
 
 import java.io.IOException;
 import java.util.List;
@@ -88,10 +90,10 @@ public void testQueryFiltering() throws IOException {
 
     private void testAggregation(Query query,
                                  CheckedConsumer<RandomIndexWriter, IOException> indexer,
-                                 Consumer<InternalGlobal> verify) throws IOException {
-        GlobalAggregationBuilder aggregationBuilder = new GlobalAggregationBuilder("_name");
+                                 Consumer<InternalFilter> verify) throws IOException {
+        AggregationBuilder builder = new FilterAggregationBuilder("f", new MatchAllQueryBuilder());
         MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType(NUMBER_FIELD, NumberFieldMapper.NumberType.LONG);
         MappedFieldType docCountFieldType = new DocCountFieldMapper.DocCountFieldType();
-        testCase(aggregationBuilder, query, indexer, verify, fieldType, docCountFieldType);
+        testCase(builder, query, indexer, verify, fieldType, docCountFieldType);
     }
 }

server/src/test/java/org/elasticsearch/search/aggregations/bucket/GlobalAggregatorTests.java

@@ -8,72 +8,66 @@
 
 package org.elasticsearch.search.aggregations.bucket;
 
+import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.document.SortedNumericDocValuesField;
-import org.apache.lucene.index.DirectoryReader;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.RandomIndexWriter;
-import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
-import org.apache.lucene.store.Directory;
+import org.apache.lucene.search.Query;
 import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.search.aggregations.AggregatorTestCase;
 import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregationBuilder;
-import org.elasticsearch.search.aggregations.bucket.global.GlobalAggregator;
 import org.elasticsearch.search.aggregations.bucket.global.InternalGlobal;
 import org.elasticsearch.search.aggregations.metrics.InternalMin;
 import org.elasticsearch.search.aggregations.metrics.MinAggregationBuilder;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.function.BiConsumer;
 
-import static java.util.Collections.singleton;
-
 public class GlobalAggregatorTests extends AggregatorTestCase {
     public void testNoDocs() throws IOException {
         testCase(iw -> {
             // Intentionally not writing any docs
-        }, (global, min) -> {
+        }, new MatchAllDocsQuery(), (global, min) -> {
             assertEquals(0, global.getDocCount());
             assertEquals(Double.POSITIVE_INFINITY, min.getValue(), 0);
         });
     }
 
     public void testSomeDocs() throws IOException {
         testCase(iw -> {
-            iw.addDocument(singleton(new SortedNumericDocValuesField("number", 7)));
-            iw.addDocument(singleton(new SortedNumericDocValuesField("number", 1)));
-        }, (global, min) -> {
+            iw.addDocument(List.of(new SortedNumericDocValuesField("number", 7)));
+            iw.addDocument(List.of(new SortedNumericDocValuesField("number", 1)));
+        }, new MatchAllDocsQuery(), (global, min) -> {
             assertEquals(2, global.getDocCount());
             assertEquals(1, min.getValue(), 0);
         });
     }
 
-    // Note that `global`'s fancy support for ignoring the query comes from special code in AggregationPhase. We don't test that here.
-
-    private void testCase(CheckedConsumer<RandomIndexWriter, IOException> buildIndex, BiConsumer<InternalGlobal, InternalMin> verify)
-            throws IOException {
-        Directory directory = newDirectory();
-        RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory);
-        buildIndex.accept(indexWriter);
-        indexWriter.close();
-
-        IndexReader indexReader = DirectoryReader.open(directory);
-        IndexSearcher indexSearcher = newSearcher(indexReader, true, true);
+    public void testIgnoresQuery() throws IOException {
+        testCase(iw -> {
+            iw.addDocument(List.of(new SortedNumericDocValuesField("number", 7)));
+            iw.addDocument(List.of(new SortedNumericDocValuesField("number", 1)));
+        }, LongPoint.newRangeQuery("number", 2, Long.MAX_VALUE), (global, min) -> {
+            assertEquals(2, global.getDocCount());
+            assertEquals(1, min.getValue(), 0);
+        });
+    }
 
+    private void testCase(
+        CheckedConsumer<RandomIndexWriter, IOException> buildIndex,
+        Query topLevelQuery,
+        BiConsumer<InternalGlobal, InternalMin> verify
+    ) throws IOException {
         GlobalAggregationBuilder aggregationBuilder = new GlobalAggregationBuilder("_name");
         aggregationBuilder.subAggregation(new MinAggregationBuilder("in_global").field("number"));
         MappedFieldType fieldType = new NumberFieldMapper.NumberFieldType("number", NumberFieldMapper.NumberType.LONG);
 
-        GlobalAggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
-        aggregator.preCollection();
-        indexSearcher.search(new MatchAllDocsQuery(), aggregator);
-        aggregator.postCollection();
-        InternalGlobal result = (InternalGlobal) aggregator.buildTopLevel();
-        verify.accept(result, (InternalMin) result.getAggregations().asMap().get("in_global"));
-
-        indexReader.close();
-        directory.close();
+        testCase(aggregationBuilder, topLevelQuery, buildIndex, (InternalGlobal result) -> {
+            InternalMin min = result.getAggregations().get("in_global");
+            verify.accept(result, min);
+        }, fieldType);
     }
 }