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;
@@ -120,7 +121,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 AbstractGeometryFieldMapper.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;
  }
 }