Add geo_bounds aggregation support for geo_shape

server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoBoundsAggregator.java

@@ -36,20 +36,20 @@
 import java.io.IOException;
 import java.util.Map;
 
-final class GeoBoundsAggregator extends MetricsAggregator {
+public class GeoBoundsAggregator extends MetricsAggregator {
 
-    static final ParseField WRAP_LONGITUDE_FIELD = new ParseField("wrap_longitude");
+    public static final ParseField WRAP_LONGITUDE_FIELD = new ParseField("wrap_longitude");
 
     private final ValuesSource.GeoPoint valuesSource;
     private final boolean wrapLongitude;
-    DoubleArray tops;
-    DoubleArray bottoms;
-    DoubleArray posLefts;
-    DoubleArray posRights;
-    DoubleArray negLefts;
-    DoubleArray negRights;
+    protected DoubleArray tops;
+    protected DoubleArray bottoms;
+    protected DoubleArray posLefts;
+    protected DoubleArray posRights;
+    protected DoubleArray negLefts;
+    protected DoubleArray negRights;
 
-    GeoBoundsAggregator(String name, SearchContext aggregationContext, Aggregator parent,
+    public GeoBoundsAggregator(String name, SearchContext aggregationContext, Aggregator parent,
             ValuesSource.GeoPoint valuesSource, boolean wrapLongitude, Map<String, Object> metadata) throws IOException {
         super(name, aggregationContext, parent, metadata);
         this.valuesSource = valuesSource;

server/src/main/java/org/elasticsearch/search/aggregations/metrics/GeoBoundsAggregatorSupplier.java

@@ -30,6 +30,6 @@
 @FunctionalInterface
 public interface GeoBoundsAggregatorSupplier extends AggregatorSupplier {
 
-    GeoBoundsAggregator build(String name, SearchContext aggregationContext, Aggregator parent,
+    MetricsAggregator build(String name, SearchContext aggregationContext, Aggregator parent,
                         ValuesSource valuesSource, boolean wrapLongitude, Map<String, Object> metadata) throws IOException;
 }

server/src/main/java/org/elasticsearch/search/aggregations/metrics/InternalGeoBounds.java

@@ -32,16 +32,16 @@
 import java.util.Objects;
 
 public class InternalGeoBounds extends InternalAggregation implements GeoBounds {
-    final double top;
-    final double bottom;
-    final double posLeft;
-    final double posRight;
-    final double negLeft;
-    final double negRight;
-    final boolean wrapLongitude;
-
-    InternalGeoBounds(String name, double top, double bottom, double posLeft, double posRight,
-                      double negLeft, double negRight, boolean wrapLongitude, Map<String, Object> metadata) {
+    public final double top;
+    public final double bottom;
+    public final double posLeft;
+    public final double posRight;
+    public final double negLeft;
+    public final double negRight;
+    public final boolean wrapLongitude;
+
+    public InternalGeoBounds(String name, double top, double bottom, double posLeft, double posRight,
+                             double negLeft, double negRight, boolean wrapLongitude, Map<String, Object> metadata) {
         super(name, metadata);
         this.top = top;
         this.bottom = bottom;

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java

@@ -14,8 +14,14 @@
 import org.elasticsearch.plugins.IngestPlugin;
 import org.elasticsearch.plugins.MapperPlugin;
 import org.elasticsearch.plugins.SearchPlugin;
+import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregationBuilder;
+import org.elasticsearch.search.aggregations.metrics.GeoBoundsAggregatorSupplier;
+import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
 import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction;
 import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction;
+import org.elasticsearch.xpack.spatial.aggregations.metrics.GeoShapeBoundsAggregator;
+import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeValuesSource;
+import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeValuesSourceType;
 import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper;
 import org.elasticsearch.xpack.spatial.index.mapper.PointFieldMapper;
 import org.elasticsearch.xpack.spatial.index.mapper.ShapeFieldMapper;
@@ -27,6 +33,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 
 import static java.util.Collections.singletonList;
 
@@ -53,8 +60,20 @@ public List<QuerySpec<?>> getQueries() {
         return singletonList(new QuerySpec<>(ShapeQueryBuilder.NAME, ShapeQueryBuilder::new, ShapeQueryBuilder::fromXContent));
     }
 
+    @Override
+    public List<Consumer<ValuesSourceRegistry>> getBareAggregatorRegistrar() {
+        return List.of(SpatialPlugin::registerGeoShapeBoundsAggregator);
+    }
+
     @Override
     public Map<String, Processor.Factory> getProcessors(Processor.Parameters parameters) {
         return Map.of(CircleProcessor.TYPE, new CircleProcessor.Factory());
     }
+
+    public static void registerGeoShapeBoundsAggregator(ValuesSourceRegistry valuesSourceRegistry) {
+        valuesSourceRegistry.register(GeoBoundsAggregationBuilder.NAME, GeoShapeValuesSourceType.INSTANCE,
+            (GeoBoundsAggregatorSupplier) (name, aggregationContext, parent, valuesSource, wrapLongitude, metadata)
+                -> new GeoShapeBoundsAggregator(name, aggregationContext, parent, (GeoShapeValuesSource) valuesSource,
+                wrapLongitude, metadata));
+    }
 }

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/aggregations/metrics/GeoShapeBoundsAggregator.java

@@ -0,0 +1,134 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.spatial.aggregations.metrics;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.common.lease.Releasables;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.common.util.DoubleArray;
+import org.elasticsearch.search.aggregations.Aggregator;
+import org.elasticsearch.search.aggregations.InternalAggregation;
+import org.elasticsearch.search.aggregations.LeafBucketCollector;
+import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
+import org.elasticsearch.search.aggregations.metrics.InternalGeoBounds;
+import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
+import org.elasticsearch.search.internal.SearchContext;
+import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeValuesSource;
+import org.elasticsearch.xpack.spatial.index.mapper.MultiGeoShapeValues;
+
+import java.io.IOException;
+import java.util.Map;
+
+public final class GeoShapeBoundsAggregator extends MetricsAggregator {
+    private final GeoShapeValuesSource valuesSource;
+    private final boolean wrapLongitude;
+    protected DoubleArray tops;
+    protected DoubleArray bottoms;
+    protected DoubleArray posLefts;
+    protected DoubleArray posRights;
+    protected DoubleArray negLefts;
+    protected DoubleArray negRights;
+
+    public GeoShapeBoundsAggregator(String name, SearchContext aggregationContext, Aggregator parent,
+                             GeoShapeValuesSource valuesSource, boolean wrapLongitude, Map<String, Object> metadata) throws IOException {
+        super(name, aggregationContext, parent, metadata);
+        this.valuesSource = valuesSource;
+        this.wrapLongitude = wrapLongitude;
+        if (valuesSource != null) {
+            final BigArrays bigArrays = context.bigArrays();
+            tops = bigArrays.newDoubleArray(1, false);
+            tops.fill(0, tops.size(), Double.NEGATIVE_INFINITY);
+            bottoms = bigArrays.newDoubleArray(1, false);
+            bottoms.fill(0, bottoms.size(), Double.POSITIVE_INFINITY);
+            posLefts = bigArrays.newDoubleArray(1, false);
+            posLefts.fill(0, posLefts.size(), Double.POSITIVE_INFINITY);
+            posRights = bigArrays.newDoubleArray(1, false);
+            posRights.fill(0, posRights.size(), Double.NEGATIVE_INFINITY);
+            negLefts = bigArrays.newDoubleArray(1, false);
+            negLefts.fill(0, negLefts.size(), Double.POSITIVE_INFINITY);
+            negRights = bigArrays.newDoubleArray(1, false);
+            negRights.fill(0, negRights.size(), Double.NEGATIVE_INFINITY);
+        }
+    }
+
+    @Override
+    public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
+            LeafBucketCollector sub) {
+        if (valuesSource == null) {
+            return LeafBucketCollector.NO_OP_COLLECTOR;
+        }
+        final BigArrays bigArrays = context.bigArrays();
+        final MultiGeoShapeValues values = valuesSource.geoShapeValues(ctx);
+        return new LeafBucketCollectorBase(sub, values) {
+            @Override
+            public void collect(int doc, long bucket) throws IOException {
+                if (bucket >= tops.size()) {
+                    long from = tops.size();
+                    tops = bigArrays.grow(tops, bucket + 1);
+                    tops.fill(from, tops.size(), Double.NEGATIVE_INFINITY);
+                    bottoms = bigArrays.resize(bottoms, tops.size());
+                    bottoms.fill(from, bottoms.size(), Double.POSITIVE_INFINITY);
+                    posLefts = bigArrays.resize(posLefts, tops.size());
+                    posLefts.fill(from, posLefts.size(), Double.POSITIVE_INFINITY);
+                    posRights = bigArrays.resize(posRights, tops.size());
+                    posRights.fill(from, posRights.size(), Double.NEGATIVE_INFINITY);
+                    negLefts = bigArrays.resize(negLefts, tops.size());
+                    negLefts.fill(from, negLefts.size(), Double.POSITIVE_INFINITY);
+                    negRights = bigArrays.resize(negRights, tops.size());
+                    negRights.fill(from, negRights.size(), Double.NEGATIVE_INFINITY);
+                }
+
+                if (values.advanceExact(doc)) {
+                    final int valuesCount = values.docValueCount();
+
+                    for (int i = 0; i < valuesCount; ++i) {
+                        MultiGeoShapeValues.GeoShapeValue value = values.nextValue();
+                        MultiGeoShapeValues.BoundingBox bounds = value.boundingBox();
+                        double top = Math.max(tops.get(bucket), bounds.top);
+                        double bottom = Math.min(bottoms.get(bucket), bounds.bottom);
+                        double posLeft = Math.min(posLefts.get(bucket), bounds.posLeft);
+                        double posRight = Math.max(posRights.get(bucket), bounds.posRight);
+                        double negLeft = Math.min(negLefts.get(bucket), bounds.negLeft);
+                        double negRight = Math.max(negRights.get(bucket), bounds.negRight);
+                        tops.set(bucket, top);
+                        bottoms.set(bucket, bottom);
+                        posLefts.set(bucket, posLeft);
+                        posRights.set(bucket, posRight);
+                        negLefts.set(bucket, negLeft);
+                        negRights.set(bucket, negRight);
+                    }
+                }
+            }
+        };
+    }
+
+
+    @Override
+    public InternalAggregation buildAggregation(long owningBucketOrdinal) {
+        if (valuesSource == null) {
+            return buildEmptyAggregation();
+        }
+        double top = tops.get(owningBucketOrdinal);
+        double bottom = bottoms.get(owningBucketOrdinal);
+        double posLeft = posLefts.get(owningBucketOrdinal);
+        double posRight = posRights.get(owningBucketOrdinal);
+        double negLeft = negLefts.get(owningBucketOrdinal);
+        double negRight = negRights.get(owningBucketOrdinal);
+        return new InternalGeoBounds(name, top, bottom, posLeft, posRight, negLeft, negRight, wrapLongitude, metadata());
+    }
+
+    @Override
+    public InternalAggregation buildEmptyAggregation() {
+        return new InternalGeoBounds(name, Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY, Double.POSITIVE_INFINITY,
+            Double.NEGATIVE_INFINITY, Double.POSITIVE_INFINITY, Double.NEGATIVE_INFINITY, wrapLongitude, metadata());
+    }
+
+    @Override
+    public void doClose() {
+        Releasables.close(tops, bottoms, posLefts, posRights, negLefts, negRights);
+    }
+}

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeValuesSource.java

@@ -29,7 +29,7 @@ public SortedBinaryDocValues bytesValues(LeafReaderContext context) throws IOExc
 
     };
 
-    abstract MultiGeoShapeValues geoShapeValues(LeafReaderContext context);
+    public abstract MultiGeoShapeValues geoShapeValues(LeafReaderContext context);
 
     @Override
     public DocValueBits docsWithValue(LeafReaderContext context) throws IOException {

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeValuesSourceType.java

@@ -25,7 +25,7 @@
 
 public class GeoShapeValuesSourceType implements Writeable, ValuesSourceType {
 
-    static GeoShapeValuesSourceType INSTANCE = new GeoShapeValuesSourceType();
+    public static GeoShapeValuesSourceType INSTANCE = new GeoShapeValuesSourceType();
 
     @Override
     public ValuesSource getEmpty() {
@@ -58,7 +58,7 @@ public ValuesSource replaceMissing(ValuesSource valuesSource, Object rawMissing,
         final MultiGeoShapeValues.GeoShapeValue missing = MultiGeoShapeValues.GeoShapeValue.missing(rawMissing.toString());
         return new GeoShapeValuesSource() {
             @Override
-            MultiGeoShapeValues geoShapeValues(LeafReaderContext context) {
+            public MultiGeoShapeValues geoShapeValues(LeafReaderContext context) {
                 MultiGeoShapeValues values = geoShapeValuesSource.geoShapeValues(context);
                 return new MultiGeoShapeValues() {
 

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.index.mapper.TypeParsers;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.VectorGeoShapeQueryProcessor;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
 
 import java.io.IOException;
 import java.util.HashMap;
@@ -170,6 +171,11 @@ public Query existsQuery(QueryShardContext context) {
             }
         }
 
+        @Override
+        public ValuesSourceType getValuesSourceType() {
+            return GeoShapeValuesSourceType.INSTANCE;
+        }
+
         @Override
         public GeoShapeWithDocValuesFieldType clone() {
             return new GeoShapeWithDocValuesFieldType(this);