SQL: Add ST_Distance function to geosql (#39973)

* SQL: Add ST_Distance function to geosql

Adds ST_Distance function that works on points. No other shapes
are supported at the moments. No optimization and conversion into
geo_distance filter is done yet.

Relates to #29872

Author: imotov

docs/reference/sql/functions/geo.asciidoc

@@ -54,4 +54,29 @@ Returns the geometry from WKT representation. The return type is geometry.
 ["source","sql",subs="attributes,macros"]
 --------------------------------------------------
 include-tagged::{sql-specs}/geo/docs.csv-spec[aswkt]
+--------------------------------------------------
+
+[[sql-functions-geo-st-distance]]
+===== `ST_Distance`
+
+.Synopsis:
+[source, sql]
+--------------------------------------------------
+ST_Distance(geometry<1>, geometry<2>)
+--------------------------------------------------
+
+*Input*:
+
+<1> source geometry
+<2> target geometry
+
+*Output*: Double
+
+.Description:
+
+Returns the distance between geometries in meters. Both geometries have to be points. The return type is double.
+
+["source","sql",subs="attributes,macros"]
+--------------------------------------------------
+include-tagged::{sql-specs}/geo/docs.csv-spec[distance]
 --------------------------------------------------

x-pack/plugin/sql/qa/multi-node/src/test/java/org/elasticsearch/xpack/sql/qa/multi_node/GeoJdbcCsvSpecIT.java

@@ -0,0 +1,16 @@
+/*
+ * 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.qa.multi_node;
+
+import org.elasticsearch.xpack.sql.qa.geo.GeoCsvSpecTestCase;
+import org.elasticsearch.xpack.sql.qa.jdbc.CsvTestUtils.CsvTestCase;
+
+public class GeoJdbcCsvSpecIT extends GeoCsvSpecTestCase {
+    public GeoJdbcCsvSpecIT(String fileName, String groupName, String testName, Integer lineNumber, CsvTestCase testCase) {
+        super(fileName, groupName, testName, lineNumber, testCase);
+    }
+}

x-pack/plugin/sql/qa/multi-node/src/test/java/org/elasticsearch/xpack/sql/qa/multi_node/GeoJdbcSqlSpecIT.java

@@ -0,0 +1,15 @@
+/*
+ * 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.qa.multi_node;
+
+import org.elasticsearch.xpack.sql.qa.geo.GeoSqlSpecTestCase;
+
+public class GeoJdbcSqlSpecIT extends GeoSqlSpecTestCase {
+    public GeoJdbcSqlSpecIT(String fileName, String groupName, String testName, Integer lineNumber, String query) {
+        super(fileName, groupName, testName, lineNumber, query);
+    }
+}

x-pack/plugin/sql/qa/src/main/resources/command.csv-spec

@@ -132,6 +132,7 @@ DATABASE         |SCALAR
 USER             |SCALAR
 ST_ASTEXT        |SCALAR
 ST_ASWKT         |SCALAR
+ST_DISTANCE      |SCALAR
 ST_GEOMFROMTEXT  |SCALAR
 ST_WKTTOSQL      |SCALAR
 SCORE            |SCORE

x-pack/plugin/sql/qa/src/main/resources/docs.csv-spec

@@ -309,6 +309,7 @@ DATABASE         |SCALAR
 USER             |SCALAR
 ST_ASTEXT        |SCALAR
 ST_ASWKT         |SCALAR
+ST_DISTANCE      |SCALAR
 ST_GEOMFROMTEXT  |SCALAR
 ST_WKTTOSQL      |SCALAR
 SCORE            |SCORE

x-pack/plugin/sql/qa/src/main/resources/geo/docs.csv-spec

@@ -25,3 +25,13 @@ SELECT CAST(ST_WKTToSQL('POINT (10 20)') AS STRING) location;
 point (10.0 20.0)
 // end::wkttosql
 ;
+
+
+selectDistance
+// tag::distance
+SELECT ST_Distance(ST_WKTToSQL('POINT (10 20)'), ST_WKTToSQL('POINT (20 30)')) distance;
+
+   distance:d
+1499101.2889383635
+// end::distance
+;

x-pack/plugin/sql/qa/src/main/resources/geo/geosql.csv-spec

@@ -146,4 +146,58 @@ Asia           |Singapore
 Asia           |Sydney
 Europe         |Amsterdam
 Europe         |Berlin
-;
+;
+
+
+selectCitiesByDistance
+SELECT region, city, ST_Distance(location, ST_WktToSQL('POINT (-71 42)')) distance FROM geo WHERE distance < 5000000 ORDER BY region, city;
+
+   region:s    |    city:s     |    distance:d
+Americas       |Chicago        |1373941.5140200066
+Americas       |Mountain View  |4335936.909375596
+Americas       |New York       |285839.6579622518
+Americas       |Phoenix        |3692895.0346903414
+Americas       |San Francisco  |4343565.010996301
+;
+
+selectCitiesByDistanceFloored
+SELECT region, city, FLOOR(ST_Distance(location, ST_WktToSQL('POINT (-71 42)'))) distance FROM geo WHERE distance < 5000000 ORDER BY region, city;
+
+   region:s    |    city:s     |    distance:l
+Americas       |Chicago        |1373941
+Americas       |Mountain View  |4335936
+Americas       |New York       |285839
+Americas       |Phoenix        |3692895
+Americas       |San Francisco  |4343565
+;
+
+selectCitiesOrderByDistance
+SELECT region, city FROM geo ORDER BY ST_Distance(location, ST_WktToSQL('POINT (-71 42)')) ;
+
+   region:s    |    city:s
+Americas       |New York
+Americas       |Chicago
+Americas       |Phoenix
+Americas       |Mountain View
+Americas       |San Francisco
+Europe         |London
+Europe         |Paris
+Europe         |Amsterdam
+Europe         |Berlin
+Europe         |Munich
+Asia           |Tokyo
+Asia           |Seoul
+Asia           |Hong Kong
+Asia           |Singapore
+Asia           |Sydney
+;
+
+groupCitiesByDistance
+SELECT COUNT(*) count, FIRST(region) region FROM geo GROUP BY FLOOR(ST_Distance(location, ST_WktToSQL('POINT (-71 42)'))/5000000);
+
+    count:l    |   region:s
+5              |Americas
+5              |Europe
+3              |Asia
+2              |Asia
+;

x-pack/plugin/sql/qa/src/main/resources/geo/geosql.sql-spec

@@ -17,8 +17,8 @@ SELECT city, ST_GEOMFROMTEXT(ST_ASWKT(location)) shape_wkt, region FROM "geo" OR
 selectRegionUsingWktToSqlWithoutConvertion
 SELECT region, city, shape, ST_GEOMFROMTEXT(region_point) region_wkt  FROM geo ORDER BY region, city;
 
-selectCitiesWithAGroupByWktToSql
+selectCitiesWithGroupByWktToSql
 SELECT COUNT(city) city_by_region, ST_GEOMFROMTEXT(region_point) region_geom FROM geo WHERE city LIKE '%a%' GROUP BY region_geom ORDER BY city_by_region;
 
-selectCitiesWithEOrderByWktToSql
+selectCitiesWithOrderByWktToSql
 SELECT region, city, UCASE(ST_ASWKT(ST_GEOMFROMTEXT(region_point))) region_wkt FROM geo WHERE city LIKE '%e%' ORDER BY region_wkt, city;

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/function/FunctionRegistry.java

@@ -46,6 +46,7 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.WeekOfYear;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.Year;
 import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StAswkt;
+import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistance;
 import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StWkttosql;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ACos;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.ASin;
@@ -254,7 +255,9 @@ private void defineDefaultFunctions() {
 
         // Geo Functions
         addToMap(def(StAswkt.class, StAswkt::new, "ST_ASWKT", "ST_ASTEXT"),
-                def(StWkttosql.class, StWkttosql::new, "ST_WKTTOSQL", "ST_GEOMFROMTEXT"));
+                def(StWkttosql.class, StWkttosql::new, "ST_WKTTOSQL", "ST_GEOMFROMTEXT"),
+                def(StDistance.class, StDistance::new, "ST_DISTANCE")
+        );
 
         // Special
         addToMap(def(Score.class, Score::new, "SCORE"));

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

@@ -12,6 +12,7 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.NonIsoDateTimeProcessor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.datetime.QuarterProcessor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.geo.GeoProcessor;
+import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistanceProcessor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StWkttosqlProcessor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.BinaryMathProcessor;
 import org.elasticsearch.xpack.sql.expression.function.scalar.math.MathProcessor;
@@ -94,8 +95,10 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         entries.add(new Entry(Processor.class, LocateFunctionProcessor.NAME, LocateFunctionProcessor::new));
         entries.add(new Entry(Processor.class, ReplaceFunctionProcessor.NAME, ReplaceFunctionProcessor::new));
         entries.add(new Entry(Processor.class, SubstringFunctionProcessor.NAME, SubstringFunctionProcessor::new));
+        // geo
         entries.add(new Entry(Processor.class, GeoProcessor.NAME, GeoProcessor::new));
         entries.add(new Entry(Processor.class, StWkttosqlProcessor.NAME, StWkttosqlProcessor::new));
+        entries.add(new Entry(Processor.class, StDistanceProcessor.NAME, StDistanceProcessor::new));
         return entries;
     }
 

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

@@ -5,7 +5,6 @@
  */
 package org.elasticsearch.xpack.sql.expression.function.scalar.geo;
 
-import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
@@ -18,9 +17,7 @@ public class GeoProcessor implements Processor {
 
     private interface GeoShapeFunction<R> {
         default R apply(Object o) {
-            if (o instanceof GeoPoint) {
-                return doApply(new GeoShape(((GeoPoint) o).getLon(), ((GeoPoint) o).getLat()));
-            } else if (o instanceof GeoShape) {
+            if (o instanceof GeoShape) {
                 return doApply((GeoShape) o);
             } else {
                 throw new SqlIllegalArgumentException("A geo_point or geo_shape is required; received [{}]", o);

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

@@ -5,21 +5,29 @@
  */
 package org.elasticsearch.xpack.sql.expression.function.scalar.geo;
 
+import org.elasticsearch.common.geo.GeoUtils;
 import org.elasticsearch.common.geo.builders.PointBuilder;
 import org.elasticsearch.common.geo.builders.ShapeBuilder;
 import org.elasticsearch.common.geo.parsers.ShapeParser;
+import org.elasticsearch.common.io.stream.NamedWriteable;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.ToXContentFragment;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 
 import java.io.IOException;
+import java.util.Objects;
 
 /**
  * Wrapper class to represent a GeoShape in SQL
  *
  * It is required to override the XContent serialization. The ShapeBuilder serializes using GeoJSON by default,
  * but in SQL we need the serialization to be WKT-based.
  */
-public class GeoShape implements ToXContentFragment {
+public class GeoShape implements ToXContentFragment, NamedWriteable {
+
+    public static final String NAME = "geo";
 
     private final ShapeBuilder<?, ?, ?> shapeBuilder;
 
@@ -31,6 +39,15 @@ public GeoShape(Object value) throws IOException {
         shapeBuilder = ShapeParser.parse(value);
     }
 
+    public GeoShape(StreamInput in) throws IOException {
+        shapeBuilder = ShapeParser.parse(in.readString());
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(shapeBuilder.toWKT());
+    }
+
     @Override
     public String toString() {
         return shapeBuilder.toWKT();
@@ -40,4 +57,41 @@ public String toString() {
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         return builder.value(shapeBuilder.toWKT());
     }
+
+    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);
+        }
+        if (shape2.shapeBuilder instanceof PointBuilder == false) {
+            throw new SqlIllegalArgumentException("distance calculation is only supported for points; received [{}]", shape2);
+        }
+        double srcLat = ((PointBuilder) shape1.shapeBuilder).latitude();
+        double srcLon = ((PointBuilder) shape1.shapeBuilder).longitude();
+        double dstLat = ((PointBuilder) shape2.shapeBuilder).latitude();
+        double dstLon = ((PointBuilder) shape2.shapeBuilder).longitude();
+        return GeoUtils.arcDistance(srcLat, srcLon, dstLat, dstLon);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        GeoShape geoShape = (GeoShape) o;
+        return shapeBuilder.equals(geoShape.shapeBuilder);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(shapeBuilder);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
 }

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

@@ -0,0 +1,82 @@
+/*
+ * 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.expression.function.scalar.geo;
+
+import org.elasticsearch.xpack.sql.expression.Expression;
+import org.elasticsearch.xpack.sql.expression.Expressions;
+import org.elasticsearch.xpack.sql.expression.FieldAttribute;
+import org.elasticsearch.xpack.sql.expression.function.scalar.BinaryScalarFunction;
+import org.elasticsearch.xpack.sql.expression.gen.pipeline.Pipe;
+import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
+import org.elasticsearch.xpack.sql.tree.NodeInfo;
+import org.elasticsearch.xpack.sql.tree.Source;
+import org.elasticsearch.xpack.sql.type.DataType;
+
+import static org.elasticsearch.xpack.sql.expression.TypeResolutions.isGeo;
+import static org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistanceProcessor.process;
+import static org.elasticsearch.xpack.sql.expression.gen.script.ParamsBuilder.paramsBuilder;
+
+/**
+ * Calculates the distance between two points
+ */
+public class StDistance extends BinaryScalarFunction {
+
+    public StDistance(Source source, Expression source1, Expression source2) {
+        super(source, source1, source2);
+    }
+
+    @Override
+    protected StDistance replaceChildren(Expression newLeft, Expression newRight) {
+        return new StDistance(source(), newLeft, newRight);
+    }
+
+    @Override
+    protected TypeResolution resolveType() {
+        if (!childrenResolved()) {
+            return new TypeResolution("Unresolved children");
+        }
+
+        TypeResolution resolution = isGeo(left(), functionName(), Expressions.ParamOrdinal.FIRST);
+        if (resolution.unresolved()) {
+            return resolution;
+        }
+
+        return isGeo(right(), functionName(), Expressions.ParamOrdinal.SECOND);
+    }
+
+    @Override
+    public DataType dataType() {
+        return DataType.DOUBLE;
+    }
+
+    @Override
+    protected NodeInfo<StDistance> info() {
+        return NodeInfo.create(this, StDistance::new, left(), right());
+    }
+
+    @Override
+    public ScriptTemplate scriptWithField(FieldAttribute field) {
+        return new ScriptTemplate(processScript("{sql}.geoDocValue(doc,{})"),
+            paramsBuilder().variable(field.exactAttribute().name()).build(),
+            dataType());
+    }
+
+    @Override
+    public Object fold() {
+        return process(left().fold(), right().fold());
+    }
+
+    @Override
+    protected Pipe makePipe() {
+        return new StDistancePipe(source(), this, Expressions.pipe(left()), Expressions.pipe(right()));
+    }
+
+    @Override
+    protected String scriptMethodName() {
+        return "stDistance";
+    }
+}