Support geotile_grid aggregation in composite agg sources (#45810)

Adds support for `geotile_grid` as a source in composite aggs. 

Part of this change includes adding a new docFormat of `GEOTILE` that formats a hashed `long` value into a geotile formatting string `zoom/x/y`.

Author: benwtrent

rest-api-spec/src/main/resources/rest-api-spec/test/search.aggregation/230_composite.yml

@@ -12,6 +12,8 @@ setup:
                     type: keyword
                   long:
                     type: long
+                  geo_point:
+                    type: geo_point
                   nested:
                     type: nested
                     properties:
@@ -38,25 +40,25 @@ setup:
       index:
         index: test
         id:    1
-        body:  { "keyword": "foo", "long": [10, 20], "nested": [{"nested_long": 10}, {"nested_long": 20}] }
+        body:  { "keyword": "foo", "long": [10, 20], "geo_point": "37.2343,-115.8067", "nested": [{"nested_long": 10}, {"nested_long": 20}] }
 
   - do:
       index:
         index: test
         id:    2
-        body:  { "keyword": ["foo", "bar"] }
+        body:  { "keyword": ["foo", "bar"], "geo_point": "41.12,-71.34" }
 
   - do:
       index:
         index: test
         id:    3
-        body:  { "keyword": "bar", "long": [100, 0], "nested": [{"nested_long": 10}, {"nested_long": 0}] }
+        body:  { "keyword": "bar", "long": [100, 0], "geo_point": "90.0,0.0", "nested": [{"nested_long": 10}, {"nested_long": 0}] }
 
   - do:
       index:
         index: test
         id:    4
-        body:  { "keyword": "bar", "long": [1000, 0], "nested": [{"nested_long": 1000}, {"nested_long": 20}] }
+        body:  { "keyword": "bar", "long": [1000, 0], "geo_point": "41.12,-71.34", "nested": [{"nested_long": 1000}, {"nested_long": 20}] }
 
   - do:
       index:
@@ -615,3 +617,87 @@ setup:
                     }
                   ]
 
+---
+"Simple Composite aggregation with GeoTile grid":
+  - skip:
+      version: " - 7.99.99"
+      reason: geotile_grid is not supported until 8.0.0
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test
+        body:
+          aggregations:
+            test:
+              composite:
+                sources: [
+                  "geo": {
+                    "geotile_grid": {
+                      "field": "geo_point",
+                      "precision": 12
+                    }
+                  },
+                  {
+                    "kw": {
+                      "terms": {
+                        "field": "keyword"
+                      }
+                    }
+                  }
+                ]
+
+  - match: {hits.total: 6}
+  - length: { aggregations.test.buckets: 4 }
+  - match: { aggregations.test.buckets.0.key.geo: "12/730/1590" }
+  - match: { aggregations.test.buckets.0.key.kw: "foo" }
+  - match: { aggregations.test.buckets.0.doc_count: 1 }
+  - match: { aggregations.test.buckets.1.key.geo: "12/1236/1533" }
+  - match: { aggregations.test.buckets.1.key.kw: "bar" }
+  - match: { aggregations.test.buckets.1.doc_count: 2 }
+  - match: { aggregations.test.buckets.2.key.geo: "12/1236/1533" }
+  - match: { aggregations.test.buckets.2.key.kw: "foo" }
+  - match: { aggregations.test.buckets.2.doc_count: 1 }
+  - match: { aggregations.test.buckets.3.key.geo: "12/2048/0" }
+  - match: { aggregations.test.buckets.3.key.kw: "bar" }
+  - match: { aggregations.test.buckets.3.doc_count: 1 }
+---
+"Simple Composite aggregation with geotile grid add aggregate after":
+  - skip:
+      version: " - 7.99.99"
+      reason: geotile_grid is not supported until 8.0.0
+  - do:
+      search:
+        rest_total_hits_as_int: true
+        index: test
+        body:
+          aggregations:
+            test:
+              composite:
+                sources: [
+                  "geo": {
+                    "geotile_grid": {
+                      "field": "geo_point",
+                      "precision": 12
+                    }
+                  },
+                  {
+                    "kw": {
+                      "terms": {
+                        "field": "keyword"
+                      }
+                    }
+                  }
+                ]
+                after: { "geo": "12/730/1590", "kw": "foo" }
+
+  - match: {hits.total: 6}
+  - length: { aggregations.test.buckets: 3 }
+  - match: { aggregations.test.buckets.0.key.geo: "12/1236/1533" }
+  - match: { aggregations.test.buckets.0.key.kw: "bar" }
+  - match: { aggregations.test.buckets.0.doc_count: 2 }
+  - match: { aggregations.test.buckets.1.key.geo: "12/1236/1533" }
+  - match: { aggregations.test.buckets.1.key.kw: "foo" }
+  - match: { aggregations.test.buckets.1.doc_count: 1 }
+  - match: { aggregations.test.buckets.2.key.geo: "12/2048/0" }
+  - match: { aggregations.test.buckets.2.key.kw: "bar" }
+  - match: { aggregations.test.buckets.2.doc_count: 1 }

server/src/main/java/org/elasticsearch/search/DocValueFormat.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.common.time.DateMathParser;
 import org.elasticsearch.geometry.utils.Geohash;
 import org.elasticsearch.index.mapper.DateFieldMapper;
+import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils;
 
 import java.io.IOException;
 import java.net.InetAddress;
@@ -246,6 +247,28 @@ public String format(double value) {
         }
     };
 
+    DocValueFormat GEOTILE = new DocValueFormat() {
+
+        @Override
+        public String getWriteableName() {
+            return "geo_tile";
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) {
+        }
+
+        @Override
+        public String format(long value) {
+            return GeoTileUtils.stringEncode(value);
+        }
+
+        @Override
+        public String format(double value) {
+            return format((long) value);
+        }
+    };
+
     DocValueFormat BOOLEAN = new DocValueFormat() {
 
         @Override

server/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -667,6 +667,7 @@ private void registerValueFormats() {
         registerValueFormat(DocValueFormat.DateTime.NAME, DocValueFormat.DateTime::new);
         registerValueFormat(DocValueFormat.Decimal.NAME, DocValueFormat.Decimal::new);
         registerValueFormat(DocValueFormat.GEOHASH.getWriteableName(), in -> DocValueFormat.GEOHASH);
+        registerValueFormat(DocValueFormat.GEOTILE.getWriteableName(), in -> DocValueFormat.GEOTILE);
         registerValueFormat(DocValueFormat.IP.getWriteableName(), in -> DocValueFormat.IP);
         registerValueFormat(DocValueFormat.RAW.getWriteableName(), in -> DocValueFormat.RAW);
         registerValueFormat(DocValueFormat.BINARY.getWriteableName(), in -> DocValueFormat.BINARY);

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeAggregator.java

@@ -42,6 +42,7 @@
 import org.elasticsearch.search.aggregations.MultiBucketCollector;
 import org.elasticsearch.search.aggregations.MultiBucketConsumerService;
 import org.elasticsearch.search.aggregations.bucket.BucketsAggregator;
+import org.elasticsearch.search.aggregations.bucket.geogrid.CellIdSource;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
@@ -299,6 +300,17 @@ private SingleDimensionValuesSource<?> createValuesSource(BigArrays bigArrays, I
                 reverseMul
             );
 
+        } else if (config.valuesSource() instanceof CellIdSource) {
+            final CellIdSource cis = (CellIdSource) config.valuesSource();
+            return new GeoTileValuesSource(
+                bigArrays,
+                config.fieldType(),
+                cis::longValues,
+                LongUnaryOperator.identity(),
+                config.format(),
+                config.missingBucket(),
+                size,
+                reverseMul);
         } else if (config.valuesSource() instanceof ValuesSource.Numeric) {
             final ValuesSource.Numeric vs = (ValuesSource.Numeric) config.valuesSource();
             if (vs.isFloatingPoint()) {

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/CompositeValuesSourceParserHelper.java

@@ -19,6 +19,7 @@
 
 package org.elasticsearch.search.aggregations.bucket.composite;
 
+import org.elasticsearch.Version;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -67,6 +68,12 @@ public static void writeTo(CompositeValuesSourceBuilder<?> builder, StreamOutput
             code = 1;
         } else if (builder.getClass() == HistogramValuesSourceBuilder.class) {
             code = 2;
+        } else if (builder.getClass() == GeoTileGridValuesSourceBuilder.class) {
+            if (out.getVersion().before(Version.V_8_0_0)) {
+                throw new IOException("Attempting to serialize [" + builder.getClass().getSimpleName()
+                    + "] to a node with unsupported version [" + out.getVersion() + "]");
+            }
+            code = 3;
         } else {
             throw new IOException("invalid builder type: " + builder.getClass().getSimpleName());
         }
@@ -83,6 +90,8 @@ public static CompositeValuesSourceBuilder<?> readFrom(StreamInput in) throws IO
                 return new DateHistogramValuesSourceBuilder(in);
             case 2:
                 return new HistogramValuesSourceBuilder(in);
+            case 3:
+                return new GeoTileGridValuesSourceBuilder(in);
             default:
                 throw new IOException("Invalid code " + code);
         }
@@ -112,6 +121,9 @@ public static CompositeValuesSourceBuilder<?> fromXContent(XContentParser parser
             case HistogramValuesSourceBuilder.TYPE:
                 builder = HistogramValuesSourceBuilder.parse(name, parser);
                 break;
+            case GeoTileGridValuesSourceBuilder.TYPE:
+                builder = GeoTileGridValuesSourceBuilder.parse(name, parser);
+                break;
             default:
                 throw new ParsingException(parser.getTokenLocation(), "invalid source type: " + type);
         }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/GeoTileGridValuesSourceBuilder.java

@@ -0,0 +1,122 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.search.aggregations.bucket.composite;
+
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.xcontent.ObjectParser;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.index.mapper.MappedFieldType;
+import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.bucket.geogrid.CellIdSource;
+import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileGridAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils;
+import org.elasticsearch.search.aggregations.support.ValueType;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
+import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
+import org.elasticsearch.search.internal.SearchContext;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class GeoTileGridValuesSourceBuilder extends CompositeValuesSourceBuilder<GeoTileGridValuesSourceBuilder> {
+    static final String TYPE = "geotile_grid";
+
+    private static final ObjectParser<GeoTileGridValuesSourceBuilder, Void> PARSER;
+    static {
+        PARSER = new ObjectParser<>(GeoTileGridValuesSourceBuilder.TYPE);
+        PARSER.declareInt(GeoTileGridValuesSourceBuilder::precision, new ParseField("precision"));
+        CompositeValuesSourceParserHelper.declareValuesSourceFields(PARSER, ValueType.NUMERIC);
+    }
+
+    static GeoTileGridValuesSourceBuilder parse(String name, XContentParser parser) throws IOException {
+        return PARSER.parse(parser, new GeoTileGridValuesSourceBuilder(name), null);
+    }
+
+    private int precision = GeoTileGridAggregationBuilder.DEFAULT_PRECISION;
+
+    GeoTileGridValuesSourceBuilder(String name) {
+        super(name);
+    }
+
+    GeoTileGridValuesSourceBuilder(StreamInput in) throws IOException {
+        super(in);
+        this.precision = in.readInt();
+    }
+
+    public GeoTileGridValuesSourceBuilder precision(int precision) {
+        this.precision = GeoTileUtils.checkPrecisionRange(precision);
+        return this;
+    }
+
+    @Override
+    public GeoTileGridValuesSourceBuilder format(String format) {
+        throw new IllegalArgumentException("[format] is not supported for [" + TYPE + "]");
+    }
+
+    @Override
+    protected void innerWriteTo(StreamOutput out) throws IOException {
+        out.writeInt(precision);
+    }
+
+    @Override
+    protected void doXContentBody(XContentBuilder builder, Params params) throws IOException {
+        builder.field("precision", precision);
+    }
+
+    @Override
+    String type() {
+        return TYPE;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), precision);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) return true;
+        if (obj == null || getClass() != obj.getClass()) return false;
+        if (super.equals(obj) == false) return false;
+        GeoTileGridValuesSourceBuilder other = (GeoTileGridValuesSourceBuilder) obj;
+        return precision == other.precision;
+    }
+
+    @Override
+    protected CompositeValuesSourceConfig innerBuild(SearchContext context, ValuesSourceConfig<?> config) throws IOException {
+        ValuesSource orig = config.toValuesSource(context.getQueryShardContext());
+        if (orig == null) {
+            orig = ValuesSource.GeoPoint.EMPTY;
+        }
+        if (orig instanceof ValuesSource.GeoPoint) {
+            ValuesSource.GeoPoint geoPoint = (ValuesSource.GeoPoint) orig;
+            // is specified in the builder.
+            final MappedFieldType fieldType = config.fieldContext() != null ? config.fieldContext().fieldType() : null;
+            CellIdSource cellIdSource = new CellIdSource(geoPoint, precision, GeoTileUtils::longEncode);
+            return new CompositeValuesSourceConfig(name, fieldType, cellIdSource, DocValueFormat.GEOTILE, order(), missingBucket());
+        } else {
+            throw new IllegalArgumentException("invalid source, expected geo_point, got " + orig.getClass().getSimpleName());
+        }
+    }
+
+}