Add bounds support for geogrid agg on shapes

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

@@ -322,7 +322,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

@@ -316,6 +316,5 @@ public double minX() {
         public double maxX() {
             return Math.max(negRight, posRight);
         }
-
     }
 }

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

@@ -31,9 +31,11 @@
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.aggregations.bucket.geogrid.BoundedGeoTileGridTiler;
 import org.elasticsearch.search.aggregations.bucket.geogrid.CellIdSource;
 import org.elasticsearch.search.aggregations.bucket.geogrid.GeoGridTiler;
 import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileGridAggregationBuilder;
+import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileGridTiler;
 import org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils;
 import org.elasticsearch.search.aggregations.support.ValueType;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
@@ -138,7 +140,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 = new GeoTileGridTiler();
+            } else {
+                tiler = new BoundedGeoTileGridTiler(geoBoundingBox);
+            }
+
+            CellIdSource cellIdSource = new CellIdSource(geoValue, precision, tiler);
             return new CompositeValuesSourceConfig(name, fieldType, cellIdSource, DocValueFormat.GEOTILE, order(),
                 missingBucket(), script() != null);
         } else {

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

@@ -0,0 +1,108 @@
+/*
+ * 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.GeoBoundingBox;
+import org.elasticsearch.common.geo.GeoRelation;
+import org.elasticsearch.geometry.Rectangle;
+import org.elasticsearch.geometry.utils.Geohash;
+import org.elasticsearch.index.fielddata.MultiGeoValues;
+
+public class BoundedGeoHashGridTiler extends GeoHashGridTiler {
+    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;
+
+    BoundedGeoHashGridTiler(GeoBoundingBox geoBoundingBox) {
+        // 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;
+        }
+    }
+
+    @Override
+    public int advancePointValue(long[] values, double x, double y, int precision, int valuesIdx) {
+        long hash = encode(x, y, precision);
+        if (cellIntersectsGeoBoundingBox(Geohash.toBoundingBox(Geohash.stringEncode(hash)))) {
+            values[valuesIdx] = hash;
+            return valuesIdx + 1;
+        }
+        return valuesIdx;
+    }
+
+    boolean cellIntersectsGeoBoundingBox(Rectangle rectangle) {
+        return (boundsTop >= rectangle.getMinY() && boundsBottom <= rectangle.getMaxY()
+            && (boundsEastLeft <= rectangle.getMaxX() && boundsEastRight >= rectangle.getMinX()
+            || (crossesDateline && boundsWestLeft <= rectangle.getMaxX() && boundsWestRight >= rectangle.getMinX())));
+    }
+
+    @Override
+    protected int setValue(CellValues docValues, MultiGeoValues.GeoValue geoValue, MultiGeoValues.BoundingBox bounds, int precision) {
+        String hash = Geohash.stringEncode(bounds.minX(), bounds.minY(), precision);
+        GeoRelation relation = relateTile(geoValue, hash);
+        if (relation != GeoRelation.QUERY_DISJOINT) {
+            docValues.resizeCell(1);
+            docValues.add(0, Geohash.longEncode(hash));
+            return 1;
+        }
+        return 0;
+    }
+
+    @Override
+    protected GeoRelation relateTile(MultiGeoValues.GeoValue geoValue, String hash) {
+        Rectangle rectangle = Geohash.toBoundingBox(hash);
+        if (cellIntersectsGeoBoundingBox(rectangle)) {
+            return geoValue.relate(rectangle);
+        } else {
+            return GeoRelation.QUERY_DISJOINT;
+        }
+    }
+
+    @Override
+    protected int setValuesForFullyContainedTile(String hash, CellValues values,
+                                               int valuesIndex, int targetPrecision) {
+        String[] hashes = Geohash.getSubGeohashes(hash);
+        for (int i = 0; i < hashes.length; i++) {
+            if (hashes[i].length() == targetPrecision ) {
+                if (cellIntersectsGeoBoundingBox(Geohash.toBoundingBox(hashes[i]))) {
+                    values.add(valuesIndex++, Geohash.longEncode(hashes[i]));
+                }
+            } else {
+                valuesIndex = setValuesForFullyContainedTile(hashes[i], values, valuesIndex, targetPrecision);
+            }
+        }
+        return valuesIndex;
+    }
+}

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

@@ -0,0 +1,107 @@
+/*
+ * 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.GeoBoundingBox;
+import org.elasticsearch.common.geo.GeoRelation;
+import org.elasticsearch.geometry.Rectangle;
+import org.elasticsearch.index.fielddata.MultiGeoValues;
+
+public class BoundedGeoTileGridTiler extends GeoTileGridTiler {
+    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;
+
+    public BoundedGeoTileGridTiler(GeoBoundingBox geoBoundingBox) {
+        // 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;
+        }
+    }
+
+    public int advancePointValue(long[] values, double x, double y, int precision, int valuesIdx) {
+        long hash = encode(x, y, precision);
+        if (cellIntersectsGeoBoundingBox(GeoTileUtils.toBoundingBox(hash))) {
+            values[valuesIdx] = hash;
+            return valuesIdx + 1;
+        }
+        return valuesIdx;
+    }
+
+    boolean cellIntersectsGeoBoundingBox(Rectangle rectangle) {
+        return (boundsTop >= rectangle.getMinY() && boundsBottom <= rectangle.getMaxY()
+            && (boundsEastLeft <= rectangle.getMaxX() && boundsEastRight >= rectangle.getMinX()
+            || (crossesDateline && boundsWestLeft <= rectangle.getMaxX() && boundsWestRight >= rectangle.getMinX())));
+    }
+
+    @Override
+    public GeoRelation relateTile(MultiGeoValues.GeoValue geoValue, int xTile, int yTile, int precision) {
+        Rectangle rectangle = GeoTileUtils.toBoundingBox(xTile, yTile, precision);
+        if (cellIntersectsGeoBoundingBox(rectangle)) {
+            return geoValue.relate(rectangle);
+        }
+        return GeoRelation.QUERY_DISJOINT;
+    }
+
+    @Override
+    protected int setValue(CellValues docValues, MultiGeoValues.GeoValue geoValue, int xTile, int yTile, int precision) {
+        if (cellIntersectsGeoBoundingBox(GeoTileUtils.toBoundingBox(xTile, yTile, precision))) {
+            docValues.resizeCell(1);
+            docValues.add(0, GeoTileUtils.longEncodeTiles(precision, xTile, yTile));
+            return 1;
+        }
+        return 0;
+    }
+
+    @Override
+    protected int setValuesForFullyContainedTile(int xTile, int yTile, int zTile, CellValues values, int valuesIndex,
+                                               int targetPrecision) {
+        zTile++;
+        for (int i = 0; i < 2; i++) {
+            for (int j = 0; j < 2; j++) {
+                int nextX = 2 * xTile + i;
+                int nextY = 2 * yTile + j;
+                if (zTile == targetPrecision) {
+                    if (cellIntersectsGeoBoundingBox(GeoTileUtils.toBoundingBox(nextX, nextY, zTile))) {
+                        values.add(valuesIndex++, GeoTileUtils.longEncodeTiles(zTile, nextX, nextY));
+                    }
+                } else {
+                    valuesIndex = setValuesForFullyContainedTile(nextX, nextY, zTile, values, valuesIndex, targetPrecision);
+                }
+            }
+        }
+        return valuesIndex;
+    }
+}

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

@@ -21,7 +21,6 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.elasticsearch.index.fielddata.MultiGeoValues;
-import org.elasticsearch.common.geo.GeoBoundingBox;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
 import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
@@ -36,13 +35,11 @@ public class CellIdSource extends ValuesSource.Numeric {
     private final ValuesSource.Geo valuesSource;
     private final int precision;
     private final GeoGridTiler encoder;
-    private final GeoBoundingBox geoBoundingBox;
 
-    public CellIdSource(ValuesSource.Geo valuesSource, int precision, GeoBoundingBox geoBoundingBox, GeoGridTiler encoder) {
+    public CellIdSource(ValuesSource.Geo valuesSource, int precision, GeoGridTiler encoder) {
         this.valuesSource = valuesSource;
         //different GeoPoints could map to the same or different hashing cells.
         this.precision = precision;
-        this.geoBoundingBox = geoBoundingBox;
         this.encoder = encoder;
     }
 
@@ -65,19 +62,10 @@ public SortedNumericDocValues longValues(LeafReaderContext ctx) {
         ValuesSourceType vs = geoValues.valuesSourceType();
         if (CoreValuesSourceType.GEOPOINT == vs) {
             // docValues are geo points
-            if (geoBoundingBox.isUnbounded()) {
-                return new UnboundedGeoPointCellValues(geoValues, precision, encoder);
-            } else {
-                return new BoundedGeoPointCellValues(geoValues, precision, encoder, geoBoundingBox);
-            }
+            return new GeoPointCellValues(geoValues, precision, encoder);
         } else if (CoreValuesSourceType.GEOSHAPE == vs || CoreValuesSourceType.GEO == vs) {
             // docValues are geo shapes
-            if (geoBoundingBox.isUnbounded()) {
-                return new GeoShapeCellValues(geoValues, precision, encoder);
-            } else {
-                // TODO(talevy): support unbounded
-                throw new IllegalArgumentException("bounded geogrid is not supported on geo_shape fields");
-            }
+            return new GeoShapeCellValues(geoValues, precision, encoder);
         } 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>.