diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java index 5cb609be59f39..36c0710de1090 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java +++ b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java @@ -34,42 +34,50 @@ */ public class GeometryTreeReader implements ShapeTreeReader { - private final int extentOffset = 8; + private static final int EXTENT_OFFSET = 8; + private final int startPosition; private final ByteBufferStreamInput input; private final CoordinateEncoder coordinateEncoder; public GeometryTreeReader(BytesRef bytesRef, CoordinateEncoder coordinateEncoder) { this.input = new ByteBufferStreamInput(ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length)); + this.startPosition = 0; + this.coordinateEncoder = coordinateEncoder; + } + + private GeometryTreeReader(ByteBufferStreamInput input, CoordinateEncoder coordinateEncoder) throws IOException { + this.input = input; + startPosition = input.position(); this.coordinateEncoder = coordinateEncoder; } public double getCentroidX() throws IOException { - input.position(0); + input.position(startPosition); return coordinateEncoder.decodeX(input.readInt()); } public double getCentroidY() throws IOException { - input.position(4); + input.position(startPosition + 4); return coordinateEncoder.decodeY(input.readInt()); } @Override public Extent getExtent() throws IOException { - input.position(extentOffset); + input.position(startPosition + EXTENT_OFFSET); Extent extent = input.readOptionalWriteable(Extent::new); if (extent != null) { return extent; } assert input.readVInt() == 1; ShapeType shapeType = input.readEnum(ShapeType.class); - ShapeTreeReader reader = getReader(shapeType, input); + ShapeTreeReader reader = getReader(shapeType, coordinateEncoder, input); return reader.getExtent(); } @Override public GeoRelation relate(Extent extent) throws IOException { GeoRelation relation = GeoRelation.QUERY_DISJOINT; - input.position(extentOffset); + input.position(startPosition + EXTENT_OFFSET); boolean hasExtent = input.readBoolean(); if (hasExtent) { Optional extentCheck = EdgeTreeReader.checkExtent(new Extent(input), extent); @@ -79,9 +87,17 @@ public GeoRelation relate(Extent extent) throws IOException { } int numTrees = input.readVInt(); + int nextPosition = input.position(); for (int i = 0; i < numTrees; i++) { + if (numTrees > 1) { + if (i > 0) { + input.position(nextPosition); + } + int pos = input.readVInt(); + nextPosition = input.position() + pos; + } ShapeType shapeType = input.readEnum(ShapeType.class); - ShapeTreeReader reader = getReader(shapeType, input); + ShapeTreeReader reader = getReader(shapeType, coordinateEncoder, input); GeoRelation shapeRelation = reader.relate(extent); if (GeoRelation.QUERY_CROSSES == shapeRelation || (GeoRelation.QUERY_DISJOINT == shapeRelation && GeoRelation.QUERY_INSIDE == relation) @@ -95,7 +111,8 @@ public GeoRelation relate(Extent extent) throws IOException { return relation; } - private static ShapeTreeReader getReader(ShapeType shapeType, ByteBufferStreamInput input) throws IOException { + private static ShapeTreeReader getReader(ShapeType shapeType, CoordinateEncoder coordinateEncoder, ByteBufferStreamInput input) + throws IOException { switch (shapeType) { case POLYGON: return new PolygonTreeReader(input); @@ -105,6 +122,8 @@ private static ShapeTreeReader getReader(ShapeType shapeType, ByteBufferStreamIn case LINESTRING: case MULTILINESTRING: return new EdgeTreeReader(input, false); + case GEOMETRYCOLLECTION: + return new GeometryTreeReader(input, coordinateEncoder); default: throw new UnsupportedOperationException("unsupported shape type [" + shapeType + "]"); } diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java index 61bda1d8c44ed..08222031516a2 100644 --- a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java +++ b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java @@ -18,8 +18,9 @@ */ package org.elasticsearch.common.geo; +import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.io.stream.BytesStreamOutput; import org.elasticsearch.common.io.stream.StreamOutput; -import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.geometry.Circle; import org.elasticsearch.geometry.Geometry; import org.elasticsearch.geometry.GeometryCollection; @@ -32,6 +33,7 @@ import org.elasticsearch.geometry.Point; import org.elasticsearch.geometry.Polygon; import org.elasticsearch.geometry.Rectangle; +import org.elasticsearch.geometry.ShapeType; import java.io.IOException; import java.util.ArrayList; @@ -43,7 +45,7 @@ * appropriate tree structure for each type of * {@link Geometry} into a byte array. */ -public class GeometryTreeWriter implements Writeable { +public class GeometryTreeWriter extends ShapeTreeWriter { private final GeometryTreeBuilder builder; private final CoordinateEncoder coordinateEncoder; @@ -53,29 +55,56 @@ public GeometryTreeWriter(Geometry geometry, CoordinateEncoder coordinateEncoder this.coordinateEncoder = coordinateEncoder; this.centroidCalculator = new CentroidCalculator(); builder = new GeometryTreeBuilder(coordinateEncoder); - geometry.visit(builder); + if (geometry.type() == ShapeType.GEOMETRYCOLLECTION) { + for (Geometry shape : (GeometryCollection) geometry) { + shape.visit(builder); + } + } else { + geometry.visit(builder); + } } - public Extent extent() { + @Override + public Extent getExtent() { return new Extent(builder.top, builder.bottom, builder.negLeft, builder.negRight, builder.posLeft, builder.posRight); } + @Override + public ShapeType getShapeType() { + return ShapeType.GEOMETRYCOLLECTION; + } + + @Override + public CentroidCalculator getCentroidCalculator() { + return centroidCalculator; + } + @Override public void writeTo(StreamOutput out) throws IOException { // only write a geometry extent for the tree if the tree // contains multiple sub-shapes - boolean prependExtent = builder.shapeWriters.size() > 1; + boolean multiShape = builder.shapeWriters.size() > 1; Extent extent = null; out.writeInt(coordinateEncoder.encodeX(centroidCalculator.getX())); out.writeInt(coordinateEncoder.encodeY(centroidCalculator.getY())); - if (prependExtent) { + if (multiShape) { extent = new Extent(builder.top, builder.bottom, builder.negLeft, builder.negRight, builder.posLeft, builder.posRight); } out.writeOptionalWriteable(extent); out.writeVInt(builder.shapeWriters.size()); - for (ShapeTreeWriter writer : builder.shapeWriters) { - out.writeEnum(writer.getShapeType()); - writer.writeTo(out); + if (multiShape) { + for (ShapeTreeWriter writer : builder.shapeWriters) { + try(BytesStreamOutput bytesStream = new BytesStreamOutput()) { + bytesStream.writeEnum(writer.getShapeType()); + writer.writeTo(bytesStream); + BytesReference bytes = bytesStream.bytes(); + out.writeVInt(bytes.length()); + bytes.writeTo(out); + } + } + } else { + out.writeEnum(builder.shapeWriters.get(0).getShapeType()); + builder.shapeWriters.get(0).writeTo(out); } } @@ -110,9 +139,7 @@ private void addWriter(ShapeTreeWriter writer) { @Override public Void visit(GeometryCollection collection) { - for (Geometry geometry : collection) { - geometry.visit(this); - } + addWriter(new GeometryTreeWriter(collection, coordinateEncoder)); return null; } diff --git a/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java b/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java index 387555b34656c..d18865ee656b5 100644 --- a/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java +++ b/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java @@ -185,7 +185,7 @@ public static GeoShapeValue missing(String missing) { try { Geometry geometry = MISSING_GEOMETRY_PARSER.fromWKT(missing); GeometryTreeWriter writer = new GeometryTreeWriter(geometry, GeoShapeCoordinateEncoder.INSTANCE); - return new GeoShapeValue(writer.extent()); + return new GeoShapeValue(writer.getExtent()); } catch (IOException | ParseException e) { throw new IllegalArgumentException("Can't apply missing value [" + missing + "]", e); } diff --git a/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java b/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java index 02c1f031ad41e..c9d18f98f2ca0 100644 --- a/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java +++ b/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java @@ -23,6 +23,7 @@ import org.elasticsearch.common.io.stream.BytesStreamOutput; import org.elasticsearch.geo.GeometryTestUtils; import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.geometry.GeometryCollection; import org.elasticsearch.geometry.Line; import org.elasticsearch.geometry.LinearRing; import org.elasticsearch.geometry.MultiLine; @@ -33,12 +34,12 @@ import org.elasticsearch.geometry.Rectangle; import org.elasticsearch.index.mapper.GeoShapeIndexer; import org.elasticsearch.index.query.LegacyGeoShapeQueryProcessor; -import org.elasticsearch.geometry.ShapeType; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.geo.RandomShapeGenerator; import java.io.IOException; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -306,10 +307,6 @@ public void testRandomGeometryIntersection() throws IOException { GeoShapeIndexer indexer = new GeoShapeIndexer(true, "test"); Geometry preparedGeometry = indexer.prepareForIndexing(geometry); - // TODO: support multi-polygons - assumeFalse("polygon crosses dateline", - ShapeType.POLYGON == geometry.type() && ShapeType.MULTIPOLYGON == preparedGeometry.type()); - for (int i = 0; i < testPointCount; i++) { int cur = i; intersects[cur] = fold(preparedGeometry, false, (g, s) -> s || intersects(g, testPoints[cur], extentSize)); @@ -329,21 +326,36 @@ private Extent bufferedExtentFromGeoPoint(double x, double y, double extentSize) } private boolean intersects(Geometry g, Point p, double extentSize) throws IOException { - return geometryTreeReader(g, GeoShapeCoordinateEncoder.INSTANCE) - .relate(bufferedExtentFromGeoPoint(p.getX(), p.getY(), extentSize)) == GeoRelation.QUERY_CROSSES; + GeoRelation relation = geometryTreeReader(g, GeoShapeCoordinateEncoder.INSTANCE) + .relate(bufferedExtentFromGeoPoint(p.getX(), p.getY(), extentSize)); + return relation == GeoRelation.QUERY_CROSSES || relation == GeoRelation.QUERY_INSIDE; } private static Geometry randomGeometryTreeGeometry() { + return randomGeometryTreeGeometry(0); + } + + private static Geometry randomGeometryTreeGeometry(int level) { @SuppressWarnings("unchecked") Function geometry = ESTestCase.randomFrom( GeometryTestUtils::randomLine, GeometryTestUtils::randomPoint, GeometryTestUtils::randomPolygon, GeometryTestUtils::randomMultiLine, - GeometryTestUtils::randomMultiPoint + GeometryTestUtils::randomMultiPoint, + level < 3 ? (b) -> randomGeometryTreeCollection(level + 1) : GeometryTestUtils::randomPoint // don't build too deep ); return geometry.apply(false); } + private static Geometry randomGeometryTreeCollection(int level) { + int size = ESTestCase.randomIntBetween(1, 10); + List shapes = new ArrayList<>(); + for (int i = 0; i < size; i++) { + shapes.add(randomGeometryTreeGeometry(level)); + } + return new GeometryCollection<>(shapes); + } + private GeometryTreeReader geometryTreeReader(Geometry geometry, CoordinateEncoder encoder) throws IOException { GeometryTreeWriter writer = new GeometryTreeWriter(geometry, encoder); BytesStreamOutput output = new BytesStreamOutput();