Test StDistance multivalue consistency and fixed two CartesianPoint bugs (#114729)

craigtaverner · davidkyle · commit c074056874f8 · 2024-10-15T11:26:58.000+01:00
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownCartesianPointIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownCartesianPointIT.java
@@ -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;
+    }
 }
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownGeoPointIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownGeoPointIT.java
@@ -36,4 +36,9 @@ protected Geometry getQueryGeometry() {
     protected String castingFunction() {
         return "TO_GEOSHAPE";
     }
+
+    @Override
+    protected double searchDistance() {
+        return 10000000;
+    }
 }
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownPointsTestCase.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownPointsTestCase.java
@@ -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;
@@ -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;
@@ -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() + ")";
     }
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownTestCase.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/spatial/SpatialPushDownTestCase.java
@@ -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.
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/BinarySpatialFunction.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/BinarySpatialFunction.java
@@ -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;
@@ -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)
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/EnableSpatialDistancePushdown.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/EnableSpatialDistancePushdown.java
@@ -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());
+        };
     }
 
     /**

Original file line number	Diff line number	Diff line change
`@@ -31,4 +31,10 @@ protected Geometry getQueryGeometry() {`
`31`	`31`	`protected String castingFunction() {`
`32`	`32`	`return "TO_CARTESIANSHAPE";`
`33`	`33`	`}`
	`34`	`+`
	`35`	`+ @Override`
	`36`	`+ protected double searchDistance() {`
	`37`	`+ // We search much larger distances for Cartesian, to ensure we actually get results from the much wider data range`
	`38`	`+ return 1e12;`
	`39`	`+ }`
`34`	`40`	`}`
Original file line number	Diff line number	Diff line change
`@@ -36,4 +36,9 @@ protected Geometry getQueryGeometry() {`
`36`	`36`	`protected String castingFunction() {`
`37`	`37`	`return "TO_GEOSHAPE";`
`38`	`38`	`}`
	`39`	`+`
	`40`	`+ @Override`
	`41`	`+ protected double searchDistance() {`
	`42`	`+ return 10000000;`
	`43`	`+ }`
`39`	`44`	`}`