Skip to content

Commit de0da5f

Browse files
imotovimotov
committed
Add support for geometrycollections to GeometryTreeReader and Writer
Add support for multi shapes and geometry collections to GeometryTreeReader and GeometryTreeWriter. Relates elastic#37206
1 parent cbc691f commit de0da5f

File tree

4 files changed

+91
-27
lines changed

4 files changed

+91
-27
lines changed

‎server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java

+27-8
Original file line numberDiff line numberDiff line change
@@ -34,42 +34,50 @@
3434
*/
3535
public class GeometryTreeReader implements ShapeTreeReader {
3636

37-
private final int extentOffset = 8;
37+
private static final int EXTENT_OFFSET = 8;
38+
private final int extentPosition;
3839
private final ByteBufferStreamInput input;
3940
private final CoordinateEncoder coordinateEncoder;
4041

4142
public GeometryTreeReader(BytesRef bytesRef, CoordinateEncoder coordinateEncoder) {
4243
this.input = new ByteBufferStreamInput(ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length));
44+
extentPosition = 0;
45+
this.coordinateEncoder = coordinateEncoder;
46+
}
47+
48+
private GeometryTreeReader(ByteBufferStreamInput input, CoordinateEncoder coordinateEncoder) throws IOException {
49+
this.input = input;
50+
extentPosition = input.position();
4351
this.coordinateEncoder = coordinateEncoder;
4452
}
4553

4654
public double getCentroidX() throws IOException {
47-
input.position(0);
55+
input.position(extentPosition);
4856
return coordinateEncoder.decodeX(input.readInt());
4957
}
5058

5159
public double getCentroidY() throws IOException {
52-
input.position(4);
60+
input.position(extentPosition + 4);
5361
return coordinateEncoder.decodeY(input.readInt());
5462
}
5563

5664
@Override
5765
public Extent getExtent() throws IOException {
58-
input.position(extentOffset);
66+
input.position(extentPosition + EXTENT_OFFSET);
5967
Extent extent = input.readOptionalWriteable(Extent::new);
6068
if (extent != null) {
6169
return extent;
6270
}
6371
assert input.readVInt() == 1;
6472
ShapeType shapeType = input.readEnum(ShapeType.class);
65-
ShapeTreeReader reader = getReader(shapeType, input);
73+
ShapeTreeReader reader = getReader(shapeType, coordinateEncoder, input);
6674
return reader.getExtent();
6775
}
6876

6977
@Override
7078
public GeoRelation relate(Extent extent) throws IOException {
7179
GeoRelation relation = GeoRelation.QUERY_DISJOINT;
72-
input.position(extentOffset);
80+
input.position(extentPosition + EXTENT_OFFSET);
7381
boolean hasExtent = input.readBoolean();
7482
if (hasExtent) {
7583
Optional<Boolean> extentCheck = EdgeTreeReader.checkExtent(new Extent(input), extent);
@@ -79,9 +87,17 @@ public GeoRelation relate(Extent extent) throws IOException {
7987
}
8088

8189
int numTrees = input.readVInt();
90+
int nextPosition = input.position();
8291
for (int i = 0; i < numTrees; i++) {
92+
if (numTrees > 1) {
93+
if (i > 0) {
94+
input.position(nextPosition);
95+
}
96+
int pos = input.readVInt();
97+
nextPosition = input.position() + pos;
98+
}
8399
ShapeType shapeType = input.readEnum(ShapeType.class);
84-
ShapeTreeReader reader = getReader(shapeType, input);
100+
ShapeTreeReader reader = getReader(shapeType, coordinateEncoder, input);
85101
GeoRelation shapeRelation = reader.relate(extent);
86102
if (GeoRelation.QUERY_CROSSES == shapeRelation ||
87103
(GeoRelation.QUERY_DISJOINT == shapeRelation && GeoRelation.QUERY_INSIDE == relation)
@@ -95,7 +111,8 @@ public GeoRelation relate(Extent extent) throws IOException {
95111
return relation;
96112
}
97113

98-
private static ShapeTreeReader getReader(ShapeType shapeType, ByteBufferStreamInput input) throws IOException {
114+
private static ShapeTreeReader getReader(ShapeType shapeType, CoordinateEncoder coordinateEncoder, ByteBufferStreamInput input)
115+
throws IOException {
99116
switch (shapeType) {
100117
case POLYGON:
101118
return new PolygonTreeReader(input);
@@ -105,6 +122,8 @@ private static ShapeTreeReader getReader(ShapeType shapeType, ByteBufferStreamIn
105122
case LINESTRING:
106123
case MULTILINESTRING:
107124
return new EdgeTreeReader(input, false);
125+
case GEOMETRYCOLLECTION:
126+
return new GeometryTreeReader(input, coordinateEncoder);
108127
default:
109128
throw new UnsupportedOperationException("unsupported shape type [" + shapeType + "]");
110129
}

‎server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java

+38-12
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,9 @@
1818
*/
1919
package org.elasticsearch.common.geo;
2020

21+
import org.elasticsearch.common.bytes.BytesReference;
22+
import org.elasticsearch.common.io.stream.BytesStreamOutput;
2123
import org.elasticsearch.common.io.stream.StreamOutput;
22-
import org.elasticsearch.common.io.stream.Writeable;
2324
import org.elasticsearch.geometry.Circle;
2425
import org.elasticsearch.geometry.Geometry;
2526
import org.elasticsearch.geometry.GeometryCollection;
@@ -32,6 +33,7 @@
3233
import org.elasticsearch.geometry.Point;
3334
import org.elasticsearch.geometry.Polygon;
3435
import org.elasticsearch.geometry.Rectangle;
36+
import org.elasticsearch.geometry.ShapeType;
3537

3638
import java.io.IOException;
3739
import java.util.ArrayList;
@@ -43,7 +45,7 @@
4345
* appropriate tree structure for each type of
4446
* {@link Geometry} into a byte array.
4547
*/
46-
public class GeometryTreeWriter implements Writeable {
48+
public class GeometryTreeWriter extends ShapeTreeWriter {
4749

4850
private final GeometryTreeBuilder builder;
4951
private final CoordinateEncoder coordinateEncoder;
@@ -53,29 +55,55 @@ public GeometryTreeWriter(Geometry geometry, CoordinateEncoder coordinateEncoder
5355
this.coordinateEncoder = coordinateEncoder;
5456
this.centroidCalculator = new CentroidCalculator();
5557
builder = new GeometryTreeBuilder(coordinateEncoder);
56-
geometry.visit(builder);
58+
if (geometry.type() == ShapeType.GEOMETRYCOLLECTION) {
59+
for (Geometry shape : (GeometryCollection<?>) geometry) {
60+
shape.visit(builder);
61+
}
62+
} else {
63+
geometry.visit(builder);
64+
}
5765
}
5866

59-
public Extent extent() {
67+
@Override
68+
public Extent getExtent() {
6069
return new Extent(builder.top, builder.bottom, builder.negLeft, builder.negRight, builder.posLeft, builder.posRight);
6170
}
6271

72+
@Override
73+
public ShapeType getShapeType() {
74+
return ShapeType.GEOMETRYCOLLECTION;
75+
}
76+
77+
@Override
78+
public CentroidCalculator getCentroidCalculator() {
79+
return centroidCalculator;
80+
}
81+
6382
@Override
6483
public void writeTo(StreamOutput out) throws IOException {
6584
// only write a geometry extent for the tree if the tree
6685
// contains multiple sub-shapes
67-
boolean prependExtent = builder.shapeWriters.size() > 1;
86+
boolean multiShape = builder.shapeWriters.size() > 1;
6887
Extent extent = null;
6988
out.writeInt(coordinateEncoder.encodeX(centroidCalculator.getX()));
7089
out.writeInt(coordinateEncoder.encodeY(centroidCalculator.getY()));
71-
if (prependExtent) {
90+
if (multiShape) {
7291
extent = new Extent(builder.top, builder.bottom, builder.negLeft, builder.negRight, builder.posLeft, builder.posRight);
7392
}
7493
out.writeOptionalWriteable(extent);
7594
out.writeVInt(builder.shapeWriters.size());
76-
for (ShapeTreeWriter writer : builder.shapeWriters) {
77-
out.writeEnum(writer.getShapeType());
78-
writer.writeTo(out);
95+
if (multiShape) {
96+
for (ShapeTreeWriter writer : builder.shapeWriters) {
97+
try(BytesStreamOutput bytesStream = new BytesStreamOutput()) {
98+
bytesStream.writeEnum(writer.getShapeType());
99+
writer.writeTo(bytesStream);
100+
BytesReference bytes = bytesStream.bytes();
101+
out.writeBytesReference(bytes);
102+
}
103+
}
104+
} else {
105+
out.writeEnum(builder.shapeWriters.get(0).getShapeType());
106+
builder.shapeWriters.get(0).writeTo(out);
79107
}
80108
}
81109

@@ -110,9 +138,7 @@ private void addWriter(ShapeTreeWriter writer) {
110138

111139
@Override
112140
public Void visit(GeometryCollection<?> collection) {
113-
for (Geometry geometry : collection) {
114-
geometry.visit(this);
115-
}
141+
addWriter(new GeometryTreeWriter(collection, coordinateEncoder));
116142
return null;
117143
}
118144

‎server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -185,7 +185,7 @@ public static GeoShapeValue missing(String missing) {
185185
try {
186186
Geometry geometry = MISSING_GEOMETRY_PARSER.fromWKT(missing);
187187
GeometryTreeWriter writer = new GeometryTreeWriter(geometry, GeoShapeCoordinateEncoder.INSTANCE);
188-
return new GeoShapeValue(writer.extent());
188+
return new GeoShapeValue(writer.getExtent());
189189
} catch (IOException | ParseException e) {
190190
throw new IllegalArgumentException("Can't apply missing value [" + missing + "]", e);
191191
}

‎server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java

+25-6
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import org.elasticsearch.common.io.stream.BytesStreamOutput;
2424
import org.elasticsearch.geo.GeometryTestUtils;
2525
import org.elasticsearch.geometry.Geometry;
26+
import org.elasticsearch.geometry.GeometryCollection;
2627
import org.elasticsearch.geometry.Line;
2728
import org.elasticsearch.geometry.LinearRing;
2829
import org.elasticsearch.geometry.MultiLine;
@@ -31,14 +32,15 @@
3132
import org.elasticsearch.geometry.Point;
3233
import org.elasticsearch.geometry.Polygon;
3334
import org.elasticsearch.geometry.Rectangle;
35+
import org.elasticsearch.geometry.ShapeType;
3436
import org.elasticsearch.index.mapper.GeoShapeIndexer;
3537
import org.elasticsearch.index.query.LegacyGeoShapeQueryProcessor;
36-
import org.elasticsearch.geometry.ShapeType;
3738
import org.elasticsearch.test.ESTestCase;
3839
import org.elasticsearch.test.geo.RandomShapeGenerator;
3940

4041
import java.io.IOException;
4142
import java.nio.ByteBuffer;
43+
import java.util.ArrayList;
4244
import java.util.Arrays;
4345
import java.util.Collections;
4446
import java.util.List;
@@ -307,8 +309,10 @@ public void testRandomGeometryIntersection() throws IOException {
307309
Geometry preparedGeometry = indexer.prepareForIndexing(geometry);
308310

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

313317
for (int i = 0; i < testPointCount; i++) {
314318
int cur = i;
@@ -329,21 +333,36 @@ private Extent bufferedExtentFromGeoPoint(double x, double y, double extentSize)
329333
}
330334

331335
private boolean intersects(Geometry g, Point p, double extentSize) throws IOException {
332-
return geometryTreeReader(g, GeoShapeCoordinateEncoder.INSTANCE)
333-
.relate(bufferedExtentFromGeoPoint(p.getX(), p.getY(), extentSize)) == GeoRelation.QUERY_CROSSES;
336+
GeoRelation relation = geometryTreeReader(g, GeoShapeCoordinateEncoder.INSTANCE)
337+
.relate(bufferedExtentFromGeoPoint(p.getX(), p.getY(), extentSize));
338+
return relation == GeoRelation.QUERY_CROSSES || relation == GeoRelation.QUERY_INSIDE;
334339
}
335340

336341
private static Geometry randomGeometryTreeGeometry() {
342+
return randomGeometryTreeGeometry(0);
343+
}
344+
345+
private static Geometry randomGeometryTreeGeometry(int level) {
337346
@SuppressWarnings("unchecked") Function<Boolean, Geometry> geometry = ESTestCase.randomFrom(
338347
GeometryTestUtils::randomLine,
339348
GeometryTestUtils::randomPoint,
340349
GeometryTestUtils::randomPolygon,
341350
GeometryTestUtils::randomMultiLine,
342-
GeometryTestUtils::randomMultiPoint
351+
GeometryTestUtils::randomMultiPoint,
352+
level < 3 ? (b) -> randomGeometryTreeCollection(level + 1) : GeometryTestUtils::randomPoint // don't build too deep
343353
);
344354
return geometry.apply(false);
345355
}
346356

357+
private static Geometry randomGeometryTreeCollection(int level) {
358+
int size = ESTestCase.randomIntBetween(1, 10);
359+
List<Geometry> shapes = new ArrayList<>();
360+
for (int i = 0; i < size; i++) {
361+
shapes.add(randomGeometryTreeGeometry(level));
362+
}
363+
return new GeometryCollection<>(shapes);
364+
}
365+
347366
private GeometryTreeReader geometryTreeReader(Geometry geometry, CoordinateEncoder encoder) throws IOException {
348367
GeometryTreeWriter writer = new GeometryTreeWriter(geometry, encoder);
349368
BytesStreamOutput output = new BytesStreamOutput();

0 commit comments

Comments
 (0)