Skip to content

Commit e1b12cb

Browse files
imotovimotov
committed
Add support for geometrycollections to GeometryTreeReader and Writer (#49608)
Adds support geometry collections to GeometryTreeReader and GeometryTreeWriter. Relates #37206
1 parent 100e119 commit e1b12cb

File tree

4 files changed

+87
-29
lines changed

4 files changed

+87
-29
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 startPosition;
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+
this.startPosition = 0;
45+
this.coordinateEncoder = coordinateEncoder;
46+
}
47+
48+
private GeometryTreeReader(ByteBufferStreamInput input, CoordinateEncoder coordinateEncoder) throws IOException {
49+
this.input = input;
50+
startPosition = input.position();
4351
this.coordinateEncoder = coordinateEncoder;
4452
}
4553

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

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

5664
@Override
5765
public Extent getExtent() throws IOException {
58-
input.position(extentOffset);
66+
input.position(startPosition + 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(startPosition + 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

+39-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,56 @@ 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.writeVInt(bytes.length());
102+
bytes.writeTo(out);
103+
}
104+
}
105+
} else {
106+
out.writeEnum(builder.shapeWriters.get(0).getShapeType());
107+
builder.shapeWriters.get(0).writeTo(out);
79108
}
80109
}
81110

@@ -110,9 +139,7 @@ private void addWriter(ShapeTreeWriter writer) {
110139

111140
@Override
112141
public Void visit(GeometryCollection<?> collection) {
113-
for (Geometry geometry : collection) {
114-
geometry.visit(this);
115-
}
142+
addWriter(new GeometryTreeWriter(collection, coordinateEncoder));
116143
return null;
117144
}
118145

‎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

+20-8
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;
@@ -33,12 +34,12 @@
3334
import org.elasticsearch.geometry.Rectangle;
3435
import org.elasticsearch.index.mapper.GeoShapeIndexer;
3536
import org.elasticsearch.index.query.LegacyGeoShapeQueryProcessor;
36-
import org.elasticsearch.geometry.ShapeType;
3737
import org.elasticsearch.test.ESTestCase;
3838
import org.elasticsearch.test.geo.RandomShapeGenerator;
3939

4040
import java.io.IOException;
4141
import java.nio.ByteBuffer;
42+
import java.util.ArrayList;
4243
import java.util.Arrays;
4344
import java.util.Collections;
4445
import java.util.List;
@@ -306,10 +307,6 @@ public void testRandomGeometryIntersection() throws IOException {
306307
GeoShapeIndexer indexer = new GeoShapeIndexer(true, "test");
307308
Geometry preparedGeometry = indexer.prepareForIndexing(geometry);
308309

309-
// TODO: support multi-polygons
310-
assumeFalse("polygon crosses dateline",
311-
ShapeType.POLYGON == geometry.type() && ShapeType.MULTIPOLYGON == preparedGeometry.type());
312-
313310
for (int i = 0; i < testPointCount; i++) {
314311
int cur = i;
315312
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)
329326
}
330327

331328
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;
329+
GeoRelation relation = geometryTreeReader(g, GeoShapeCoordinateEncoder.INSTANCE)
330+
.relate(bufferedExtentFromGeoPoint(p.getX(), p.getY(), extentSize));
331+
return relation == GeoRelation.QUERY_CROSSES || relation == GeoRelation.QUERY_INSIDE;
334332
}
335333

336334
private static Geometry randomGeometryTreeGeometry() {
335+
return randomGeometryTreeGeometry(0);
336+
}
337+
338+
private static Geometry randomGeometryTreeGeometry(int level) {
337339
@SuppressWarnings("unchecked") Function<Boolean, Geometry> geometry = ESTestCase.randomFrom(
338340
GeometryTestUtils::randomLine,
339341
GeometryTestUtils::randomPoint,
340342
GeometryTestUtils::randomPolygon,
341343
GeometryTestUtils::randomMultiLine,
342-
GeometryTestUtils::randomMultiPoint
344+
GeometryTestUtils::randomMultiPoint,
345+
level < 3 ? (b) -> randomGeometryTreeCollection(level + 1) : GeometryTestUtils::randomPoint // don't build too deep
343346
);
344347
return geometry.apply(false);
345348
}
346349

350+
private static Geometry randomGeometryTreeCollection(int level) {
351+
int size = ESTestCase.randomIntBetween(1, 10);
352+
List<Geometry> shapes = new ArrayList<>();
353+
for (int i = 0; i < size; i++) {
354+
shapes.add(randomGeometryTreeGeometry(level));
355+
}
356+
return new GeometryCollection<>(shapes);
357+
}
358+
347359
private GeometryTreeReader geometryTreeReader(Geometry geometry, CoordinateEncoder encoder) throws IOException {
348360
GeometryTreeWriter writer = new GeometryTreeWriter(geometry, encoder);
349361
BytesStreamOutput output = new BytesStreamOutput();

0 commit comments

Comments
 (0)