Skip to content

Commit

Permalink
Test StDistance multivalue consistency and fixed two CartesianPoint b…
Browse files Browse the repository at this point in the history
  • Loading branch information
craigtaverner authored Oct 14, 2024
1 parent f5188af commit 35e79f8
Show file tree
Hide file tree
Showing 6 changed files with 87 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -31,4 +31,10 @@ protected Geometry getQueryGeometry() {
protected String castingFunction() {
return "TO_CARTESIANSHAPE";
}

@Override
protected double searchDistance() {
// We search much larger distances for Cartesian, to ensure we actually get results from the much wider data range
return 1e12;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -36,4 +36,9 @@ protected Geometry getQueryGeometry() {
protected String castingFunction() {
return "TO_GEOSHAPE";
}

@Override
protected double searchDistance() {
return 10000000;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

package org.elasticsearch.xpack.esql.spatial;

import org.elasticsearch.geometry.Geometry;
import org.elasticsearch.geometry.Point;
import org.elasticsearch.geometry.utils.GeometryValidator;
import org.elasticsearch.geometry.utils.WellKnownText;
Expand All @@ -20,6 +21,7 @@
import java.io.IOException;
import java.text.ParseException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Locale;

import static org.hamcrest.Matchers.closeTo;
Expand Down Expand Up @@ -108,6 +110,61 @@ protected void assertFunction(String spatialFunction, String wkt, long expected,
}
}

public void testPushedDownDistanceSingleValue() throws RuntimeException {
assertPushedDownDistance(false);
}

public void testPushedDownDistanceMultiValue() throws RuntimeException {
assertPushedDownDistance(true);
}

private void assertPushedDownDistance(boolean multiValue) throws RuntimeException {
initIndexes();
for (int i = 0; i < random().nextInt(50, 100); i++) {
if (multiValue) {
final String[] values = new String[randomIntBetween(1, 5)];
for (int j = 0; j < values.length; j++) {
values[j] = "\"" + WellKnownText.toWKT(getIndexGeometry()) + "\"";
}
index("indexed", i + "", "{\"location\" : " + Arrays.toString(values) + " }");
index("not-indexed", i + "", "{\"location\" : " + Arrays.toString(values) + " }");
} else {
final String value = WellKnownText.toWKT(getIndexGeometry());
index("indexed", i + "", "{\"location\" : \"" + value + "\" }");
index("not-indexed", i + "", "{\"location\" : \"" + value + "\" }");
}
}

refresh("indexed", "not-indexed");

for (int i = 0; i < 10; i++) {
final Geometry geometry = getIndexGeometry();
final String wkt = WellKnownText.toWKT(geometry);
assertDistanceFunction(wkt);
}
}

protected abstract double searchDistance();

protected void assertDistanceFunction(String wkt) {
String spatialFunction = "ST_DISTANCE";
String castingFunction = castingFunction().replaceAll("SHAPE", "POINT");
final String query1 = String.format(Locale.ROOT, """
FROM indexed | WHERE %s(location, %s("%s")) < %.1f | STATS COUNT(*)
""", spatialFunction, castingFunction, wkt, searchDistance());
final String query2 = String.format(Locale.ROOT, """
FROM not-indexed | WHERE %s(location, %s("%s")) < %.1f | STATS COUNT(*)
""", spatialFunction, castingFunction, wkt, searchDistance());
try (
EsqlQueryResponse response1 = EsqlQueryRequestBuilder.newRequestBuilder(client()).query(query1).get();
EsqlQueryResponse response2 = EsqlQueryRequestBuilder.newRequestBuilder(client()).query(query2).get();
) {
Object indexedResult = response1.response().column(0).iterator().next();
Object notIndexedResult = response2.response().column(0).iterator().next();
assertEquals(spatialFunction, indexedResult, notIndexedResult);
}
}

private String toString(CentroidCalculator centroid) {
return "Centroid (x:" + centroid.getX() + ", y:" + centroid.getY() + ")";
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
/**
* Base class to check that a query than can be pushed down gives the same result
* if it is actually pushed down and when it is executed by the compute engine,
*
* <p>
* For doing that we create two indices, one fully indexed and another with index
* and doc values disabled. Then we index the same data in both indices and we check
* that the same ES|QL queries produce the same results in both.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_POINT;
import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_SHAPE;
import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_POINT;
import static org.elasticsearch.xpack.esql.core.type.DataType.GEO_SHAPE;
import static org.elasticsearch.xpack.esql.core.type.DataType.isNull;
Expand Down Expand Up @@ -203,7 +205,7 @@ public void setCrsType(DataType dataType) {
}

private static final String[] GEO_TYPE_NAMES = new String[] { GEO_POINT.typeName(), GEO_SHAPE.typeName() };
private static final String[] CARTESIAN_TYPE_NAMES = new String[] { GEO_POINT.typeName(), GEO_SHAPE.typeName() };
private static final String[] CARTESIAN_TYPE_NAMES = new String[] { CARTESIAN_POINT.typeName(), CARTESIAN_SHAPE.typeName() };

protected static boolean spatialCRSCompatible(DataType spatialDataType, DataType otherDataType) {
return DataType.isSpatialGeo(spatialDataType) && DataType.isSpatialGeo(otherDataType)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -216,31 +216,40 @@ private Expression rewriteDistanceFilter(
Number number,
ComparisonType comparisonType
) {
DataType shapeDataType = getShapeDataType(spatialExp);
Geometry geometry = SpatialRelatesUtils.makeGeometryFromLiteral(literalExp);
if (geometry instanceof Point point) {
double distance = number.doubleValue();
Source source = comparison.source();
if (comparisonType.lt) {
distance = comparisonType.eq ? distance : Math.nextDown(distance);
return new SpatialIntersects(source, spatialExp, makeCircleLiteral(point, distance, literalExp));
return new SpatialIntersects(source, spatialExp, makeCircleLiteral(point, distance, literalExp, shapeDataType));
} else if (comparisonType.gt) {
distance = comparisonType.eq ? distance : Math.nextUp(distance);
return new SpatialDisjoint(source, spatialExp, makeCircleLiteral(point, distance, literalExp));
return new SpatialDisjoint(source, spatialExp, makeCircleLiteral(point, distance, literalExp, shapeDataType));
} else if (comparisonType.eq) {
return new And(
source,
new SpatialIntersects(source, spatialExp, makeCircleLiteral(point, distance, literalExp)),
new SpatialDisjoint(source, spatialExp, makeCircleLiteral(point, Math.nextDown(distance), literalExp))
new SpatialIntersects(source, spatialExp, makeCircleLiteral(point, distance, literalExp, shapeDataType)),
new SpatialDisjoint(source, spatialExp, makeCircleLiteral(point, Math.nextDown(distance), literalExp, shapeDataType))
);
}
}
return comparison;
}

private Literal makeCircleLiteral(Point point, double distance, Expression literalExpression) {
private Literal makeCircleLiteral(Point point, double distance, Expression literalExpression, DataType shapeDataType) {
var circle = new Circle(point.getX(), point.getY(), distance);
var wkb = WellKnownBinary.toWKB(circle, ByteOrder.LITTLE_ENDIAN);
return new Literal(literalExpression.source(), new BytesRef(wkb), DataType.GEO_SHAPE);
return new Literal(literalExpression.source(), new BytesRef(wkb), shapeDataType);
}

private DataType getShapeDataType(Expression expression) {
return switch (expression.dataType()) {
case GEO_POINT, GEO_SHAPE -> DataType.GEO_SHAPE;
case CARTESIAN_POINT, CARTESIAN_SHAPE -> DataType.CARTESIAN_SHAPE;
default -> throw new IllegalArgumentException("Unsupported spatial data type: " + expression.dataType());
};
}

/**
Expand Down

0 comments on commit 35e79f8

Please sign in to comment.