Add geo_shape support for geotile_grid and geohash_grid (#55966)

this commit adds aggregation support for the geo_shape field
type on geo*_grid aggregations.

it introduces a Tiler for both tiles and hashes that enables a new type of
ValuesSource to replace the GeoPoint's CellIdSource. This makes it possible
for the existing Aggregator to be re-used, so no new implementations of
the grid aggregators are added.

Author: talevy

libs/geo/src/main/java/org/elasticsearch/geometry/utils/Geohash.java

@@ -53,6 +53,20 @@ public class Geohash {
     /** Bit encoded representation of the latitude of north pole */
     private static final long MAX_LAT_BITS = (0x1L << (PRECISION * 5 / 2)) - 1;
 
+    // Below code is adapted from the spatial4j library (GeohashUtils.java) Apache 2.0 Licensed
+    private static final double[] precisionToLatHeight, precisionToLonWidth;
+    static {
+        precisionToLatHeight = new double[PRECISION + 1];
+        precisionToLonWidth = new double[PRECISION + 1];
+        precisionToLatHeight[0] = 90*2;
+        precisionToLonWidth[0] = 180*2;
+        boolean even = false;
+        for(int i = 1; i <= PRECISION; i++) {
+            precisionToLatHeight[i] = precisionToLatHeight[i-1] / (even ? 8 : 4);
+            precisionToLonWidth[i] = precisionToLonWidth[i-1] / (even ? 4 : 8);
+            even = ! even;
+        }
+    }
 
     // no instance:
     private Geohash() {
@@ -97,6 +111,16 @@ public static Rectangle toBoundingBox(final String geohash) {
         }
     }
 
+    /** Array of geohashes one level below the baseGeohash. Sorted. */
+    public static String[] getSubGeohashes(String baseGeohash) {
+        String[] hashes = new String[BASE_32.length];
+        for (int i = 0; i < BASE_32.length; i++) {//note: already sorted
+            char c = BASE_32[i];
+            hashes[i] = baseGeohash+c;
+        }
+        return hashes;
+    }
+
     /**
      * Calculate all neighbors of a given geohash cell.
      *
@@ -201,6 +225,13 @@ public static final String getNeighbor(String geohash, int level, int dx, int dy
         }
     }
 
+    /**
+     * Encode a string geohash to the geohash based long format (lon/lat interleaved, 4 least significant bits = level)
+     */
+    public static final long longEncode(String hash) {
+        return longEncode(hash, hash.length());
+    }
+
     /**
      * Encode lon/lat to the geohash based long format (lon/lat interleaved, 4 least significant bits = level)
      */
@@ -290,6 +321,16 @@ public static long mortonEncode(final String hash) {
         return BitUtil.flipFlop(l);
     }
 
+    /** approximate width of geohash tile for a specific precision in degrees */
+    public static double lonWidthInDegrees(int precision) {
+        return precisionToLonWidth[precision];
+    }
+
+    /** approximate height of geohash tile for a specific precision in degrees */
+    public static double latHeightInDegrees(int precision) {
+        return precisionToLatHeight[precision];
+    }
+
     private static long encodeLatLon(final double lat, final double lon) {
         // encode lat/lon flipping the sign bit so negative ints sort before positive ints
         final int latEnc = encodeLatitude(lat) ^ 0x80000000;

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

@@ -28,6 +28,7 @@
 import org.elasticsearch.search.aggregations.LeafBucketCollector;
 import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
 import org.elasticsearch.search.aggregations.bucket.BucketsAggregator;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -43,10 +44,11 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid> extends Bucke
 
     protected final int requiredSize;
     protected final int shardSize;
-    protected final CellIdSource valuesSource;
+    protected final ValuesSource.Numeric valuesSource;
     protected final LongHash bucketOrds;
+    protected SortedNumericDocValues values;
 
-    GeoGridAggregator(String name, AggregatorFactories factories, CellIdSource valuesSource,
+    GeoGridAggregator(String name, AggregatorFactories factories, ValuesSource.Numeric valuesSource,
                       int requiredSize, int shardSize, SearchContext aggregationContext,
                       Aggregator parent, Map<String, Object> metadata) throws IOException {
         super(name, factories, aggregationContext, parent, metadata);
@@ -67,7 +69,7 @@ public ScoreMode scoreMode() {
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
             final LeafBucketCollector sub) throws IOException {
-        final SortedNumericDocValues values = valuesSource.longValues(ctx);
+        values = valuesSource.longValues(ctx);
         return new LeafBucketCollectorBase(sub, null) {
             @Override
             public void collect(int doc, long bucket) throws IOException {

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

@@ -20,6 +20,7 @@
 
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -32,9 +33,9 @@
  */
 public class GeoHashGridAggregator extends GeoGridAggregator<InternalGeoHashGrid> {
 
-    GeoHashGridAggregator(String name, AggregatorFactories factories, CellIdSource valuesSource,
-                          int requiredSize, int shardSize, SearchContext aggregationContext,
-                          Aggregator parent, Map<String, Object> metadata) throws IOException {
+    public GeoHashGridAggregator(String name, AggregatorFactories factories, ValuesSource.Numeric valuesSource,
+                                 int requiredSize, int shardSize, SearchContext aggregationContext,
+                                 Aggregator parent, Map<String, Object> metadata) throws IOException {
         super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, metadata);
     }
 

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

@@ -21,6 +21,7 @@
 
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -33,9 +34,9 @@
  */
 public class GeoTileGridAggregator extends GeoGridAggregator<InternalGeoTileGrid> {
 
-    GeoTileGridAggregator(String name, AggregatorFactories factories, CellIdSource valuesSource,
-                          int requiredSize, int shardSize, SearchContext aggregationContext,
-                          Aggregator parent, Map<String, Object> metadata) throws IOException {
+    public GeoTileGridAggregator(String name, AggregatorFactories factories, ValuesSource.Numeric valuesSource,
+                                 int requiredSize, int shardSize, SearchContext aggregationContext,
+                                 Aggregator parent, Map<String, Object> metadata) throws IOException {
         super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, metadata);
     }
 

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

@@ -18,12 +18,15 @@
  */
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
+import org.apache.lucene.geo.GeoEncodingUtils;
+import org.apache.lucene.util.SloppyMath;
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.util.ESSloppyMath;
 import org.elasticsearch.common.xcontent.ObjectParser.ValueType;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
+import org.elasticsearch.geometry.Rectangle;
 
 import java.io.IOException;
 import java.util.Locale;
@@ -43,6 +46,8 @@ public final class GeoTileUtils {
 
     private GeoTileUtils() {}
 
+    private static final double PI_DIV_2 = Math.PI / 2;
+
     /**
      * Largest number of tiles (precision) to use.
      * This value cannot be more than (64-5)/2 = 29, because 5 bits are used for zoom level itself (0-31)
@@ -53,6 +58,18 @@ private GeoTileUtils() {}
      */
     public static final int MAX_ZOOM = 29;
 
+    /**
+     * The geo-tile map is clipped at 85.05112878 to 90 and -85.05112878 to -90
+     */
+    public static final double LATITUDE_MASK = 85.0511287798066;
+
+    /**
+     * Since shapes are encoded, their boundaries are to be compared to against the encoded/decoded values of <code>LATITUDE_MASK</code>
+     */
+    public static final double NORMALIZED_LATITUDE_MASK = GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(LATITUDE_MASK));
+    public static final double NORMALIZED_NEGATIVE_LATITUDE_MASK =
+        GeoEncodingUtils.decodeLatitude(GeoEncodingUtils.encodeLatitude(-LATITUDE_MASK));
+
     /**
      * Bit position of the zoom value within hash - zoom is stored in the most significant 6 bits of a long number.
      */
@@ -63,6 +80,7 @@ private GeoTileUtils() {}
      */
     private static final long X_Y_VALUE_MASK = (1L << MAX_ZOOM) - 1;
 
+
     /**
      * Parse an integer precision (zoom level). The {@link ValueType#INT} allows it to be a number or a string.
      *
@@ -90,37 +108,65 @@ public static int checkPrecisionRange(int precision) {
     }
 
     /**
-     * Encode lon/lat to the geotile based long format.
-     * The resulting hash contains interleaved tile X and Y coordinates.
-     * The precision itself is also encoded as a few high bits.
+     * Calculates the x-coordinate in the tile grid for the specified longitude given
+     * the number of tile columns for a pre-determined zoom-level.
+     *
+     * @param longitude the longitude to use when determining the tile x-coordinate
+     * @param tiles     the number of tiles per row for a pre-determined zoom-level
      */
-    public static long longEncode(double longitude, double latitude, int precision) {
-        // Mathematics for this code was adapted from https://wiki.openstreetmap.org/wiki/Slippy_map_tilenames#Java
-
-        // Number of tiles for the current zoom level along the X and Y axis
-        final long tiles = 1 << checkPrecisionRange(precision);
-
-        long xTile = (long) Math.floor((normalizeLon(longitude) + 180) / 360 * tiles);
+    public static int getXTile(double longitude, long tiles) {
+        // normalizeLon treats this as 180, which is not friendly for tile mapping
+        if (longitude == -180) {
+            return 0;
+        }
 
-        double latSin = Math.sin(Math.toRadians(normalizeLat(latitude)));
-        long yTile = (long) Math.floor((0.5 - (Math.log((1 + latSin) / (1 - latSin)) / (4 * Math.PI))) * tiles);
+        int xTile = (int) Math.floor((normalizeLon(longitude) + 180) / 360 * tiles);
 
         // Edge values may generate invalid values, and need to be clipped.
         // For example, polar regions (above/below lat 85.05112878) get normalized.
         if (xTile < 0) {
-            xTile = 0;
+            return 0;
         }
         if (xTile >= tiles) {
-            xTile = tiles - 1;
+            return (int) tiles - 1;
         }
+
+        return xTile;
+    }
+
+    /**
+     * Calculates the y-coordinate in the tile grid for the specified longitude given
+     * the number of tile rows for pre-determined zoom-level.
+     *
+     * @param latitude  the latitude to use when determining the tile y-coordinate
+     * @param tiles     the number of tiles per column for a pre-determined zoom-level
+     */
+    public static int getYTile(double latitude, long tiles) {
+        double latSin = SloppyMath.cos(PI_DIV_2 - Math.toRadians(normalizeLat(latitude)));
+        int yTile = (int) Math.floor((0.5 - (Math.log((1 + latSin) / (1 - latSin)) / (4 * Math.PI))) * tiles);
+
         if (yTile < 0) {
             yTile = 0;
         }
         if (yTile >= tiles) {
-            yTile = tiles - 1;
+            return (int) tiles - 1;
         }
 
-        return longEncode((long) precision, xTile, yTile);
+        return yTile;
+    }
+
+    /**
+     * Encode lon/lat to the geotile based long format.
+     * The resulting hash contains interleaved tile X and Y coordinates.
+     * The precision itself is also encoded as a few high bits.
+     */
+    public static long longEncode(double longitude, double latitude, int precision) {
+        // Mathematics for this code was adapted from https://wiki.openstreetmap.org/wiki/Slippy_map_tilenames#Java
+        // Number of tiles for the current zoom level along the X and Y axis
+        final long tiles = 1 << checkPrecisionRange(precision);
+        long xTile = getXTile(longitude, tiles);
+        long yTile = getYTile(latitude, tiles);
+        return longEncodeTiles(precision, xTile, yTile);
     }
 
     /**
@@ -131,7 +177,14 @@ public static long longEncode(double longitude, double latitude, int precision)
      */
     public static long longEncode(String hashAsString) {
         int[] parsed = parseHash(hashAsString);
-        return longEncode((long)parsed[0], (long)parsed[1], (long)parsed[2]);
+        return longEncode((long) parsed[0], (long) parsed[1], (long) parsed[2]);
+    }
+
+    public static long longEncodeTiles(int precision, long xTile, long yTile) {
+        // Zoom value is placed in front of all the bits used for the geotile
+        // e.g. when max zoom is 29, the largest index would use 58 bits (57th..0th),
+        // leaving 5 bits unused for zoom. See MAX_ZOOM comment above.
+        return ((long) precision << ZOOM_SHIFT) | (xTile << MAX_ZOOM) | yTile;
     }
 
     /**
@@ -193,6 +246,23 @@ static GeoPoint keyToGeoPoint(String hashAsString) {
         return zxyToGeoPoint(hashAsInts[0], hashAsInts[1], hashAsInts[2]);
     }
 
+    public static Rectangle toBoundingBox(long hash) {
+        int[] hashAsInts = parseHash(hash);
+        return toBoundingBox(hashAsInts[1], hashAsInts[2], hashAsInts[0]);
+    }
+
+    public static Rectangle toBoundingBox(int xTile, int yTile, int precision) {
+        final double tiles = validateZXY(precision, xTile, yTile);
+        final double minN = Math.PI - (2.0 * Math.PI * (yTile + 1)) / tiles;
+        final double maxN = Math.PI - (2.0 * Math.PI * (yTile)) / tiles;
+        final double minY = Math.toDegrees(ESSloppyMath.atan(ESSloppyMath.sinh(minN)));
+        final double minX = ((xTile) / tiles * 360.0) - 180;
+        final double maxY = Math.toDegrees(ESSloppyMath.atan(ESSloppyMath.sinh(maxN)));
+        final double maxX = ((xTile + 1) / tiles * 360.0) - 180;
+
+        return new Rectangle(minX, maxX, maxY, minY);
+    }
+
     /**
      * Validates Zoom, X, and Y values, and returns the total number of allowed tiles along the x/y axis.
      */

server/src/test/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoTileUtilsTests.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
 import org.elasticsearch.common.geo.GeoPoint;
+import org.elasticsearch.geometry.Rectangle;
 import org.elasticsearch.test.ESTestCase;
 
 import static org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils.MAX_ZOOM;
@@ -28,8 +29,10 @@
 import static org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils.keyToGeoPoint;
 import static org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils.longEncode;
 import static org.elasticsearch.search.aggregations.bucket.geogrid.GeoTileUtils.stringEncode;
+import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
 
 public class GeoTileUtilsTests extends ESTestCase {
 
@@ -219,8 +222,8 @@ public void testGeoTileAsLongRoutines() {
      * so ensure they are clipped correctly.
      */
     public void testSingularityAtPoles() {
-        double minLat = -85.05112878;
-        double maxLat = 85.05112878;
+        double minLat = -GeoTileUtils.LATITUDE_MASK;
+        double maxLat = GeoTileUtils.LATITUDE_MASK;
         double lon = randomIntBetween(-180, 180);
         double lat = randomBoolean()
             ? randomDoubleBetween(-90, minLat, true)
@@ -231,4 +234,23 @@ public void testSingularityAtPoles() {
         String clippedTileIndex = stringEncode(longEncode(lon, clippedLat, zoom));
         assertEquals(tileIndex, clippedTileIndex);
     }
+
+    public void testPointToTile() {
+        int zoom = randomIntBetween(0, MAX_ZOOM);
+        int tiles = 1 << zoom;
+        int xTile = randomIntBetween(0, zoom);
+        int yTile = randomIntBetween(0, zoom);
+        Rectangle rectangle = GeoTileUtils.toBoundingBox(xTile, yTile, zoom);
+        // check corners
+        assertThat(GeoTileUtils.getXTile(rectangle.getMinX(), tiles), equalTo(xTile));
+        assertThat(GeoTileUtils.getXTile(rectangle.getMaxX(), tiles), equalTo(Math.min(tiles - 1, xTile + 1)));
+        assertThat(GeoTileUtils.getYTile(rectangle.getMaxY(), tiles), anyOf(equalTo(yTile - 1), equalTo(yTile)));
+        assertThat(GeoTileUtils.getYTile(rectangle.getMinY(), tiles), anyOf(equalTo(yTile + 1), equalTo(yTile)));
+        // check point inside
+        double x = randomDoubleBetween(rectangle.getMinX(), rectangle.getMaxX(), false);
+        double y = randomDoubleBetween(rectangle.getMinY(), rectangle.getMaxY(), false);
+        assertThat(GeoTileUtils.getXTile(x, tiles), equalTo(xTile));
+        assertThat(GeoTileUtils.getYTile(y, tiles), equalTo(yTile));
+
+    }
 }

x-pack/plugin/core/src/main/java/org/elasticsearch/license/XPackLicenseState.java

@@ -89,6 +89,8 @@ public enum Feature {
 
         SPATIAL_GEO_CENTROID(OperationMode.GOLD, true),
 
+        SPATIAL_GEO_GRID(OperationMode.GOLD, true),
+
         ANALYTICS(OperationMode.MISSING, true);
 
         final OperationMode minimumOperationMode;

x-pack/plugin/spatial/build.gradle

@@ -27,6 +27,7 @@ restResources {
 
 testClusters.integTest {
   setting 'xpack.license.self_generated.type', 'trial'
+  setting 'indices.breaker.request.limit', '25kb'
   testDistribution = 'DEFAULT'
 }