Add support for distance queries on shape queries

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

@@ -6,8 +6,6 @@
 package org.elasticsearch.xpack.spatial.index.mapper;
 
 import org.apache.lucene.document.XYShape;
-import org.apache.lucene.geo.XYLine;
-import org.apache.lucene.geo.XYPolygon;
 import org.apache.lucene.index.IndexableField;
 import org.elasticsearch.geometry.Circle;
 import org.elasticsearch.geometry.Geometry;
@@ -75,8 +73,7 @@ public Void visit(GeometryCollection<?> collection) {
 
         @Override
         public Void visit(Line line) {
-            float[][] vertices = lineToFloatArray(line.getX(), line.getY());
-            addFields(XYShape.createIndexableFields(name, new XYLine(vertices[0], vertices[1])));
+            addFields(XYShape.createIndexableFields(name, ShapeUtils.toLuceneXYLine(line)));
             return null;
         }
 
@@ -111,50 +108,24 @@ public Void visit(MultiPolygon multiPolygon) {
 
         @Override
         public Void visit(Point point) {
-            addFields(XYShape.createIndexableFields(name, (float)point.getX(), (float)point.getY()));
+            addFields(XYShape.createIndexableFields(name, (float) point.getX(), (float) point.getY()));
             return null;
         }
 
         @Override
         public Void visit(Polygon polygon) {
-            addFields(XYShape.createIndexableFields(name, toLucenePolygon(polygon)));
+            addFields(XYShape.createIndexableFields(name, ShapeUtils.toLuceneXYPolygon(polygon)));
             return null;
         }
 
         @Override
         public Void visit(Rectangle r) {
-            XYPolygon p = new XYPolygon(
-                new float[]{(float)r.getMinX(), (float)r.getMaxX(), (float)r.getMaxX(), (float)r.getMinX(), (float)r.getMinX()},
-                new float[]{(float)r.getMinY(), (float)r.getMinY(), (float)r.getMaxY(), (float)r.getMaxY(), (float)r.getMinY()});
-            addFields(XYShape.createIndexableFields(name, p));
+            addFields(XYShape.createIndexableFields(name, ShapeUtils.toLuceneXYPolygon(r)));
             return null;
         }
 
         private void addFields(IndexableField[] fields) {
             this.fields.addAll(Arrays.asList(fields));
         }
     }
-
-    public static XYPolygon toLucenePolygon(Polygon polygon) {
-        XYPolygon[] holes = new XYPolygon[polygon.getNumberOfHoles()];
-        LinearRing ring;
-        float[][] vertices;
-        for(int i = 0; i<holes.length; i++) {
-            ring = polygon.getHole(i);
-            vertices = lineToFloatArray(ring.getX(), ring.getY());
-            holes[i] = new XYPolygon(vertices[0], vertices[1]);
-        }
-        ring = polygon.getPolygon();
-        vertices = lineToFloatArray(ring.getX(), ring.getY());
-        return new XYPolygon(vertices[0], vertices[1], holes);
-    }
-
-    private static float[][] lineToFloatArray(double[] x, double[] y) {
-        float[][] result = new float[2][x.length];
-        for (int i = 0; i < x.length; ++i) {
-            result[0][i] = (float)x[i];
-            result[1][i] = (float)y[i];
-        }
-        return result;
-    }
 }

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

@@ -0,0 +1,68 @@
+/*
+ * 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.index.mapper;
+
+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 ShapeUtils {
+
+    public static org.apache.lucene.geo.XYPolygon toLuceneXYPolygon(Polygon polygon) {
+        org.apache.lucene.geo.XYPolygon[] holes = new org.apache.lucene.geo.XYPolygon[polygon.getNumberOfHoles()];
+        for(int i = 0; i<holes.length; i++) {
+            holes[i] = new org.apache.lucene.geo.XYPolygon(
+                doubleArrayToFloatArray(polygon.getHole(i).getX()),
+                doubleArrayToFloatArray(polygon.getHole(i).getY()));
+        }
+        return new org.apache.lucene.geo.XYPolygon(
+            doubleArrayToFloatArray(polygon.getPolygon().getX()),
+            doubleArrayToFloatArray(polygon.getPolygon().getY()), holes);
+    }
+
+    public static org.apache.lucene.geo.XYPolygon toLuceneXYPolygon(Rectangle r) {
+        return new org.apache.lucene.geo.XYPolygon(
+            new float[]{(float) r.getMinX(), (float) r.getMaxX(), (float) r.getMaxX(), (float) r.getMinX(), (float) r.getMinX()},
+            new float[]{(float) r.getMinY(), (float) r.getMinY(), (float) r.getMaxY(), (float) r.getMaxY(), (float) r.getMinY()});
+    }
+
+    public static org.apache.lucene.geo.XYRectangle toLuceneXYRectangle(Rectangle r) {
+        return new org.apache.lucene.geo.XYRectangle((float) r.getMinX(), (float) r.getMaxX(),
+                                                     (float) r.getMinY(), (float) r.getMaxY());
+    }
+
+    public static org.apache.lucene.geo.XYPoint toLuceneXYPoint(Point point) {
+        return new org.apache.lucene.geo.XYPoint((float) point.getX(), (float) point.getY());
+    }
+
+    public static org.apache.lucene.geo.XYLine toLuceneXYLine(Line line) {
+        return new org.apache.lucene.geo.XYLine(
+            doubleArrayToFloatArray(line.getX()),
+            doubleArrayToFloatArray(line.getY()));
+    }
+
+    public static org.apache.lucene.geo.XYCircle toLuceneXYCircle(Circle circle) {
+        return new org.apache.lucene.geo.XYCircle((float) circle.getX(), (float) circle.getY(), (float) circle.getRadiusMeters());
+    }
+
+    private ShapeUtils() {
+    }
+
+    private static float[] doubleArrayToFloatArray(double[] array) {
+        float[] result = new float[array.length];
+        for (int i = 0; i < array.length; ++i) {
+            result[i] = (float) array[i];
+        }
+        return result;
+    }
+
+}

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/ShapeQueryBuilder.java

@@ -5,6 +5,7 @@
  */
 package org.elasticsearch.xpack.spatial.index.query;
 
+import org.apache.lucene.search.ConstantScoreQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.geo.builders.ShapeBuilder;
 import org.elasticsearch.common.geo.parsers.ShapeParser;
@@ -121,7 +122,7 @@ public Query buildShapeQuery(QueryShardContext context, MappedFieldType fieldTyp
         }
 
         final AbstractGeometryFieldMapper.AbstractGeometryFieldType ft = (AbstractGeometryFieldMapper.AbstractGeometryFieldType) fieldType;
-        return ft.geometryQueryBuilder().process(shape, ft.name(), relation, context);
+        return  new ConstantScoreQuery(ft.geometryQueryBuilder().process(shape, ft.name(), relation, context));
     }
 
     @Override

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/query/ShapeQueryProcessor.java

@@ -5,13 +5,8 @@
  */
 package org.elasticsearch.xpack.spatial.index.query;
 
-import org.apache.lucene.document.ShapeField;
 import org.apache.lucene.document.XYShape;
-import org.apache.lucene.geo.XYLine;
-import org.apache.lucene.geo.XYPolygon;
-import org.apache.lucene.search.BooleanClause;
-import org.apache.lucene.search.BooleanQuery;
-import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.geo.XYGeometry;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.Version;
@@ -33,24 +28,26 @@
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.xpack.spatial.index.mapper.ShapeFieldMapper;
+import org.elasticsearch.xpack.spatial.index.mapper.ShapeUtils;
+
+import java.util.ArrayList;
+import java.util.List;
 
-import static org.elasticsearch.xpack.spatial.index.mapper.ShapeIndexer.toLucenePolygon;
 
 public class ShapeQueryProcessor implements AbstractSearchableGeometryFieldType.QueryProcessor {
 
     @Override
     public Query process(Geometry shape, String fieldName, ShapeRelation relation, QueryShardContext context) {
         validateIsShapeFieldType(fieldName, context);
-        if (shape == null) {
-            return new MatchNoDocsQuery();
-        }
         // CONTAINS queries are not supported by VECTOR strategy for indices created before version 7.5.0 (Lucene 8.3.0);
         if (relation == ShapeRelation.CONTAINS && context.indexVersionCreated().before(Version.V_7_5_0)) {
             throw new QueryShardException(context,
                 ShapeRelation.CONTAINS + " query relation not supported for Field [" + fieldName + "].");
         }
-        // wrap geometry Query as a ConstantScoreQuery
-        return new ConstantScoreQuery(shape.visit(new ShapeVisitor(context, fieldName, relation)));
+        if (shape == null) {
+            return new MatchNoDocsQuery();
+        }
+        return getVectorQueryFromShape(shape, fieldName, relation, context);
     }
 
     private void validateIsShapeFieldType(String fieldName, QueryShardContext context) {
@@ -61,115 +58,107 @@ private void validateIsShapeFieldType(String fieldName, QueryShardContext contex
         }
     }
 
-    private class ShapeVisitor implements GeometryVisitor<Query, RuntimeException> {
-        QueryShardContext context;
-        String fieldName;
-        ShapeRelation relation;
+    private Query getVectorQueryFromShape(Geometry queryShape, String fieldName, ShapeRelation relation, QueryShardContext context) {
+        final LuceneGeometryCollector visitor = new LuceneGeometryCollector(fieldName, context);
+        queryShape.visit(visitor);
+        final List<XYGeometry> geometries = visitor.geometries();
+        if (geometries.size() == 0) {
+            return new MatchNoDocsQuery();
+        }
+        return XYShape.newGeometryQuery(fieldName, relation.getLuceneRelation(),
+            geometries.toArray(new XYGeometry[geometries.size()]));
+    }
+
+    private static class LuceneGeometryCollector implements GeometryVisitor<Void, RuntimeException> {
+        private final List<XYGeometry> 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.fieldName = fieldName;
-            this.relation = relation;
         }
 
-        @Override
-        public Query visit(Circle circle) {
-            throw new QueryShardException(context, "Field [" + fieldName + "] found and unknown shape Circle");
+        List<XYGeometry> 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(ShapeUtils.toLuceneXYCircle(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(Line line) {
-            return XYShape.newLineQuery(fieldName, relation.getLuceneRelation(),
-                new XYLine(doubleArrayToFloatArray(line.getX()), doubleArrayToFloatArray(line.getY())));
+        public Void visit(Line line) {
+            if (line.isEmpty() == false) {
+                geometries.add(ShapeUtils.toLuceneXYLine(line));
+            }
+            return null;
         }
 
         @Override
-        public Query visit(LinearRing ring) {
-            throw new QueryShardException(context, "Field [" + fieldName + "] found and 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) {
-            XYLine[] lines = new XYLine[multiLine.size()];
-            for (int i=0; i<multiLine.size(); i++) {
-                lines[i] = new XYLine(doubleArrayToFloatArray(multiLine.get(i).getX()),
-                    doubleArrayToFloatArray(multiLine.get(i).getY()));
+        public Void visit(MultiLine multiLine) {
+            for (Line line : multiLine) {
+                visit(line);
             }
-            return XYShape.newLineQuery(fieldName, relation.getLuceneRelation(), lines);
+            return null;
         }
 
         @Override
-        public Query visit(MultiPoint multiPoint) {
-            float[][] points = new float[multiPoint.size()][2];
-            for (int i = 0; i < multiPoint.size(); i++) {
-                points[i] = new float[] {(float) multiPoint.get(i).getX(), (float) multiPoint.get(i).getY()};
+        public Void visit(MultiPoint multiPoint) {
+            for (Point point : multiPoint) {
+                visit(point);
             }
-            return XYShape.newPointQuery(fieldName, relation.getLuceneRelation(), points);
+            return null;
         }
 
         @Override
-        public Query visit(MultiPolygon multiPolygon) {
-            XYPolygon[] polygons = new XYPolygon[multiPolygon.size()];
-            for (int i=0; i<multiPolygon.size(); i++) {
-                polygons[i] = toLucenePolygon(multiPolygon.get(i));
+        public Void visit(MultiPolygon multiPolygon) {
+            for (Polygon polygon : multiPolygon) {
+                visit(polygon);
             }
-            return visitMultiPolygon(polygons);
-        }
-
-        private Query visitMultiPolygon(XYPolygon... polygons) {
-            return XYShape.newPolygonQuery(fieldName, relation.getLuceneRelation(), polygons);
+            return null;
         }
 
         @Override
-        public Query visit(Point point) {
-            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(ShapeUtils.toLuceneXYPoint(point));
             }
-            float[][] pointArray  = new float[][] {{(float)point.getX(), (float)point.getY()}};
-            return XYShape.newPointQuery(fieldName, luceneRelation, pointArray);
-        }
+            return null;
 
-        @Override
-        public Query visit(Polygon polygon) {
-            return XYShape.newPolygonQuery(fieldName, relation.getLuceneRelation(), toLucenePolygon(polygon));
         }
 
         @Override
-        public Query visit(Rectangle r) {
-            return XYShape.newBoxQuery(fieldName, relation.getLuceneRelation(),
-                (float)r.getMinX(), (float)r.getMaxX(), (float)r.getMinY(), (float)r.getMaxY());
+        public Void visit(Polygon polygon) {
+            if (polygon.isEmpty() == false) {
+                geometries.add(ShapeUtils.toLuceneXYPolygon(polygon));
+            }
+            return null;
         }
-    }
 
-    private static float[] doubleArrayToFloatArray(double[] array) {
-        float[] result = new float[array.length];
-        for (int i = 0; i < array.length; ++i) {
-            result[i] = (float) array[i];
+        @Override
+        public Void visit(Rectangle r) {
+            if (r.isEmpty() == false) {
+                geometries.add(ShapeUtils.toLuceneXYRectangle(r));
+            }
+            return null;
         }
-        return result;
     }
 }