Skip to content

Commit e5e3d44

Browse files
ncordonncordon
committed
Checkpoint
1 parent ed96b60 commit e5e3d44

File tree

1 file changed

+53
-86
lines changed
  • x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial

1 file changed

+53
-86
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StSimplify.java‎

Lines changed: 53 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,13 @@
1111
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
1212
import org.elasticsearch.common.io.stream.StreamInput;
1313
import org.elasticsearch.common.io.stream.StreamOutput;
14-
import org.elasticsearch.compute.data.Block;
15-
import org.elasticsearch.compute.data.BytesRefVectorBlock;
16-
import org.elasticsearch.compute.data.DoubleVector;
17-
import org.elasticsearch.compute.data.Page;
14+
import org.elasticsearch.compute.ann.ConvertEvaluator;
15+
import org.elasticsearch.compute.ann.Evaluator;
16+
import org.elasticsearch.compute.ann.Fixed;
17+
import org.elasticsearch.compute.data.BytesRefBlock;
1818
import org.elasticsearch.compute.operator.EvalOperator;
1919
import org.elasticsearch.geometry.Geometry;
20+
import org.elasticsearch.geometry.utils.GeometryValidator;
2021
import org.elasticsearch.geometry.utils.StandardValidator;
2122
import org.elasticsearch.geometry.utils.WellKnownBinary;
2223
import org.elasticsearch.geometry.utils.WellKnownText;
@@ -38,7 +39,6 @@
3839
import java.util.List;
3940

4041
import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_SHAPE;
41-
import static org.elasticsearch.xpack.esql.expression.function.scalar.spatial.SpatialRelatesUtils.makeGeometryFromLiteral;
4242

4343
public class StSimplify extends EsqlScalarFunction {
4444
public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
@@ -70,11 +70,7 @@ public StSimplify(
7070
}
7171

7272
private StSimplify(StreamInput in) throws IOException {
73-
this(
74-
Source.readFrom((PlanStreamInput) in),
75-
in.readNamedWriteable(Expression.class),
76-
in.readNamedWriteable(Expression.class)
77-
);
73+
this(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(Expression.class), in.readNamedWriteable(Expression.class));
7874
}
7975

8076
@Override
@@ -104,83 +100,54 @@ public void writeTo(StreamOutput out) throws IOException {
104100
out.writeNamedWriteable(tolerance);
105101
}
106102

103+
private static class GeoSimplifier {
104+
static GeometryValidator validator = StandardValidator.instance(true);
105+
static WKTReader reader = new WKTReader();
106+
static WKTWriter writer = new WKTWriter();
107+
108+
public static BytesRef geoSourceAndConstantTolerance(BytesRef inputGeometry, @Fixed double inputTolerance) {
109+
String wkt = WellKnownText.fromWKB(inputGeometry.bytes, inputGeometry.offset, inputGeometry.length);
110+
try {
111+
org.locationtech.jts.geom.Geometry jtsGeometry = reader.read(wkt);
112+
org.locationtech.jts.geom.Geometry simplifiedGeometry = DouglasPeuckerSimplifier.simplify(jtsGeometry, inputTolerance);
113+
String simplifiedWkt = writer.write(simplifiedGeometry);
114+
Geometry esGeometryResult = WellKnownText.fromWKT(validator, false, simplifiedWkt);
115+
return new BytesRef(WellKnownBinary.toWKB(esGeometryResult, ByteOrder.LITTLE_ENDIAN));
116+
} catch (Exception e) {
117+
throw new RuntimeException(e);
118+
}
119+
}
120+
}
121+
107122
@Override
108123
public EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
109-
return dvrCtx -> {
110-
EvalOperator.ExpressionEvaluator geometryEvaluator = toEvaluator.apply(geometry).get(dvrCtx);
111-
EvalOperator.ExpressionEvaluator toleranceEvaluator = toEvaluator.apply(tolerance).get(dvrCtx);
112-
113-
return new EvalOperator.ExpressionEvaluator() {
114-
@Override
115-
public void close() {
116-
117-
}
118-
119-
@Override
120-
public Block eval(Page page) {
121-
var isGeometryFoldable = geometry.foldable();
122-
var positionCount = page.getPositionCount();
123-
var validator = StandardValidator.instance(true);
124-
WKTReader reader = new WKTReader();
125-
WKTWriter writer = new WKTWriter();
126-
DoubleVector tolerances = (DoubleVector) toleranceEvaluator.eval(page).asVector();
127-
128-
// TODO We are not extracting non foldable geometries
129-
// TODO We are not using the tolerance
130-
try (var result = dvrCtx.blockFactory().newBytesRefVectorBuilder(positionCount)) {
131-
if (isGeometryFoldable) {
132-
var esGeometry = makeGeometryFromLiteral(toEvaluator.foldCtx(), geometry);
133-
String wkt = WellKnownText.toWKT(esGeometry);
134-
135-
try {
136-
org.locationtech.jts.geom.Geometry jtsGeometry = reader.read(wkt);
137-
138-
for (int p = 0; p < positionCount; p++) {
139-
double distanceTolerance = tolerances.getDouble(p);
140-
org.locationtech.jts.geom.Geometry simplifiedGeometry = DouglasPeuckerSimplifier.simplify(jtsGeometry, distanceTolerance);
141-
String simplifiedWkt = writer.write(simplifiedGeometry);
142-
Geometry esGeometryResult = WellKnownText.fromWKT(validator, false, simplifiedWkt);
143-
144-
result.appendBytesRef(new BytesRef(WellKnownBinary.toWKB(esGeometryResult, ByteOrder.LITTLE_ENDIAN)));
145-
}
146-
} catch (Exception e) {
147-
throw new RuntimeException(e);
148-
}
149-
} else {
150-
// Geometry is non foldable
151-
BytesRefVectorBlock block = (BytesRefVectorBlock) geometryEvaluator.eval(page);
152-
var bytesRefVector = block.asVector();
153-
154-
try {
155-
for (int p = 0; p < positionCount; p++) {
156-
double distanceTolerance = tolerances.getDouble(p);
157-
var destRef = bytesRefVector.getBytesRef(p, new BytesRef());
158-
var wkt = WellKnownText.fromWKB(destRef.bytes, destRef.offset, destRef.length);
159-
org.locationtech.jts.geom.Geometry jtsGeometry = reader.read(wkt);
160-
org.locationtech.jts.geom.Geometry simplifiedGeometry = DouglasPeuckerSimplifier.simplify(
161-
jtsGeometry,
162-
distanceTolerance
163-
);
164-
String simplifiedWkt = writer.write(simplifiedGeometry);
165-
Geometry esGeometryResult = WellKnownText.fromWKT(validator, false, simplifiedWkt);
166-
result.appendBytesRef(new BytesRef(WellKnownBinary.toWKB(esGeometryResult, ByteOrder.LITTLE_ENDIAN)));
167-
}
168-
} catch (Exception e) {
169-
throw new RuntimeException(e);
170-
}
171-
block.close();
172-
}
173-
tolerances.close();
174-
175-
return result.build().asBlock();
176-
}
177-
}
178-
179-
@Override
180-
public long baseRamBytesUsed() {
181-
return 0;
182-
}
183-
};
184-
};
124+
EvalOperator.ExpressionEvaluator.Factory geometryEvaluator = toEvaluator.apply(geometry);
125+
126+
if (tolerance.foldable() == false) {
127+
throw new IllegalArgumentException("tolerance must be foldable");
128+
}
129+
double inputTolerance = (double) tolerance.fold(toEvaluator.foldCtx());
130+
131+
if (geometry.foldable()) {
132+
BytesRef inputGeometry = (BytesRef) geometry.fold(toEvaluator.foldCtx());
133+
return new StSimplifyFoldableGeoAndConstantToleranceEvaluator.Factory(source(), inputGeometry, inputTolerance);
134+
}
135+
return new StSimplifyNonFoldableGeoAndConstantToleranceEvaluator.Factory(source(), geometryEvaluator, inputTolerance);
136+
}
137+
138+
@Evaluator(extraName = "NonFoldableGeoAndConstantTolerance", warnExceptions = { IllegalArgumentException.class })
139+
static BytesRef processNonFoldableGeoAndConstantTolerance(
140+
BytesRef inputGeometry,
141+
@Fixed double inputTolerance
142+
) {
143+
return GeoSimplifier.geoSourceAndConstantTolerance(inputGeometry, inputTolerance);
144+
}
145+
146+
@Evaluator(extraName = "FoldableGeoAndConstantTolerance", warnExceptions = { IllegalArgumentException.class })
147+
static BytesRef processFoldableGeoAndConstantTolerance(
148+
@Fixed BytesRef inputGeometry,
149+
@Fixed double inputTolerance
150+
) {
151+
return GeoSimplifier.geoSourceAndConstantTolerance(inputGeometry, inputTolerance);
185152
}
186153
}

0 commit comments

Comments
 (0)