Skip to content

Commit 6b7b297

Browse files
iveraseiverase
authored
Simplify triangle tree writer (#64572)
This change delays the decoding of the indexed triangles to the tree writer.
1 parent 479cc16 commit 6b7b297

File tree

13 files changed

+120
-171
lines changed

13 files changed

+120
-171
lines changed

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeoShapeValues.java

Lines changed: 7 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,22 +6,17 @@
66

77
package org.elasticsearch.xpack.spatial.index.fielddata;
88

9-
import org.apache.lucene.document.ShapeField;
10-
import org.apache.lucene.index.IndexableField;
11-
import org.apache.lucene.store.ByteBuffersDataOutput;
12-
import org.apache.lucene.util.BytesRef;
139
import org.elasticsearch.geometry.Geometry;
1410
import org.elasticsearch.geometry.Rectangle;
1511
import org.elasticsearch.geometry.utils.GeographyValidator;
1612
import org.elasticsearch.geometry.utils.WellKnownText;
1713
import org.elasticsearch.index.mapper.GeoShapeIndexer;
1814
import org.elasticsearch.search.aggregations.support.ValuesSourceType;
15+
import org.elasticsearch.xpack.spatial.index.mapper.BinaryGeoShapeDocValuesField;
1916
import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType;
2017

2118
import java.io.IOException;
2219
import java.text.ParseException;
23-
import java.util.Arrays;
24-
import java.util.List;
2520

2621
/**
2722
* A stateful lightweight per document geo values.
@@ -136,35 +131,17 @@ public double lon() {
136131

137132
public static GeoShapeValue missing(String missing) {
138133
try {
139-
Geometry geometry = MISSING_GEOMETRY_PARSER.fromWKT(missing);
140-
ShapeField.DecodedTriangle[] triangles = toDecodedTriangles(geometry);
141-
GeometryDocValueWriter writer =
142-
new GeometryDocValueWriter(Arrays.asList(triangles), CoordinateEncoder.GEO,
143-
new CentroidCalculator(geometry));
144-
ByteBuffersDataOutput output = new ByteBuffersDataOutput();
145-
writer.writeTo(output);
146-
GeometryDocValueReader reader = new GeometryDocValueReader();
147-
reader.reset(new BytesRef(output.toArrayCopy(), 0, Math.toIntExact(output.size())));
134+
final GeoShapeIndexer indexer = new GeoShapeIndexer(true, "missing");
135+
final Geometry geometry = indexer.prepareForIndexing(MISSING_GEOMETRY_PARSER.fromWKT(missing));
136+
final BinaryGeoShapeDocValuesField field =
137+
new BinaryGeoShapeDocValuesField(missing, indexer.indexShape(null, geometry), new CentroidCalculator(geometry));
138+
final GeometryDocValueReader reader = new GeometryDocValueReader();
139+
reader.reset(field.binaryValue());
148140
return new GeoShapeValue(reader);
149141
} catch (IOException | ParseException e) {
150142
throw new IllegalArgumentException("Can't apply missing value [" + missing + "]", e);
151143
}
152144
}
153-
154-
private static ShapeField.DecodedTriangle[] toDecodedTriangles(Geometry geometry) {
155-
GeoShapeIndexer indexer = new GeoShapeIndexer(true, "test");
156-
geometry = indexer.prepareForIndexing(geometry);
157-
List<IndexableField> fields = indexer.indexShape(null, geometry);
158-
ShapeField.DecodedTriangle[] triangles = new ShapeField.DecodedTriangle[fields.size()];
159-
final byte[] scratch = new byte[7 * Integer.BYTES];
160-
for (int i = 0; i < fields.size(); i++) {
161-
BytesRef bytesRef = fields.get(i).binaryValue();
162-
assert bytesRef.length == 7 * Integer.BYTES;
163-
System.arraycopy(bytesRef.bytes, bytesRef.offset, scratch, 0, 7 * Integer.BYTES);
164-
ShapeField.decodeTriangle(scratch, triangles[i] = new ShapeField.DecodedTriangle());
165-
}
166-
return triangles;
167-
}
168145
}
169146

170147
public static class BoundingBox {

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeometryDocValueWriter.java

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,35 +6,32 @@
66

77
package org.elasticsearch.xpack.spatial.index.fielddata;
88

9-
import org.apache.lucene.document.ShapeField;
9+
import org.apache.lucene.index.IndexableField;
1010
import org.apache.lucene.store.ByteBuffersDataOutput;
11+
import org.apache.lucene.util.BytesRef;
1112

1213
import java.io.IOException;
1314
import java.util.List;
1415

1516
/**
16-
* This is a tree-writer that serializes a list of {@link ShapeField.DecodedTriangle} as an interval tree
17+
* This is a tree-writer that serializes a list of {@link IndexableField} as an interval tree
1718
* into a byte array.
1819
*/
1920
public class GeometryDocValueWriter {
2021

21-
private final TriangleTreeWriter treeWriter;
22-
private final CoordinateEncoder coordinateEncoder;
23-
private final CentroidCalculator centroidCalculator;
24-
25-
public GeometryDocValueWriter(List<ShapeField.DecodedTriangle> triangles, CoordinateEncoder coordinateEncoder,
26-
CentroidCalculator centroidCalculator) {
27-
this.coordinateEncoder = coordinateEncoder;
28-
this.centroidCalculator = centroidCalculator;
29-
this.treeWriter = new TriangleTreeWriter(triangles);
22+
private GeometryDocValueWriter() {
3023
}
3124

32-
/*** Serialize the interval tree in the provided data output */
33-
public void writeTo(ByteBuffersDataOutput out) throws IOException {
25+
/*** Serialize the triangle tree in a BytesRef */
26+
public static BytesRef write(List<IndexableField> fields,
27+
CoordinateEncoder coordinateEncoder,
28+
CentroidCalculator centroidCalculator) throws IOException {
29+
final ByteBuffersDataOutput out = new ByteBuffersDataOutput();
3430
out.writeInt(coordinateEncoder.encodeX(centroidCalculator.getX()));
3531
out.writeInt(coordinateEncoder.encodeY(centroidCalculator.getY()));
3632
centroidCalculator.getDimensionalShapeType().writeTo(out);
3733
out.writeVLong(Double.doubleToLongBits(centroidCalculator.sumWeight()));
38-
treeWriter.writeTo(out);
34+
TriangleTreeWriter.writeTo(out, fields);
35+
return new BytesRef(out.toArrayCopy(), 0, Math.toIntExact(out.size()));
3936
}
4037
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/TriangleTreeWriter.java

Lines changed: 30 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@
77
package org.elasticsearch.xpack.spatial.index.fielddata;
88

99
import org.apache.lucene.document.ShapeField;
10+
import org.apache.lucene.index.IndexableField;
1011
import org.apache.lucene.store.ByteBuffersDataOutput;
1112
import org.apache.lucene.util.ArrayUtil;
13+
import org.apache.lucene.util.BytesRef;
1214

1315
import java.io.IOException;
1416
import java.util.Comparator;
@@ -20,36 +22,43 @@
2022
*/
2123
class TriangleTreeWriter {
2224

23-
private final TriangleTreeNode node;
24-
private final Extent extent;
25-
26-
TriangleTreeWriter(List<ShapeField.DecodedTriangle> triangles) {
27-
this.extent = new Extent();
28-
this.node = build(triangles);
25+
private TriangleTreeWriter() {
2926
}
3027

3128
/*** Serialize the interval tree in the provided data output */
32-
public void writeTo(ByteBuffersDataOutput out) throws IOException {
29+
public static void writeTo(ByteBuffersDataOutput out, List<IndexableField> fields) throws IOException {
30+
final Extent extent = new Extent();
31+
final TriangleTreeNode node = build(fields, extent); ;
3332
extent.writeCompressed(out);
3433
node.writeTo(out);
3534
}
3635

37-
private TriangleTreeNode build(List<ShapeField.DecodedTriangle> triangles) {
38-
if (triangles.size() == 1) {
39-
TriangleTreeNode triangleTreeNode = new TriangleTreeNode(triangles.get(0));
36+
private static TriangleTreeNode build(List<IndexableField> fields, Extent extent) {
37+
final byte[] scratch = new byte[7 * Integer.BYTES];
38+
if (fields.size() == 1) {
39+
final TriangleTreeNode triangleTreeNode = new TriangleTreeNode(toDecodedTriangle(fields.get(0), scratch));
4040
extent.addRectangle(triangleTreeNode.minX, triangleTreeNode.minY, triangleTreeNode.maxX, triangleTreeNode.maxY);
4141
return triangleTreeNode;
4242
}
43-
TriangleTreeNode[] nodes = new TriangleTreeNode[triangles.size()];
44-
for (int i = 0; i < triangles.size(); i++) {
45-
nodes[i] = new TriangleTreeNode(triangles.get(i));
43+
final TriangleTreeNode[] nodes = new TriangleTreeNode[fields.size()];
44+
for (int i = 0; i < fields.size(); i++) {
45+
nodes[i] = new TriangleTreeNode(toDecodedTriangle(fields.get(i), scratch));
4646
extent.addRectangle(nodes[i].minX, nodes[i].minY, nodes[i].maxX, nodes[i].maxY);
4747
}
48-
return createTree(nodes, 0, triangles.size() - 1, true);
48+
return createTree(nodes, 0, fields.size() - 1, true);
49+
}
50+
51+
private static ShapeField.DecodedTriangle toDecodedTriangle(IndexableField field, byte[] scratch) {
52+
final BytesRef bytesRef = field.binaryValue();
53+
assert bytesRef.length == 7 * Integer.BYTES;
54+
System.arraycopy(bytesRef.bytes, bytesRef.offset, scratch, 0, 7 * Integer.BYTES);
55+
final ShapeField.DecodedTriangle decodedTriangle = new ShapeField.DecodedTriangle();
56+
ShapeField.decodeTriangle(scratch, decodedTriangle);
57+
return decodedTriangle;
4958
}
5059

5160
/** Creates tree from sorted components (with range low and high inclusive) */
52-
private TriangleTreeNode createTree(TriangleTreeNode[] components, int low, int high, boolean splitX) {
61+
private static TriangleTreeNode createTree(TriangleTreeNode[] components, int low, int high, boolean splitX) {
5362
if (low > high) {
5463
return null;
5564
}
@@ -95,16 +104,13 @@ private static class TriangleTreeNode {
95104
private TriangleTreeNode right;
96105
/** root node of edge tree */
97106
private final ShapeField.DecodedTriangle component;
98-
/** component type */
99-
private final ShapeField.DecodedTriangle.TYPE type;
100107

101108
private TriangleTreeNode(ShapeField.DecodedTriangle component) {
102109
this.minY = Math.min(Math.min(component.aY, component.bY), component.cY);
103110
this.maxY = Math.max(Math.max(component.aY, component.bY), component.cY);
104111
this.minX = Math.min(Math.min(component.aX, component.bX), component.cX);
105112
this.maxX = Math.max(Math.max(component.aX, component.bX), component.cX);
106113
this.component = component;
107-
this.type = component.type;
108114
}
109115

110116
private void writeTo(ByteBuffersDataOutput out) throws IOException {
@@ -141,9 +147,9 @@ private void writeMetadata(ByteBuffersDataOutput out) {
141147
byte metadata = 0;
142148
metadata |= (left != null) ? (1 << 0) : 0;
143149
metadata |= (right != null) ? (1 << 1) : 0;
144-
if (type == ShapeField.DecodedTriangle.TYPE.POINT) {
150+
if (component.type == ShapeField.DecodedTriangle.TYPE.POINT) {
145151
metadata |= (1 << 2);
146-
} else if (type == ShapeField.DecodedTriangle.TYPE.LINE) {
152+
} else if (component.type == ShapeField.DecodedTriangle.TYPE.LINE) {
147153
metadata |= (1 << 3);
148154
metadata |= (component.ab) ? (1 << 4) : 0;
149155
} else {
@@ -157,12 +163,12 @@ private void writeMetadata(ByteBuffersDataOutput out) {
157163
private void writeComponent(ByteBuffersDataOutput out) throws IOException {
158164
out.writeVLong((long) maxX - component.aX);
159165
out.writeVLong((long) maxY - component.aY);
160-
if (type == ShapeField.DecodedTriangle.TYPE.POINT) {
166+
if (component.type == ShapeField.DecodedTriangle.TYPE.POINT) {
161167
return;
162168
}
163169
out.writeVLong((long) maxX - component.bX);
164170
out.writeVLong((long) maxY - component.bY);
165-
if (type == ShapeField.DecodedTriangle.TYPE.LINE) {
171+
if (component.type == ShapeField.DecodedTriangle.TYPE.LINE) {
166172
return;
167173
}
168174
out.writeVLong((long) maxX - component.cX);
@@ -196,10 +202,10 @@ private int nodeSize(boolean includeBox, int parentMaxX, int parentMaxY, ByteBuf
196202

197203
private int componentSize(ByteBuffersDataOutput scratchBuffer) throws IOException {
198204
scratchBuffer.reset();
199-
if (type == ShapeField.DecodedTriangle.TYPE.POINT) {
205+
if (component.type == ShapeField.DecodedTriangle.TYPE.POINT) {
200206
scratchBuffer.writeVLong((long) maxX - component.aX);
201207
scratchBuffer.writeVLong((long) maxY - component.aY);
202-
} else if (type == ShapeField.DecodedTriangle.TYPE.LINE) {
208+
} else if (component.type == ShapeField.DecodedTriangle.TYPE.LINE) {
203209
scratchBuffer.writeVLong((long) maxX - component.aX);
204210
scratchBuffer.writeVLong((long) maxY - component.aY);
205211
scratchBuffer.writeVLong((long) maxX - component.bX);

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/BinaryGeoShapeDocValuesField.java

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,7 @@
66

77
package org.elasticsearch.xpack.spatial.index.mapper;
88

9-
import org.apache.lucene.document.ShapeField;
10-
import org.apache.lucene.store.ByteBuffersDataOutput;
9+
import org.apache.lucene.index.IndexableField;
1110
import org.apache.lucene.util.BytesRef;
1211
import org.elasticsearch.ElasticsearchException;
1312
import org.elasticsearch.index.mapper.CustomDocValuesField;
@@ -17,33 +16,29 @@
1716

1817
import java.io.IOException;
1918
import java.util.ArrayList;
20-
import java.util.Arrays;
2119
import java.util.List;
2220

2321
public class BinaryGeoShapeDocValuesField extends CustomDocValuesField {
2422

25-
private final List<ShapeField.DecodedTriangle> triangles;
23+
private final List<IndexableField> fields;
2624
private final CentroidCalculator centroidCalculator;
2725

28-
public BinaryGeoShapeDocValuesField(String name, ShapeField.DecodedTriangle[] triangles, CentroidCalculator centroidCalculator) {
26+
public BinaryGeoShapeDocValuesField(String name, List<IndexableField> fields, CentroidCalculator centroidCalculator) {
2927
super(name);
30-
this.triangles = new ArrayList<>(triangles.length);
28+
this.fields = new ArrayList<>(fields.size());
3129
this.centroidCalculator = centroidCalculator;
32-
this.triangles.addAll(Arrays.asList(triangles));
30+
this.fields.addAll(fields);
3331
}
3432

35-
public void add(ShapeField.DecodedTriangle[] triangles, CentroidCalculator centroidCalculator) {
36-
this.triangles.addAll(Arrays.asList(triangles));
33+
public void add( List<IndexableField> fields, CentroidCalculator centroidCalculator) {
34+
this.fields.addAll(fields);
3735
this.centroidCalculator.addFrom(centroidCalculator);
3836
}
3937

4038
@Override
4139
public BytesRef binaryValue() {
4240
try {
43-
final GeometryDocValueWriter writer = new GeometryDocValueWriter(triangles, CoordinateEncoder.GEO, centroidCalculator);
44-
final ByteBuffersDataOutput output = new ByteBuffersDataOutput();
45-
writer.writeTo(output);
46-
return new BytesRef(output.toArrayCopy(), 0, Math.toIntExact(output.size()));
41+
return GeometryDocValueWriter.write(fields, CoordinateEncoder.GEO, centroidCalculator);
4742
} catch (IOException e) {
4843
throw new ElasticsearchException("failed to encode shape", e);
4944
}

‎x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,8 @@
77
package org.elasticsearch.xpack.spatial.index.mapper;
88

99
import org.apache.lucene.document.LatLonShape;
10-
import org.apache.lucene.document.ShapeField;
1110
import org.apache.lucene.index.IndexableField;
1211
import org.apache.lucene.search.Query;
13-
import org.apache.lucene.util.BytesRef;
1412
import org.elasticsearch.Version;
1513
import org.elasticsearch.common.Explicit;
1614
import org.elasticsearch.common.geo.GeometryParser;
@@ -119,25 +117,16 @@ public GeoShapeWithDocValuesFieldMapper build(BuilderContext context) {
119117

120118
@Override
121119
@SuppressWarnings({"rawtypes", "unchecked"})
122-
protected void addDocValuesFields(String name, Geometry shape, List fields, ParseContext context) {
120+
protected void addDocValuesFields(String name, Geometry shape, List<IndexableField> fields, ParseContext context) {
123121
CentroidCalculator calculator = new CentroidCalculator(shape);
124-
final byte[] scratch = new byte[7 * Integer.BYTES];
125-
// doc values are generated from the indexed fields.
126-
ShapeField.DecodedTriangle[] triangles = new ShapeField.DecodedTriangle[fields.size()];
127-
for (int i = 0; i < fields.size(); i++) {
128-
BytesRef bytesRef = ((List<IndexableField>)fields).get(i).binaryValue();
129-
assert bytesRef.length == 7 * Integer.BYTES;
130-
System.arraycopy(bytesRef.bytes, bytesRef.offset, scratch, 0, 7 * Integer.BYTES);
131-
ShapeField.decodeTriangle(scratch, triangles[i] = new ShapeField.DecodedTriangle());
132-
}
133122
BinaryGeoShapeDocValuesField docValuesField =
134123
(BinaryGeoShapeDocValuesField) context.doc().getByKey(name);
135124
if (docValuesField == null) {
136-
docValuesField = new BinaryGeoShapeDocValuesField(name, triangles, calculator);
125+
docValuesField = new BinaryGeoShapeDocValuesField(name, fields, calculator);
137126
context.doc().addWithKey(name, docValuesField);
138127

139128
} else {
140-
docValuesField.add(triangles, calculator);
129+
docValuesField.add(fields, calculator);
141130
}
142131
}
143132

‎x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/aggregations/metrics/GeoShapeCentroidAggregatorTests.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import org.elasticsearch.xpack.spatial.LocalStateSpatialPlugin;
3434
import org.elasticsearch.xpack.spatial.index.fielddata.CentroidCalculator;
3535
import org.elasticsearch.xpack.spatial.index.fielddata.DimensionalShapeType;
36-
import org.elasticsearch.xpack.spatial.index.mapper.BinaryGeoShapeDocValuesField;
3736
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper.GeoShapeWithDocValuesFieldType;
3837
import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType;
3938
import org.elasticsearch.xpack.spatial.util.GeoTestUtils;
@@ -158,7 +157,7 @@ public void testSingleValuedField() throws Exception {
158157
for (Geometry geometry : geometries) {
159158
Document document = new Document();
160159
CentroidCalculator calculator = new CentroidCalculator(geometry);
161-
document.add(new BinaryGeoShapeDocValuesField("field", GeoTestUtils.toDecodedTriangles(geometry), calculator));
160+
document.add(GeoTestUtils.binaryGeoShapeDocValuesField("field", geometry));
162161
w.addDocument(document);
163162
if (targetShapeType.compareTo(calculator.getDimensionalShapeType()) == 0) {
164163
double weight = calculator.sumWeight();

‎x-pack/plugin/spatial/src/test/java/org/elasticsearch/xpack/spatial/index/fielddata/GeometryDocValueTests.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ public void testRectangleShape() throws IOException {
9393
int minY = randomIntBetween(-40, -1);
9494
int maxY = randomIntBetween(1, 40);
9595
Geometry rectangle = new Rectangle(minX, maxX, maxY, minY);
96-
GeometryDocValueReader reader = GeoTestUtils.GeometryDocValueReader(rectangle, CoordinateEncoder.GEO);
96+
GeometryDocValueReader reader = GeoTestUtils.geometryDocValueReader(rectangle, CoordinateEncoder.GEO);
9797

9898
Extent expectedExtent = getExtentFromBox(minX, minY, maxX, maxY);
9999
assertThat(expectedExtent, equalTo(reader.getExtent()));
@@ -114,7 +114,7 @@ private static Extent getExtentFromBox(double bottomLeftX, double bottomLeftY, d
114114
}
115115

116116
private static void assertDimensionalShapeType(Geometry geometry, DimensionalShapeType expected) throws IOException {
117-
GeometryDocValueReader reader = GeoTestUtils.GeometryDocValueReader(geometry, CoordinateEncoder.GEO);
117+
GeometryDocValueReader reader = GeoTestUtils.geometryDocValueReader(geometry, CoordinateEncoder.GEO);
118118
assertThat(reader.getDimensionalShapeType(), equalTo(expected));
119119
}
120120

0 commit comments

Comments
 (0)