Add bounds support for geogrid agg on shapes

server/src/main/java/org/elasticsearch/common/geo/TriangleTreeReader.java

@@ -320,7 +320,6 @@ private GeoRelation relateTriangle(int aX, int aY, boolean ab, int bX, int bY, b
             int tMinY = StrictMath.min(StrictMath.min(aY, bY), cY);
             int tMaxY = StrictMath.max(StrictMath.max(aY, bY), cY);
 
-
             // 1. check bounding boxes are disjoint, where north and east boundaries are not considered as crossing
             if (tMaxX <= minX || tMinX > maxX || tMinY > maxY || tMaxY <= minY) {
                 return GeoRelation.QUERY_DISJOINT;

server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java

@@ -42,6 +42,7 @@
 import java.text.ParseException;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * A stateful lightweight per document set of geo values.
@@ -58,7 +59,7 @@
  * The set of values associated with a document might contain duplicates and
  * comes in a non-specified order.
  */
-public abstract class MultiGeoValues {
+public abstract class MultiGeoValues <G extends MultiGeoValues.GeoValue> {
 
     /**
      * Creates a new {@link MultiGeoValues} instance
@@ -87,7 +88,7 @@ protected MultiGeoValues() {
      *
      * @return the next value for the current docID set to {@link #advanceExact(int)}.
      */
-    public abstract GeoValue nextValue() throws IOException;
+    public abstract G nextValue() throws IOException;
 
     public static class GeoPointValue implements GeoValue {
         private final GeoPoint geoPoint;
@@ -317,5 +318,22 @@ public double maxX() {
             return Math.max(negRight, posRight);
         }
 
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            BoundingBox that = (BoundingBox) o;
+            return Double.compare(that.top, top) == 0 &&
+                Double.compare(that.bottom, bottom) == 0 &&
+                Double.compare(that.negLeft, negLeft) == 0 &&
+                Double.compare(that.negRight, negRight) == 0 &&
+                Double.compare(that.posLeft, posLeft) == 0 &&
+                Double.compare(that.posRight, posRight) == 0;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(top, bottom, negLeft, negRight, posLeft, posRight);
+        }
     }
 }

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

@@ -138,7 +138,15 @@ protected CompositeValuesSourceConfig innerBuild(QueryShardContext queryShardCon
             ValuesSource.Geo geoValue = (ValuesSource.Geo) orig;
             // is specified in the builder.
             final MappedFieldType fieldType = config.fieldContext() != null ? config.fieldContext().fieldType() : null;
-            CellIdSource cellIdSource = new CellIdSource(geoValue, precision, geoBoundingBox, GeoGridTiler.GeoTileGridTiler.INSTANCE);
+
+            final GeoGridTiler tiler;
+            if (geoBoundingBox.isUnbounded()) {
+                tiler = GeoGridTiler.GeoTileGridTiler.INSTANCE;
+            } else {
+                tiler = GeoGridTiler.GeoTileGridTiler.BOUNDED_INSTANCE;
+            }
+
+            CellIdSource cellIdSource = new CellIdSource(geoValue, precision, geoBoundingBox, tiler);
             return new CompositeValuesSourceConfig(name, fieldType, cellIdSource, DocValueFormat.GEOTILE, order(),
                 missingBucket(), script() != null);
         } else {

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/BoundedGeoShapeCellValues.java

@@ -0,0 +1,121 @@
+/*
+ * 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.geogrid;
+
+import org.elasticsearch.common.geo.DimensionalShapeType;
+import org.elasticsearch.common.geo.GeoBoundingBox;
+import org.elasticsearch.common.geo.GeoRelation;
+import org.elasticsearch.geometry.Rectangle;
+import org.elasticsearch.index.fielddata.MultiGeoValues;
+
+/**
+ * A bounded representation of {@link UnboundedGeoShapeCellValues}
+ */
+class BoundedGeoShapeCellValues extends UnboundedGeoShapeCellValues {
+    private BoundedGeoValue boundedGeoValue;
+
+    protected BoundedGeoShapeCellValues(MultiGeoValues geoValues, int precision, GeoGridTiler tiler, GeoBoundingBox geoBoundingBox) {
+        super(geoValues, precision, tiler);
+        this.boundedGeoValue = new BoundedGeoValue(geoBoundingBox);
+    }
+
+    @Override
+    int advanceValue(MultiGeoValues.GeoValue target, int valuesIdx) {
+        boundedGeoValue.reset(target);
+        return super.advanceValue(boundedGeoValue, valuesIdx);
+    }
+
+    static class BoundedGeoValue implements MultiGeoValues.GeoValue {
+        private MultiGeoValues.GeoValue innerValue;
+        private final double boundsTop;
+        private final double boundsBottom;
+        private final double boundsWestLeft;
+        private final double boundsWestRight;
+        private final double boundsEastLeft;
+        private final double boundsEastRight;
+        private final boolean crossesDateline;
+
+        BoundedGeoValue(GeoBoundingBox geoBoundingBox) {
+            this.innerValue = null;
+
+            // split geoBoundingBox into west and east boxes
+            boundsTop = geoBoundingBox.top();
+            boundsBottom = geoBoundingBox.bottom();
+            if (geoBoundingBox.right() < geoBoundingBox.left()) {
+                boundsWestLeft = -180;
+                boundsWestRight = geoBoundingBox.right();
+                boundsEastLeft = geoBoundingBox.left();
+                boundsEastRight = 180;
+                crossesDateline = true;
+            } else { // only set east bounds
+                boundsEastLeft = geoBoundingBox.left();
+                boundsEastRight = geoBoundingBox.right();
+                boundsWestLeft = 0;
+                boundsWestRight = 0;
+                crossesDateline = false;
+            }
+        }
+
+        void reset(MultiGeoValues.GeoValue innerValue) {
+            this.innerValue = innerValue;
+        }
+
+        @Override
+        public GeoRelation relate(Rectangle rectangle) {
+            if (boundsTop >= rectangle.getMinY() && boundsBottom <= rectangle.getMaxY()
+                    && (boundsEastLeft <= rectangle.getMaxX() && boundsEastRight >= rectangle.getMinX()
+                    || (crossesDateline && boundsWestLeft <= rectangle.getMaxX() && boundsWestRight >= rectangle.getMinX()))) {
+                GeoRelation relation = innerValue.relate(rectangle);
+                // bounded relations can never be sure to be inside until each tile and sub-tile is checked
+                // explicitly against the geoBoundingBox
+                if (relation == GeoRelation.QUERY_INSIDE) {
+                    return GeoRelation.QUERY_CROSSES;
+                }
+                return relation;
+            } else {
+                return GeoRelation.QUERY_DISJOINT;
+            }
+        }
+
+        @Override
+        public double lat() {
+            return innerValue.lat();
+        }
+
+        @Override
+        public double lon() {
+            return innerValue.lon();
+        }
+
+        @Override
+        public MultiGeoValues.BoundingBox boundingBox() {
+            return innerValue.boundingBox();
+        }
+
+        @Override
+        public DimensionalShapeType dimensionalShapeType() {
+            return innerValue.dimensionalShapeType();
+        }
+
+        @Override
+        public double weight() {
+            return innerValue.weight();
+        }
+    }
+}

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/CellIdSource.java

@@ -73,10 +73,9 @@ public SortedNumericDocValues longValues(LeafReaderContext ctx) {
         } else if (CoreValuesSourceType.GEOSHAPE == vs || CoreValuesSourceType.GEO == vs) {
             // docValues are geo shapes
             if (geoBoundingBox.isUnbounded()) {
-                return new GeoShapeCellValues(geoValues, precision, encoder);
+                return new UnboundedGeoShapeCellValues(geoValues, precision, encoder);
             } else {
-                // TODO(talevy): support unbounded
-                throw new IllegalArgumentException("bounded geogrid is not supported on geo_shape fields");
+                return new BoundedGeoShapeCellValues(geoValues, precision, encoder, geoBoundingBox);
             }
         } else {
             throw new IllegalArgumentException("unsupported geo type");

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/CellValues.java

@@ -56,6 +56,19 @@ public boolean advanceExact(int docId) throws IOException {
         }
     }
 
+    // for testing
+    protected long[] getValues() {
+        return values;
+    }
+
+    protected void add(int idx, long value) {
+        values[idx] = value;
+    }
+
+    void resizeCell(int newSize) {
+        resize(newSize);
+    }
+
     /**
      * Sets the appropriate long-encoded value for <code>target</code>
      * in <code>values</code>.