Skip to content

Commit

Permalink
Add support for geometrycollections to GeometryTreeReader and Writer (#…
Browse files Browse the repository at this point in the history 
…49608)

Adds support geometry collections to GeometryTreeReader
and GeometryTreeWriter.

Relates #37206
  • Loading branch information
@imotov
imotov committed Nov 27, 2019
1 parent 100e119 commit e1b12cb
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 29 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 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);
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
51 changes: 39 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,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);
}
}

Expand Down Expand Up @@ -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;
}

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
28 changes: 20 additions & 8 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 @@ -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;
Expand Down Expand Up @@ -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));
Expand All @@ -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();
Expand Down

0 comments on commit e1b12cb

Please sign in to comment.