Skip to content

Commit

Permalink
Support serialization of GeoShapeValue and CartesianShapeValue (#99299)
Browse files Browse the repository at this point in the history
* Support serialization of GeoShapeValue and CartesianShapeValue

* Improve test assert message

* Code review suggestions

* Use BytesReference in order to benefit from future plans for improved serialization

* Removed copyBytes by moving the writing closer to the actual byte[]

* Fixed unused import

* Simplfy read/write BytesReference by maintaining passed in BytesRef

* Added unit tests for serialization/deserialization

* Fixed license header, reduced duplicate code, reduced flaky tests
  • Loading branch information
craigtaverner committed Nov 15, 2023
1 parent db53786 commit 2a90ff4
Show file tree
Hide file tree
Showing 11 changed files with 258 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ static TransportVersion def(int id) {
public static final TransportVersion ML_TRAINED_MODEL_PREFIX_STRINGS_ADDED = def(8_535_00_0);
public static final TransportVersion COUNTED_KEYWORD_ADDED = def(8_536_00_0);

public static final TransportVersion SHAPE_VALUE_SERIALIZATION_ADDED = def(8_537_00_0);
/*
* STOP! READ THIS FIRST! No, really,
* ____ _____ ___ ____ _ ____ _____ _ ____ _____ _ _ ___ ____ _____ ___ ____ ____ _____ _
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@
import org.elasticsearch.xpack.spatial.action.SpatialStatsTransportAction;
import org.elasticsearch.xpack.spatial.action.SpatialUsageTransportAction;
import org.elasticsearch.xpack.spatial.common.CartesianBoundingBox;
import org.elasticsearch.xpack.spatial.index.fielddata.CartesianShapeValues;
import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues;
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeScriptFieldType;
import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper;
import org.elasticsearch.xpack.spatial.index.mapper.PointFieldMapper;
Expand Down Expand Up @@ -201,7 +203,11 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet

@Override
public List<GenericNamedWriteableSpec> getGenericNamedWriteables() {
return List.of(new GenericNamedWriteableSpec("CartesianBoundingBox", CartesianBoundingBox::new));
return List.of(
new GenericNamedWriteableSpec("CartesianBoundingBox", CartesianBoundingBox::new),
new GenericNamedWriteableSpec("GeoShapeValue", GeoShapeValues.GeoShapeValue::new),
new GenericNamedWriteableSpec("CartesianShapeValue", CartesianShapeValues.CartesianShapeValue::new)
);
}

private static void registerGeoShapeBoundsAggregator(ValuesSourceRegistry.Builder builder) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
import org.apache.lucene.geo.Component2D;
import org.apache.lucene.geo.XYGeometry;
import org.apache.lucene.geo.XYPoint;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.geometry.utils.GeometryValidator;
import org.elasticsearch.geometry.utils.StandardValidator;
import org.elasticsearch.search.aggregations.support.ValuesSourceType;
Expand Down Expand Up @@ -66,6 +67,11 @@ public CartesianShapeValue() {
super(CoordinateEncoder.CARTESIAN, CartesianPoint::new);
}

public CartesianShapeValue(StreamInput in) throws IOException {
this();
this.reset(in);
}

@Override
protected Component2D centroidAsComponent2D() throws IOException {
return XYGeometry.create(new XYPoint((float) getX(), (float) getY()));
Expand All @@ -80,5 +86,10 @@ protected Component2D centroidAsComponent2D() throws IOException {
public GeoRelation relate(XYGeometry geometry) throws IOException {
return relate(XYGeometry.create(geometry));
}

@Override
public String getWriteableName() {
return "CartesianShapeValue";
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import org.apache.lucene.geo.Point;
import org.elasticsearch.common.geo.GeoPoint;
import org.elasticsearch.common.geo.Orientation;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.geometry.utils.GeographyValidator;
import org.elasticsearch.geometry.utils.GeometryValidator;
import org.elasticsearch.index.mapper.GeoShapeIndexer;
Expand Down Expand Up @@ -69,6 +70,11 @@ public GeoShapeValue() {
this.tile2DVisitor = new Tile2DVisitor();
}

public GeoShapeValue(StreamInput in) throws IOException {
this();
reset(in);
}

@Override
protected Component2D centroidAsComponent2D() throws IOException {
return LatLonGeometry.create(new Point(getY(), getX()));
Expand All @@ -93,5 +99,10 @@ public GeoRelation relate(int minX, int maxX, int minY, int maxY) throws IOExcep
public GeoRelation relate(LatLonGeometry geometry) throws IOException {
return relate(LatLonGeometry.create(geometry));
}

@Override
public String getWriteableName() {
return "GeoShapeValue";
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ public class GeometryDocValueReader {
private final Extent extent;
private int treeOffset;
private int docValueOffset;
private BytesRef bytesRef;

public GeometryDocValueReader() {
this.extent = new Extent();
Expand All @@ -50,6 +51,7 @@ public GeometryDocValueReader() {
* reset the geometry.
*/
public void reset(BytesRef bytesRef) throws IOException {
this.bytesRef = bytesRef; // Needed only for supporting Writable, maintaining original offset, not adjusted on from input
this.input.reset(bytesRef.bytes, bytesRef.offset, bytesRef.length);
docValueOffset = bytesRef.offset;
treeOffset = 0;
Expand Down Expand Up @@ -109,4 +111,7 @@ public void visit(TriangleTreeVisitor visitor) throws IOException {
}
}

public BytesRef getBytesRef() {
return bytesRef;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,14 @@
import org.apache.lucene.document.ShapeField;
import org.apache.lucene.geo.Component2D;
import org.apache.lucene.util.BytesRef;
import org.elasticsearch.TransportVersion;
import org.elasticsearch.TransportVersions;
import org.elasticsearch.common.bytes.BytesArray;
import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.geo.SpatialPoint;
import org.elasticsearch.common.io.stream.GenericNamedWriteable;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.geometry.Geometry;
import org.elasticsearch.geometry.utils.GeometryValidator;
import org.elasticsearch.geometry.utils.WellKnownText;
Expand Down Expand Up @@ -90,7 +97,7 @@ public T missing(String missing) {
* thin wrapper around a {@link GeometryDocValueReader} which encodes / decodes values using
* the provided decoder (could be geo or cartesian)
*/
protected abstract static class ShapeValue implements ToXContentFragment {
protected abstract static class ShapeValue implements ToXContentFragment, GenericNamedWriteable {
private final GeometryDocValueReader reader;
private final BoundingBox boundingBox;
private final CoordinateEncoder encoder;
Expand All @@ -113,6 +120,11 @@ public void reset(BytesRef bytesRef) throws IOException {
this.boundingBox.reset(reader.getExtent(), encoder);
}

protected void reset(StreamInput in) throws IOException {
BytesReference bytes = in.readBytesReference();
reset(bytes.toBytesRef());
}

public BoundingBox boundingBox() {
return boundingBox;
}
Expand Down Expand Up @@ -187,6 +199,29 @@ public double getX() throws IOException {
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
throw new IllegalArgumentException("cannot write xcontent for geo_shape doc value");
}

@Override
public TransportVersion getMinimalSupportedVersion() {
return TransportVersions.SHAPE_VALUE_SERIALIZATION_ADDED;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeBytesReference(new BytesArray(reader.getBytesRef()));
}

@Override
public boolean equals(Object obj) {
if (obj instanceof ShapeValue other) {
return reader.getBytesRef().equals(other.reader.getBytesRef());
}
return false;
}

@Override
public int hashCode() {
return reader.getBytesRef().hashCode();
}
}

public static class BoundingBox {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,11 @@ public void testGenericNamedWriteables() {
.filter(e -> e.categoryClass.equals(GenericNamedWriteable.class))
.map(e -> e.name)
.collect(Collectors.toSet());
assertThat("Expect both Geo and Cartesian BoundingBox", names, equalTo(Set.of("GeoBoundingBox", "CartesianBoundingBox")));
assertThat(
"Expect both Geo and Cartesian BoundingBox and ShapeValue",
names,
equalTo(Set.of("GeoBoundingBox", "CartesianBoundingBox", "GeoShapeValue", "CartesianShapeValue"))
);
}

private SpatialPlugin getPluginWithOperationMode(License.OperationMode operationMode) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

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

import org.elasticsearch.geometry.Rectangle;
import org.elasticsearch.xpack.spatial.util.GeoTestUtils;

import java.io.IOException;

public class CartesianShapeValuesGenericWriteableTests extends ShapeValuesGenericWriteableTests<CartesianShapeValues.CartesianShapeValue> {

@Override
protected String shapeValueName() {
return "CartesianShapeValue";
}

@Override
protected GenericWriteableWrapper createTestInstance() {
try {
double minX = randomDoubleBetween(-Float.MAX_VALUE, 0, false);
double minY = randomDoubleBetween(-Float.MAX_VALUE, 0, false);
double maxX = randomDoubleBetween(minX + 10, Float.MAX_VALUE, false);
double maxY = randomDoubleBetween(minY + 10, Float.MAX_VALUE, false);
Rectangle rectangle = new Rectangle(minX, maxX, maxY, minY);
CartesianShapeValues.CartesianShapeValue shapeValue = GeoTestUtils.cartesianShapeValue(rectangle);
return new GenericWriteableWrapper(shapeValue);
} catch (IOException e) {
throw new RuntimeException(e);
}
}

@Override
protected GenericWriteableWrapper mutateInstance(GenericWriteableWrapper instance) throws IOException {
ShapeValues.ShapeValue shapeValue = instance.shapeValue();
ShapeValues.BoundingBox bbox = shapeValue.boundingBox();
double height = bbox.maxY() - bbox.minY();
double width = bbox.maxX() - bbox.minX();
double xs = width * 0.001;
double ys = height * 0.001;
Rectangle rectangle = new Rectangle(bbox.minX() + xs, bbox.maxX() - xs, bbox.maxY() - ys, bbox.minY() + ys);
return new GenericWriteableWrapper(GeoTestUtils.cartesianShapeValue(rectangle));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

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

import org.elasticsearch.common.geo.GeoBoundingBox;
import org.elasticsearch.geometry.Rectangle;
import org.elasticsearch.xpack.spatial.util.GeoTestUtils;

import java.io.IOException;

public class GeoShapeValuesGenericWriteableTests extends ShapeValuesGenericWriteableTests<GeoShapeValues.GeoShapeValue> {

@Override
protected String shapeValueName() {
return "GeoShapeValue";
}

@Override
protected GenericWriteableWrapper createTestInstance() {
try {
GeoBoundingBox bbox = GeoTestUtils.randomBBox();
Rectangle rectangle = new Rectangle(bbox.left(), bbox.right(), bbox.top(), bbox.bottom());
GeoShapeValues.GeoShapeValue shapeValue = GeoTestUtils.geoShapeValue(rectangle);
return new GenericWriteableWrapper(shapeValue);
} catch (IOException e) {
throw new RuntimeException(e);
}
}

@Override
protected GenericWriteableWrapper mutateInstance(GenericWriteableWrapper instance) throws IOException {
ShapeValues.ShapeValue shapeValue = instance.shapeValue();
ShapeValues.BoundingBox bbox = shapeValue.boundingBox();
double height = bbox.maxY() - bbox.minY();
double width = bbox.maxX() - bbox.minX();
double xs = width * 0.001;
double ys = height * 0.001;
Rectangle rectangle = new Rectangle(bbox.minX() + xs, bbox.maxX() - xs, bbox.maxY() - ys, bbox.minY() + ys);
return new GenericWriteableWrapper(GeoTestUtils.geoShapeValue(rectangle));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0; you may not use this file except in compliance with the Elastic License
* 2.0.
*/

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

import org.elasticsearch.TransportVersion;
import org.elasticsearch.TransportVersions;
import org.elasticsearch.common.io.stream.BytesStreamOutput;
import org.elasticsearch.common.io.stream.GenericNamedWriteable;
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
import org.elasticsearch.common.io.stream.StreamInput;
import org.elasticsearch.common.io.stream.StreamOutput;
import org.elasticsearch.common.io.stream.Writeable;
import org.elasticsearch.test.AbstractWireTestCase;

import java.io.IOException;
import java.util.List;

import static org.hamcrest.Matchers.containsString;

public abstract class ShapeValuesGenericWriteableTests<T extends ShapeValues.ShapeValue> extends AbstractWireTestCase<
ShapeValuesGenericWriteableTests.GenericWriteableWrapper> {

/**
* Wrapper around a GeoShapeValue to verify that it round-trips via {@code writeGenericValue} and {@code readGenericValue}
*/
public record GenericWriteableWrapper(ShapeValues.ShapeValue shapeValue) implements Writeable {
@Override
public void writeTo(StreamOutput out) throws IOException {
out.writeGenericValue(shapeValue);
}

public static GenericWriteableWrapper readFrom(StreamInput in) throws IOException {
return new GenericWriteableWrapper((ShapeValues.ShapeValue) in.readGenericValue());
}
}

private static final NamedWriteableRegistry NAMED_WRITEABLE_REGISTRY = new NamedWriteableRegistry(
List.of(
new NamedWriteableRegistry.Entry(
GenericNamedWriteable.class,
GeoShapeValues.GeoShapeValue.class.getSimpleName(),
GeoShapeValues.GeoShapeValue::new
),
new NamedWriteableRegistry.Entry(
GenericNamedWriteable.class,
CartesianShapeValues.CartesianShapeValue.class.getSimpleName(),
CartesianShapeValues.CartesianShapeValue::new
)
)
);

@Override
protected NamedWriteableRegistry writableRegistry() {
return NAMED_WRITEABLE_REGISTRY;
}

@Override
protected GenericWriteableWrapper copyInstance(GenericWriteableWrapper instance, TransportVersion version) throws IOException {
return copyInstance(instance, writableRegistry(), StreamOutput::writeWriteable, GenericWriteableWrapper::readFrom, version);
}

protected abstract String shapeValueName();

public void testSerializationFailsWithOlderVersion() {
TransportVersion older = TransportVersions.KNN_AS_QUERY_ADDED;
assert older.before(TransportVersions.SHAPE_VALUE_SERIALIZATION_ADDED);
final var testInstance = createTestInstance().shapeValue();
try (var output = new BytesStreamOutput()) {
output.setTransportVersion(older);
assertThat(
expectThrows(Throwable.class, () -> output.writeGenericValue(testInstance)).getMessage(),
containsString("[" + shapeValueName() + "] requires minimal transport version")
);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import org.elasticsearch.xcontent.XContentFactory;
import org.elasticsearch.xcontent.XContentParser;
import org.elasticsearch.xcontent.XContentType;
import org.elasticsearch.xpack.spatial.index.fielddata.CartesianShapeValues;
import org.elasticsearch.xpack.spatial.index.fielddata.CentroidCalculator;
import org.elasticsearch.xpack.spatial.index.fielddata.CoordinateEncoder;
import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues;
Expand Down Expand Up @@ -70,6 +71,12 @@ public static GeoShapeValues.GeoShapeValue geoShapeValue(Geometry geometry) thro
return value;
}

public static CartesianShapeValues.CartesianShapeValue cartesianShapeValue(Geometry geometry) throws IOException {
CartesianShapeValues.CartesianShapeValue value = new CartesianShapeValues.CartesianShapeValue();
value.reset(binaryCartesianShapeDocValuesField("test", geometry).binaryValue());
return value;
}

public static GeoBoundingBox randomBBox() {
Rectangle rectangle = GeometryTestUtils.randomRectangle();
return new GeoBoundingBox(
Expand Down

0 comments on commit 2a90ff4

Please sign in to comment.