Skip to content

Commit

Permalink
Add support for geometrycollections to GeometryTreeReader and Writer
Browse files Browse the repository at this point in the history 
Add support for multi shapes and geometry collections to
GeometryTreeReader and GeometryTreeWriter.

Relates elastic#37206
  • Loading branch information
@imotov
imotov committed Nov 26, 2019
1 parent cbc691f commit de0da5f
Show file tree
Hide file tree
Showing 4 changed files with 91 additions and 27 deletions.
35 changes: 27 additions & 8 deletions server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,42 +34,50 @@
*/
public class GeometryTreeReader implements ShapeTreeReader {

private final int extentOffset = 8;
private static final int EXTENT_OFFSET = 8;
private final int extentPosition;
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));
extentPosition = 0;
this.coordinateEncoder = coordinateEncoder;
}

private GeometryTreeReader(ByteBufferStreamInput input, CoordinateEncoder coordinateEncoder) throws IOException {
this.input = input;
extentPosition = input.position();
this.coordinateEncoder = coordinateEncoder;
}

public double getCentroidX() throws IOException {
input.position(0);
input.position(extentPosition);
return coordinateEncoder.decodeX(input.readInt());
}

public double getCentroidY() throws IOException {
input.position(4);
input.position(extentPosition + 4);
return coordinateEncoder.decodeY(input.readInt());
}

@Override
public Extent getExtent() throws IOException {
input.position(extentOffset);
input.position(extentPosition + 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(extentPosition + EXTENT_OFFSET);
boolean hasExtent = input.readBoolean();
if (hasExtent) {
Optional<Boolean> extentCheck = EdgeTreeReader.checkExtent(new Extent(input), extent);
Expand All @@ -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)
Expand All @@ -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);
Expand All @@ -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 + "]");
}
Expand Down
50 changes: 38 additions & 12 deletions server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -53,29 +55,55 @@ 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.writeBytesReference(bytes);
}
}
} else {
out.writeEnum(builder.shapeWriters.get(0).getShapeType());
builder.shapeWriters.get(0).writeTo(out);
}
}

Expand Down Expand Up @@ -110,9 +138,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;
}

Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
Expand Down
31 changes: 25 additions & 6 deletions server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -31,14 +32,15 @@
import org.elasticsearch.geometry.Point;
import org.elasticsearch.geometry.Polygon;
import org.elasticsearch.geometry.Rectangle;
import org.elasticsearch.geometry.ShapeType;
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;
Expand Down Expand Up @@ -307,8 +309,10 @@ public void testRandomGeometryIntersection() throws IOException {
Geometry preparedGeometry = indexer.prepareForIndexing(geometry);

// TODO: support multi-polygons
assumeFalse("polygon crosses dateline",
ShapeType.POLYGON == geometry.type() && ShapeType.MULTIPOLYGON == preparedGeometry.type());
assumeTrue("polygon crosses dateline",
fold(preparedGeometry, 0, (g, c) -> c + (ShapeType.POLYGON == g.type() ? 1 : 0)).equals(
fold(geometry, 0, (g, c) -> c + (ShapeType.POLYGON == g.type() ? 1 : 0)))
);

for (int i = 0; i < testPointCount; i++) {
int cur = i;
Expand All @@ -329,21 +333,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();
Expand Down

0 comments on commit de0da5f

Please sign in to comment.