Adds rewrite phase to aggregations (#25495)

* Adds rewrite phase to aggregations

This change adds aggregations to the rewrite performed by the `SearchSourceBuilder`. This means that `AggregationBuilder`s are able to implement a `rewrite()` method where they can return a new `AggregationBuilder` which is functionally the same but in a more primitive form. This is exactly analogous to the rewrite done by the `QueryBuilder`s.

The first aggregation to implement the rewrite are the filter and filters aggregations so they can rewrite the filters they contain.

Closes #17676

* Removes rewrite from PipelineAggregationBuilder

Rewrite is based on shard level information. Since pipeline aggregation are run in the reduce phase it doesn’t make sense to rewrite them on the shards. In fact eventually we shouldn’t be transporting them to the shards at all and should be retaining them on the coordinating node for execution in the reduce phase

* Addresses review comments

* addresses more review comments

* Fixed imports

2 files changed, 62 insertions, 14 deletions

diff --git a/core/src/test/java/org/elasticsearch/indices/IndicesRequestCacheIT.java b/core/src/test/java/org/elasticsearch/indices/IndicesRequestCacheIT.java
index 1e97d4dd57..7f8fe5d40a 100644
--- a/core/src/test/java/org/elasticsearch/indices/IndicesRequestCacheIT.java
+++ b/core/src/test/java/org/elasticsearch/indices/IndicesRequestCacheIT.java
@@ -23,7 +23,6 @@ import org.elasticsearch.action.admin.indices.alias.Alias;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.cluster.metadata.IndexMetaData;
-import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
 import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;
@@ -37,8 +36,8 @@ import org.joda.time.format.DateTimeFormat;
 import java.util.List;
 
 import static org.elasticsearch.search.aggregations.AggregationBuilders.dateHistogram;
-import static org.elasticsearch.search.aggregations.AggregationBuilders.filter;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.dateRange;
+import static org.elasticsearch.search.aggregations.AggregationBuilders.filter;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse;
 import static org.hamcrest.Matchers.equalTo;
@@ -411,18 +410,7 @@ public class IndicesRequestCacheIT extends ESIntegTestCase {
         assertThat(client().admin().indices().prepareStats("index").setRequestCache(true).get().getTotal().getRequestCache().getMissCount(),
                 equalTo(0L));
 
-        // If the request has an aggregation containing now we should not cache
-        final SearchResponse r4 = client().prepareSearch("index").setSearchType(SearchType.QUERY_THEN_FETCH).setSize(0)
-                .setRequestCache(true).setQuery(QueryBuilders.rangeQuery("s").gte("2016-03-20").lte("2016-03-26"))
-                .addAggregation(filter("foo", QueryBuilders.rangeQuery("s").from("now-10y").to("now"))).get();
-        assertSearchResponse(r4);
-        assertThat(r4.getHits().getTotalHits(), equalTo(7L));
-        assertThat(client().admin().indices().prepareStats("index").setRequestCache(true).get().getTotal().getRequestCache().getHitCount(),
-                equalTo(0L));
-        assertThat(client().admin().indices().prepareStats("index").setRequestCache(true).get().getTotal().getRequestCache().getMissCount(),
-                equalTo(0L));
-
-        // If the request has an aggregation containng now we should not cache
+        // If the request has an non-filter aggregation containing now we should not cache
         final SearchResponse r5 = client().prepareSearch("index").setSearchType(SearchType.QUERY_THEN_FETCH).setSize(0)
                 .setRequestCache(true).setQuery(QueryBuilders.rangeQuery("s").gte("2016-03-20").lte("2016-03-26"))
                 .addAggregation(dateRange("foo").field("s").addRange("now-10y", "now")).get();
@@ -442,6 +430,17 @@ public class IndicesRequestCacheIT extends ESIntegTestCase {
                 equalTo(0L));
         assertThat(client().admin().indices().prepareStats("index").setRequestCache(true).get().getTotal().getRequestCache().getMissCount(),
                 equalTo(2L));
+
+        // If the request has a filter aggregation containing now we should cache since it gets rewritten
+        final SearchResponse r4 = client().prepareSearch("index").setSearchType(SearchType.QUERY_THEN_FETCH).setSize(0)
+                .setRequestCache(true).setQuery(QueryBuilders.rangeQuery("s").gte("2016-03-20").lte("2016-03-26"))
+                .addAggregation(filter("foo", QueryBuilders.rangeQuery("s").from("now-10y").to("now"))).get();
+        assertSearchResponse(r4);
+        assertThat(r4.getHits().getTotalHits(), equalTo(7L));
+        assertThat(client().admin().indices().prepareStats("index").setRequestCache(true).get().getTotal().getRequestCache().getHitCount(),
+                equalTo(0L));
+        assertThat(client().admin().indices().prepareStats("index").setRequestCache(true).get().getTotal().getRequestCache().getMissCount(),
+                equalTo(4L));
     }
 
     public void testCacheWithFilteredAlias() {
diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/AggregatorFactoriesTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/AggregatorFactoriesTests.java
index 4d5af51ec1..f113801965 100644
--- a/core/src/test/java/org/elasticsearch/search/aggregations/AggregatorFactoriesTests.java
+++ b/core/src/test/java/org/elasticsearch/search/aggregations/AggregatorFactoriesTests.java
@@ -19,15 +19,25 @@
 package org.elasticsearch.search.aggregations;
 
 import org.elasticsearch.common.ParsingException;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
 import org.elasticsearch.common.xcontent.XContent;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.env.Environment;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryRewriteContext;
+import org.elasticsearch.index.query.TermsQueryBuilder;
+import org.elasticsearch.index.query.WrapperQueryBuilder;
+import org.elasticsearch.script.Script;
 import org.elasticsearch.search.SearchModule;
+import org.elasticsearch.search.aggregations.bucket.filter.FilterAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregatorBuilders;
+import org.elasticsearch.search.aggregations.pipeline.bucketscript.BucketScriptPipelineAggregationBuilder;
 import org.elasticsearch.test.AbstractQueryTestCase;
 import org.elasticsearch.test.ESTestCase;
 
@@ -39,6 +49,7 @@ import java.util.regex.Pattern;
 import static java.util.Collections.emptyList;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 
 public class AggregatorFactoriesTests extends ESTestCase {
     private String[] currentTypes;
@@ -236,6 +247,44 @@ public class AggregatorFactoriesTests extends ESTestCase {
         assertThat(e.toString(), containsString("Expected [START_OBJECT] under [field], but got a [VALUE_STRING] in [tag_count]"));
     }
 
+    public void testRewrite() throws Exception {
+        XContentType xContentType = randomFrom(XContentType.values());
+        BytesReference bytesReference;
+        try (XContentBuilder builder = XContentFactory.contentBuilder(xContentType)) {
+            builder.startObject();
+            {
+                builder.startObject("terms");
+                {
+                    builder.array("title", "foo");
+                }
+                builder.endObject();
+            }
+            builder.endObject();
+            bytesReference = builder.bytes();
+        }
+        FilterAggregationBuilder filterAggBuilder = new FilterAggregationBuilder("titles", new WrapperQueryBuilder(bytesReference));
+        BucketScriptPipelineAggregationBuilder pipelineAgg = new BucketScriptPipelineAggregationBuilder("const", new Script("1"));
+        AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(filterAggBuilder)
+                .addPipelineAggregator(pipelineAgg);
+        AggregatorFactories.Builder rewritten = builder
+                .rewrite(new QueryRewriteContext(null, null, null, xContentRegistry, null, null, () -> 0L));
+        assertNotSame(builder, rewritten);
+        List<AggregationBuilder> aggregatorFactories = rewritten.getAggregatorFactories();
+        assertEquals(1, aggregatorFactories.size());
+        assertThat(aggregatorFactories.get(0), instanceOf(FilterAggregationBuilder.class));
+        FilterAggregationBuilder rewrittenFilterAggBuilder = (FilterAggregationBuilder) aggregatorFactories.get(0);
+        assertNotSame(filterAggBuilder, rewrittenFilterAggBuilder);
+        assertNotEquals(filterAggBuilder, rewrittenFilterAggBuilder);
+        // Check the filter was rewritten from a wrapper query to a terms query
+        QueryBuilder rewrittenFilter = rewrittenFilterAggBuilder.getFilter();
+        assertThat(rewrittenFilter, instanceOf(TermsQueryBuilder.class));
+        
+        // Check that a further rewrite returns the same aggregation factories builder
+        AggregatorFactories.Builder secondRewritten = rewritten
+                .rewrite(new QueryRewriteContext(null, null, null, xContentRegistry, null, null, () -> 0L));
+        assertSame(rewritten, secondRewritten);
+    }
+
     @Override
     protected NamedXContentRegistry xContentRegistry() {
         return xContentRegistry;