SQL: Convert ST_Distance into query when possible

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/geo/GeoShape.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ToXContentFragment;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.geo.geometry.Geometry;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 
 import java.io.IOException;
@@ -58,6 +59,10 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         return builder.value(shapeBuilder.toWKT());
     }
 
+    public Geometry toGeometry() {
+        return shapeBuilder.buildGeometry();
+    }
+
     public static double distance(GeoShape shape1, GeoShape shape2) {
         if (shape1.shapeBuilder instanceof PointBuilder == false) {
             throw new SqlIllegalArgumentException("distance calculation is only supported for points; received [{}]", shape1);

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/scalar/geo/StDistance.java

@@ -65,6 +65,10 @@ public ScriptTemplate scriptWithField(FieldAttribute field) {
             dataType());
     }
 
+    public StDistance swapLeftAndRight() {
+        return new StDistance(source(), right(), left());
+    }
+
     @Override
     public Object fold() {
         return process(left().fold(), right().fold());

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/optimizer/Optimizer.java

@@ -42,6 +42,7 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.Cast;
 import org.elasticsearch.xpack.sql.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.sql.expression.function.scalar.ScalarFunctionAttribute;
+import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistance;
 import org.elasticsearch.xpack.sql.expression.predicate.BinaryOperator;
 import org.elasticsearch.xpack.sql.expression.predicate.BinaryPredicate;
 import org.elasticsearch.xpack.sql.expression.predicate.Negatable;
@@ -134,6 +135,8 @@ protected Iterable<RuleExecutor<LogicalPlan>.Batch> batches() {
                 // needs to occur before BinaryComparison combinations (see class)
                 new PropagateEquals(),
                 new CombineBinaryComparisons(),
+                // Geo
+                new StDistanceLiteralsOnTheRight(),
                 // prune/elimination
                 new PruneFilters(),
                 new PruneOrderBy(),
@@ -1407,6 +1410,22 @@ private Expression literalToTheRight(BinaryOperator<?, ?, ?, ?> be) {
         }
     }
 
+
+    static class StDistanceLiteralsOnTheRight extends OptimizerExpressionRule {
+
+        StDistanceLiteralsOnTheRight() {
+            super(TransformDirection.UP);
+        }
+
+        @Override
+        protected Expression rule(Expression e) {
+            return e instanceof StDistance ? literalToTheRight((StDistance) e) : e;
+        }
+
+        private Expression literalToTheRight(StDistance sd) {
+            return sd.left() instanceof Literal && !(sd.right() instanceof Literal) ? sd.swapLeftAndRight() : sd;
+        }
+    }
     /**
      * Propagate Equals to eliminate conjuncted Ranges.
      * When encountering a different Equals or non-containing {@link Range}, the conjunction becomes false.

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/QueryTranslator.java

@@ -5,6 +5,8 @@
  */
 package org.elasticsearch.xpack.sql.planner;
 
+import org.elasticsearch.geo.geometry.Geometry;
+import org.elasticsearch.geo.geometry.Point;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 import org.elasticsearch.xpack.sql.expression.Attribute;
@@ -38,6 +40,8 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeFunction;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.DateTimeHistogramFunction;
+import org.elasticsearch.xpack.sql.expression.function.scalar.geo.GeoShape;
+import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistance;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
 import org.elasticsearch.xpack.sql.expression.literal.Intervals;
 import org.elasticsearch.xpack.sql.expression.predicate.Range;
@@ -85,6 +89,7 @@
 import org.elasticsearch.xpack.sql.querydsl.agg.TopHitsAgg;
 import org.elasticsearch.xpack.sql.querydsl.query.BoolQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.ExistsQuery;
+import org.elasticsearch.xpack.sql.querydsl.query.GeoDistanceQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.MatchQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.MultiMatchQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.NestedQuery;
@@ -675,6 +680,21 @@ private static Query translateQuery(BinaryComparison bc) {
                 return new RangeQuery(source, name, value, true, null, false, format);
             }
             if (bc instanceof LessThan) {
+                if (bc.left() instanceof StDistance && value instanceof Number) {
+                    // Special case for ST_Distance translatable into geo_distance query
+                    StDistance stDistance = (StDistance) bc.left();
+                    if (stDistance.left() instanceof FieldAttribute && stDistance.right().foldable()) {
+                        Object geoShape = valueOf(stDistance.right());
+                        if (geoShape instanceof GeoShape) {
+                            Geometry geometry = ((GeoShape) geoShape).toGeometry();
+                            if (geometry instanceof Point) {
+                                String field = nameOf(stDistance.left());
+                                return new GeoDistanceQuery(source, field, ((Number) value).doubleValue(),
+                                    ((Point) geometry).getLat(), ((Point) geometry).getLon());
+                            }
+                        }
+                    }
+                }
                 return new RangeQuery(source, name, null, false, value, false, format);
             }
             if (bc instanceof LessThanOrEqual) {
@@ -966,6 +986,9 @@ public QueryTranslation translate(Expression exp, boolean onAggs) {
 
         protected static Query handleQuery(ScalarFunction sf, Expression field, Supplier<Query> query) {
             Query q = query.get();
+            if (field instanceof StDistance && q instanceof GeoDistanceQuery) {
+                return wrapIfNested(q, ((StDistance) field).left());
+            }
             if (field instanceof FieldAttribute) {
                 return wrapIfNested(q, field);
             }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/querydsl/query/GeoDistanceQuery.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.sql.querydsl.query;
+
+import org.elasticsearch.common.unit.DistanceUnit;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.xpack.sql.tree.Source;
+
+import java.util.Objects;
+
+public class GeoDistanceQuery extends LeafQuery {
+
+    private final String field;
+    private final double lat;
+    private final double lon;
+    private final double distance;
+
+    public GeoDistanceQuery(Source source, String field, double distance, double lat, double lon) {
+        super(source);
+        this.field = field;
+        this.distance = distance;
+        this.lat = lat;
+        this.lon = lon;
+    }
+
+    public String field() {
+        return field;
+    }
+
+    public double lat() {
+        return lat;
+    }
+
+    public double lon() {
+        return lon;
+    }
+
+    public double distance() {
+        return distance;
+    }
+
+    @Override
+    public QueryBuilder asBuilder() {
+        return QueryBuilders.geoDistanceQuery(field).distance(distance, DistanceUnit.METERS).point(lat, lon);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(field, distance, lat, lon);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (this == obj) {
+            return true;
+        }
+
+        if (obj == null || getClass() != obj.getClass()) {
+            return false;
+        }
+
+        GeoDistanceQuery other = (GeoDistanceQuery) obj;
+        return Objects.equals(field, other.field) &&
+            Objects.equals(distance, other.distance) &&
+            Objects.equals(lat, other.lat) &&
+            Objects.equals(lon, other.lon);
+    }
+
+    @Override
+    protected String innerToString() {
+        return field + ":" + "(" + distance + "," + "(" + lat + ", " +  lon + "))";
+    }
+}

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/optimizer/OptimizerTests.java

@@ -31,6 +31,7 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.IsoWeekOfYear;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.MonthOfYear;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.Year;
+import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistance;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ACos;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ASin;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ATan;
@@ -81,6 +82,7 @@
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.ReplaceFoldableAttributes;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.ReplaceMinMaxWithTopHits;
 import org.elasticsearch.xpack.sql.optimizer.Optimizer.SimplifyConditional;
+import org.elasticsearch.xpack.sql.optimizer.Optimizer.StDistanceLiteralsOnTheRight;
 import org.elasticsearch.xpack.sql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.sql.plan.logical.Filter;
 import org.elasticsearch.xpack.sql.plan.logical.LocalRelation;
@@ -622,6 +624,15 @@ public void testLiteralsOnTheRight() {
         assertEquals(FIVE, nullEquals.right());
     }
 
+    public void testLiteralsOnTheRightInStDistance() {
+        Alias a = new Alias(EMPTY, "a", L(10));
+        Expression result = new StDistanceLiteralsOnTheRight().rule(new StDistance(EMPTY, FIVE, a));
+        assertTrue(result instanceof StDistance);
+        StDistance sd = (StDistance) result;
+        assertEquals(a, sd.left());
+        assertEquals(FIVE, sd.right());
+    }
+
     public void testBoolSimplifyNotIsNullAndNotIsNotNull() {
         BooleanSimplification simplification = new BooleanSimplification();
         assertTrue(simplification.rule(new Not(EMPTY, new IsNull(EMPTY, ONE))) instanceof IsNotNull);

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/planner/QueryTranslatorTests.java

@@ -39,6 +39,7 @@
 import org.elasticsearch.xpack.sql.querydsl.agg.GroupByDateHistogram;
 import org.elasticsearch.xpack.sql.querydsl.query.BoolQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.ExistsQuery;
+import org.elasticsearch.xpack.sql.querydsl.query.GeoDistanceQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.NotQuery;
 import org.elasticsearch.xpack.sql.querydsl.query.Query;
 import org.elasticsearch.xpack.sql.querydsl.query.QueryStringQuery;
@@ -596,22 +597,53 @@ public void testTranslateStWktToSql() {
         assertEquals("[{v=keyword}, {v=point (10.0 20.0)}]", aggFilter.scriptTemplate().params().toString());
     }
 
-    public void testTranslateStDistance() {
-        LogicalPlan p = plan("SELECT shape FROM test WHERE ST_Distance(shape, ST_WKTToSQL('point (10 20)')) > 20");
+    public void testTranslateStDistanceToScript() {
+        String operator = randomFrom(">", ">=", "<=");
+        String operatorFunction = null;
+        switch (operator) {
+            case ">":
+                operatorFunction = "gt";
+                break;
+            case ">=":
+                operatorFunction = "gte";
+                break;
+            case "<=":
+                operatorFunction = "lte";
+                break;
+            default:
+                fail("Unexpected operator [" + operator + "]");
+        }
+        LogicalPlan p = plan("SELECT shape FROM test WHERE ST_Distance(shape, ST_WKTToSQL('point (10 20)')) " + operator + " 20");
         assertThat(p, instanceOf(Project.class));
         assertThat(p.children().get(0), instanceOf(Filter.class));
         Expression condition = ((Filter) p.children().get(0)).condition();
         assertFalse(condition.foldable());
-        QueryTranslation translation = QueryTranslator.toQuery(condition, true);
-        assertNull(translation.query);
-        AggFilter aggFilter = translation.aggFilter;
-
+        QueryTranslation translation = QueryTranslator.toQuery(condition, false);
+        assertNull(translation.aggFilter);
+        assertTrue(translation.query instanceof ScriptQuery);
+        ScriptQuery sc = (ScriptQuery) translation.query;
         assertEquals("InternalSqlScriptUtils.nullSafeFilter(" +
-                "InternalSqlScriptUtils.gt(" +
+                "InternalSqlScriptUtils." + operatorFunction + "(" +
                 "InternalSqlScriptUtils.stDistance(" +
                 "InternalSqlScriptUtils.geoDocValue(doc,params.v0),InternalSqlScriptUtils.stWktToSql(params.v1)),params.v2))",
-            aggFilter.scriptTemplate().toString());
-        assertEquals("[{v=shape}, {v=point (10.0 20.0)}, {v=20}]", aggFilter.scriptTemplate().params().toString());
+            sc.script().toString());
+        assertEquals("[{v=shape}, {v=point (10.0 20.0)}, {v=20}]", sc.script().params().toString());
+    }
+
+    public void testTranslateStDistanceToQuery() {
+        LogicalPlan p = plan("SELECT shape FROM test WHERE ST_Distance(shape, ST_WKTToSQL('point (10 20)')) < 25");
+        assertThat(p, instanceOf(Project.class));
+        assertThat(p.children().get(0), instanceOf(Filter.class));
+        Expression condition = ((Filter) p.children().get(0)).condition();
+        assertFalse(condition.foldable());
+        QueryTranslation translation = QueryTranslator.toQuery(condition, false);
+        assertNull(translation.aggFilter);
+        assertTrue(translation.query instanceof GeoDistanceQuery);
+        GeoDistanceQuery gq = (GeoDistanceQuery) translation.query;
+        assertEquals("shape", gq.field());
+        assertEquals(20.0, gq.lat(), 0.00001);
+        assertEquals(10.0, gq.lon(), 0.00001);
+        assertEquals(25.0, gq.distance(), 0.00001);
     }
 
     public void testTranslateCoalesce_GroupBy_Painless() {