Add support for geometrycollections to GeometryTreeReader and Writer

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<Boolean> 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 + "]");
         }

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;
         }
 

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);
             }

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<Boolean, Geometry> 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<Geometry> 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();