Add support for distance queries on geo_shape queries (#53466) (#53795)

iverase · web-flow · commit 4f1b2fd2b15f · 2020-03-19T15:21:58.000+01:00
With the upgrade to Lucene 8.5, LatLonShape field has support for distance queries. This change implements this new feature and removes the limitation.
diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeoShapeUtils.java b/server/src/main/java/org/elasticsearch/common/geo/GeoShapeUtils.java
@@ -0,0 +1,66 @@
+/*
+ * 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.common.geo;
+
+import org.elasticsearch.geometry.Circle;
+import org.elasticsearch.geometry.Line;
+import org.elasticsearch.geometry.Point;
+import org.elasticsearch.geometry.Polygon;
+import org.elasticsearch.geometry.Rectangle;
+
+
+/**
+ * Utility class that transforms Elasticsearch geometry objects to the Lucene representation
+ */
+public class GeoShapeUtils {
+
+    public static org.apache.lucene.geo.Polygon toLucenePolygon(Polygon polygon) {
+        org.apache.lucene.geo.Polygon[] holes = new org.apache.lucene.geo.Polygon[polygon.getNumberOfHoles()];
+        for(int i = 0; i<holes.length; i++) {
+            holes[i] = new org.apache.lucene.geo.Polygon(polygon.getHole(i).getY(), polygon.getHole(i).getX());
+        }
+        return new org.apache.lucene.geo.Polygon(polygon.getPolygon().getY(), polygon.getPolygon().getX(), holes);
+    }
+
+    public static org.apache.lucene.geo.Polygon toLucenePolygon(Rectangle r) {
+        return new org.apache.lucene.geo.Polygon(
+            new double[]{r.getMinLat(), r.getMinLat(), r.getMaxLat(), r.getMaxLat(), r.getMinLat()},
+            new double[]{r.getMinLon(), r.getMaxLon(), r.getMaxLon(), r.getMinLon(), r.getMinLon()});
+    }
+
+    public static org.apache.lucene.geo.Rectangle toLuceneRectangle(Rectangle r) {
+        return new org.apache.lucene.geo.Rectangle(r.getMinLat(), r.getMaxLat(), r.getMinLon(), r.getMaxLon());
+    }
+
+    public static org.apache.lucene.geo.Point toLucenePoint(Point point) {
+        return new org.apache.lucene.geo.Point(point.getLat(), point.getLon());
+    }
+
+    public static org.apache.lucene.geo.Line toLuceneLine(Line line) {
+        return new org.apache.lucene.geo.Line(line.getLats(), line.getLons());
+    }
+
+    public static org.apache.lucene.geo.Circle toLuceneCircle(Circle circle) {
+        return new org.apache.lucene.geo.Circle(circle.getLat(), circle.getLon(), circle.getRadiusMeters());
+    }
+
+    private GeoShapeUtils() {
+    }
+
+}
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/GeoShapeIndexer.java b/server/src/main/java/org/elasticsearch/index/mapper/GeoShapeIndexer.java
@@ -24,6 +24,7 @@
 import org.apache.lucene.index.IndexableField;
 import org.elasticsearch.common.geo.GeoLineDecomposer;
 import org.elasticsearch.common.geo.GeoPolygonDecomposer;
+import org.elasticsearch.common.geo.GeoShapeUtils;
 import org.elasticsearch.common.geo.GeoShapeType;
 import org.elasticsearch.geometry.Circle;
 import org.elasticsearch.geometry.Geometry;
@@ -213,7 +214,7 @@ public Void visit(GeometryCollection<?> collection) {
 
         @Override
         public Void visit(Line line) {
-            addFields(LatLonShape.createIndexableFields(name, new org.apache.lucene.geo.Line(line.getY(), line.getX())));
+            addFields(LatLonShape.createIndexableFields(name, GeoShapeUtils.toLuceneLine(line)));
             return null;
         }
 
@@ -254,16 +255,13 @@ public Void visit(Point point) {
 
         @Override
         public Void visit(Polygon polygon) {
-            addFields(LatLonShape.createIndexableFields(name, toLucenePolygon(polygon)));
+            addFields(LatLonShape.createIndexableFields(name, GeoShapeUtils.toLucenePolygon(polygon)));
             return null;
         }
 
         @Override
         public Void visit(Rectangle r) {
-            org.apache.lucene.geo.Polygon p = new org.apache.lucene.geo.Polygon(
-                new double[]{r.getMinY(), r.getMinY(), r.getMaxY(), r.getMaxY(), r.getMinY()},
-                new double[]{r.getMinX(), r.getMaxX(), r.getMaxX(), r.getMinX(), r.getMinX()});
-            addFields(LatLonShape.createIndexableFields(name, p));
+            addFields(LatLonShape.createIndexableFields(name, GeoShapeUtils.toLucenePolygon(r)));
             return null;
         }
 
@@ -272,11 +270,4 @@ private void addFields(IndexableField[] fields) {
         }
     }
 
-    public static org.apache.lucene.geo.Polygon toLucenePolygon(Polygon polygon) {
-        org.apache.lucene.geo.Polygon[] holes = new org.apache.lucene.geo.Polygon[polygon.getNumberOfHoles()];
-        for(int i = 0; i<holes.length; i++) {
-            holes[i] = new org.apache.lucene.geo.Polygon(polygon.getHole(i).getY(), polygon.getHole(i).getX());
-        }
-        return new org.apache.lucene.geo.Polygon(polygon.getPolygon().getY(), polygon.getPolygon().getX(), holes);
-    }
 }
diff --git a/server/src/main/java/org/elasticsearch/index/query/VectorGeoPointShapeQueryProcessor.java b/server/src/main/java/org/elasticsearch/index/query/VectorGeoPointShapeQueryProcessor.java
@@ -27,6 +27,7 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.geo.GeoPolygonDecomposer;
 import org.elasticsearch.common.geo.GeoShapeType;
+import org.elasticsearch.common.geo.GeoShapeUtils;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.geometry.Circle;
 import org.elasticsearch.geometry.Geometry;
@@ -46,8 +47,6 @@
 
 import java.util.ArrayList;
 
-import static org.elasticsearch.index.mapper.GeoShapeIndexer.toLucenePolygon;
-
 public class VectorGeoPointShapeQueryProcessor implements AbstractSearchableGeometryFieldType.QueryProcessor {
 
     @Override
@@ -145,7 +144,7 @@ private Query visit(ArrayList<Polygon> collector) {
             org.apache.lucene.geo.Polygon[] lucenePolygons =
                 new org.apache.lucene.geo.Polygon[collector.size()];
             for (int i = 0; i < collector.size(); i++) {
-                lucenePolygons[i] = toLucenePolygon(collector.get(i));
+                lucenePolygons[i] = GeoShapeUtils.toLucenePolygon(collector.get(i));
             }
             Query query = LatLonPoint.newPolygonQuery(fieldName, lucenePolygons);
             if (fieldType.hasDocValues()) {
diff --git a/server/src/main/java/org/elasticsearch/index/query/VectorGeoShapeQueryProcessor.java b/server/src/main/java/org/elasticsearch/index/query/VectorGeoShapeQueryProcessor.java
@@ -20,31 +20,32 @@
 package org.elasticsearch.index.query;
 
 import org.apache.lucene.document.LatLonShape;
-import org.apache.lucene.document.ShapeField;
-import org.apache.lucene.geo.Line;
-import org.apache.lucene.geo.Polygon;
-import org.apache.lucene.search.BooleanClause;
-import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.geo.LatLonGeometry;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.Version;
+import org.elasticsearch.common.geo.GeoLineDecomposer;
+import org.elasticsearch.common.geo.GeoPolygonDecomposer;
+import org.elasticsearch.common.geo.GeoShapeUtils;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.geometry.Circle;
 import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.geometry.GeometryCollection;
 import org.elasticsearch.geometry.GeometryVisitor;
+import org.elasticsearch.geometry.Line;
 import org.elasticsearch.geometry.LinearRing;
 import org.elasticsearch.geometry.MultiLine;
 import org.elasticsearch.geometry.MultiPoint;
 import org.elasticsearch.geometry.MultiPolygon;
 import org.elasticsearch.geometry.Point;
+import org.elasticsearch.geometry.Polygon;
 import org.elasticsearch.geometry.Rectangle;
 import org.elasticsearch.index.mapper.AbstractSearchableGeometryFieldType;
-import org.elasticsearch.index.mapper.GeoShapeFieldMapper;
-import org.elasticsearch.index.mapper.GeoShapeIndexer;
-import org.elasticsearch.index.mapper.MappedFieldType;
 
-import static org.elasticsearch.index.mapper.GeoShapeIndexer.toLucenePolygon;
+
+import java.util.ArrayList;
+import java.util.List;
+
 
 public class VectorGeoShapeQueryProcessor implements AbstractSearchableGeometryFieldType.QueryProcessor {
 
@@ -59,127 +60,126 @@ public Query process(Geometry shape, String fieldName, ShapeRelation relation, Q
         return getVectorQueryFromShape(shape, fieldName, relation, context);
     }
 
-    protected Query getVectorQueryFromShape(
-        Geometry queryShape, String fieldName, ShapeRelation relation, QueryShardContext context) {
-        GeoShapeIndexer geometryIndexer = new GeoShapeIndexer(true, fieldName);
-
-        Geometry processedShape = geometryIndexer.prepareForIndexing(queryShape);
-
-        if (processedShape == null) {
+    private Query getVectorQueryFromShape(Geometry queryShape, String fieldName, ShapeRelation relation, QueryShardContext context) {
+        final LuceneGeometryCollector visitor = new LuceneGeometryCollector(fieldName, context);
+        queryShape.visit(visitor);
+        final List<LatLonGeometry> geometries = visitor.geometries();
+        if (geometries.size() == 0) {
             return new MatchNoDocsQuery();
         }
-        return processedShape.visit(new ShapeVisitor(context, fieldName, relation));
+        return LatLonShape.newGeometryQuery(fieldName, relation.getLuceneRelation(),
+            geometries.toArray(new LatLonGeometry[geometries.size()]));
     }
 
-    private class ShapeVisitor implements GeometryVisitor<Query, RuntimeException> {
-        QueryShardContext context;
-        MappedFieldType fieldType;
-        String fieldName;
-        ShapeRelation relation;
+    private static class LuceneGeometryCollector implements GeometryVisitor<Void, RuntimeException> {
+        private final List<LatLonGeometry> geometries = new ArrayList<>();
+        private final String name;
+        private final QueryShardContext context;
 
-        ShapeVisitor(QueryShardContext context, String fieldName, ShapeRelation relation) {
+        private LuceneGeometryCollector(String name, QueryShardContext context) {
+            this.name = name;
             this.context = context;
-            this.fieldType = context.fieldMapper(fieldName);
-            this.fieldName = fieldName;
-            this.relation = relation;
         }
 
-        @Override
-        public Query visit(Circle circle) {
-            throw new QueryShardException(context, "Field [" + fieldName + "] found an unknown shape Circle");
+        List<LatLonGeometry> geometries() {
+            return geometries;
         }
 
         @Override
-        public Query visit(GeometryCollection<?> collection) {
-            BooleanQuery.Builder bqb = new BooleanQuery.Builder();
-            visit(bqb, collection);
-            return bqb.build();
-        }
-
-        private void visit(BooleanQuery.Builder bqb, GeometryCollection<?> collection) {
-            BooleanClause.Occur occur;
-            if (relation == ShapeRelation.CONTAINS || relation == ShapeRelation.DISJOINT) {
-                // all shapes must be disjoint / must be contained in relation to the indexed shape.
-                occur = BooleanClause.Occur.MUST;
-            } else {
-                // at least one shape must intersect / contain the indexed shape.
-                occur = BooleanClause.Occur.SHOULD;
+        public Void visit(Circle circle) {
+            if (circle.isEmpty() == false) {
+                geometries.add(GeoShapeUtils.toLuceneCircle(circle));
             }
+            return null;
+        }
+
+        @Override
+        public Void visit(GeometryCollection<?> collection) {
             for (Geometry shape : collection) {
-                bqb.add(shape.visit(this), occur);
+                shape.visit(this);
             }
+            return null;
         }
 
         @Override
-        public Query visit(org.elasticsearch.geometry.Line line) {
-            validateIsGeoShapeFieldType();
-            return LatLonShape.newLineQuery(fieldName, relation.getLuceneRelation(), new Line(line.getY(), line.getX()));
+        public Void visit(org.elasticsearch.geometry.Line line) {
+            if (line.isEmpty() == false) {
+                List<org.elasticsearch.geometry.Line> collector = new ArrayList<>();
+                GeoLineDecomposer.decomposeLine(line, collector);
+                collectLines(collector);
+            }
+            return null;
         }
 
         @Override
-        public Query visit(LinearRing ring) {
-            throw new QueryShardException(context, "Field [" + fieldName + "] found an unsupported shape LinearRing");
+        public Void visit(LinearRing ring) {
+            throw new QueryShardException(context, "Field [" + name + "] found and unsupported shape LinearRing");
         }
 
         @Override
-        public Query visit(MultiLine multiLine) {
-            validateIsGeoShapeFieldType();
-            Line[] lines = new Line[multiLine.size()];
-            for (int i = 0; i < multiLine.size(); i++) {
-                lines[i] = new Line(multiLine.get(i).getY(), multiLine.get(i).getX());
-            }
-            return LatLonShape.newLineQuery(fieldName, relation.getLuceneRelation(), lines);
+        public Void visit(MultiLine multiLine) {
+            List<org.elasticsearch.geometry.Line> collector = new ArrayList<>();
+            GeoLineDecomposer.decomposeMultiLine(multiLine, collector);
+            collectLines(collector);
+            return null;
         }
 
         @Override
-        public Query visit(MultiPoint multiPoint) {
-            double[][] points = new double[multiPoint.size()][2];
-            for (int i = 0; i < multiPoint.size(); i++) {
-                points[i] = new double[] {multiPoint.get(i).getLat(), multiPoint.get(i).getLon()};
+        public Void visit(MultiPoint multiPoint) {
+            for (Point point : multiPoint) {
+                visit(point);
             }
-            return LatLonShape.newPointQuery(fieldName, relation.getLuceneRelation(), points);
+            return null;
         }
 
         @Override
-        public Query visit(MultiPolygon multiPolygon) {
-            Polygon[] polygons = new Polygon[multiPolygon.size()];
-            for (int i = 0; i < multiPolygon.size(); i++) {
-                polygons[i] = toLucenePolygon(multiPolygon.get(i));
+        public Void visit(MultiPolygon multiPolygon) {
+            if (multiPolygon.isEmpty() == false) {
+                List<org.elasticsearch.geometry.Polygon> collector = new ArrayList<>();
+                GeoPolygonDecomposer.decomposeMultiPolygon(multiPolygon, true, collector);
+                collectPolygons(collector);
             }
-            return LatLonShape.newPolygonQuery(fieldName, relation.getLuceneRelation(), polygons);
+            return null;
         }
 
         @Override
-        public Query visit(Point point) {
-            validateIsGeoShapeFieldType();
-            ShapeField.QueryRelation luceneRelation = relation.getLuceneRelation();
-            if (luceneRelation == ShapeField.QueryRelation.CONTAINS) {
-                // contains and intersects are equivalent but the implementation of
-                // intersects is more efficient.
-                luceneRelation = ShapeField.QueryRelation.INTERSECTS;
+        public Void visit(Point point) {
+            if (point.isEmpty() == false) {
+                geometries.add(GeoShapeUtils.toLucenePoint(point));
             }
-            return LatLonShape.newPointQuery(fieldName, luceneRelation,
-                new double[] {point.getY(), point.getX()});
+            return null;
+
         }
 
         @Override
-        public Query visit(org.elasticsearch.geometry.Polygon polygon) {
-            return LatLonShape.newPolygonQuery(fieldName, relation.getLuceneRelation(), toLucenePolygon(polygon));
+        public Void visit(org.elasticsearch.geometry.Polygon polygon) {
+            if (polygon.isEmpty() == false) {
+                List<org.elasticsearch.geometry.Polygon> collector = new ArrayList<>();
+                GeoPolygonDecomposer.decomposePolygon(polygon, true, collector);
+                collectPolygons(collector);
+            }
+            return null;
         }
 
         @Override
-        public Query visit(Rectangle r) {
-            return LatLonShape.newBoxQuery(fieldName, relation.getLuceneRelation(),
-                r.getMinY(), r.getMaxY(), r.getMinX(), r.getMaxX());
+        public Void visit(Rectangle r) {
+            if (r.isEmpty() == false) {
+                geometries.add(GeoShapeUtils.toLuceneRectangle(r));
+            }
+            return null;
         }
 
-        private void validateIsGeoShapeFieldType() {
-            if (fieldType instanceof GeoShapeFieldMapper.GeoShapeFieldType == false) {
-                throw new QueryShardException(context, "Expected " + GeoShapeFieldMapper.CONTENT_TYPE
-                    + " field type for Field [" + fieldName + "] but found " + fieldType.typeName());
+        private void collectLines(List<org.elasticsearch.geometry.Line> geometryLines) {
+            for (Line line: geometryLines) {
+                geometries.add(GeoShapeUtils.toLuceneLine(line));
             }
         }
-    }
 
+        private void collectPolygons(List<org.elasticsearch.geometry.Polygon> geometryPolygons) {
+            for (Polygon polygon : geometryPolygons) {
+                geometries.add(GeoShapeUtils.toLucenePolygon(polygon));
+            }
+        }
+    }
 }
 
diff --git a/server/src/test/java/org/elasticsearch/search/geo/GeoShapeQueryTests.java b/server/src/test/java/org/elasticsearch/search/geo/GeoShapeQueryTests.java
