QL: Turn eager query translations lazy (#69205)

Before the `ExpressionTranslators` did some unnecessary
`Expression` -> `Query` translations, where the result queries were thrown
away by later translations (by the `wrapFunctionQuery`).

This change adds laziness, so translations don't execute unnecessarily.

Follows #68783 and #68788

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

@@ -84,7 +84,7 @@ final class QueryTranslator {
 
         public static Query doTranslate(InsensitiveBinaryComparison bc, TranslatorHandler handler) {
             checkInsensitiveComparison(bc);
-            return handler.wrapFunctionQuery(bc, bc.left(), translate(bc, handler));
+            return handler.wrapFunctionQuery(bc, bc.left(), () -> translate(bc, handler));
         }
 
         public static void checkInsensitiveComparison(InsensitiveBinaryComparison bc) {
@@ -163,7 +163,7 @@ final class QueryTranslator {
                 }
             }
 
-            return handler.wrapFunctionQuery(f, f, new ScriptQuery(f.source(), f.asScript()));
+            return handler.wrapFunctionQuery(f, f, () -> new ScriptQuery(f.source(), f.asScript()));
         }
     }
 

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/ExpressionTranslators.java

@@ -216,15 +216,17 @@ public final class ExpressionTranslators {
         }
 
         public static Query doTranslate(IsNotNull isNotNull, TranslatorHandler handler) {
-            Query query = null;
+            return handler.wrapFunctionQuery(isNotNull, isNotNull.field(), () -> translate(isNotNull, handler));
+        }
 
+        private static Query translate(IsNotNull isNotNull, TranslatorHandler handler) {
+            Query query = null;
             if (isNotNull.field() instanceof FieldAttribute) {
                 query = new ExistsQuery(isNotNull.source(), handler.nameOf(isNotNull.field()));
             } else {
                 query = new ScriptQuery(isNotNull.source(), isNotNull.asScript());
             }
-
-            return handler.wrapFunctionQuery(isNotNull, isNotNull.field(), query);
+            return query;
         }
     }
 
@@ -236,6 +238,10 @@ public final class ExpressionTranslators {
         }
 
         public static Query doTranslate(IsNull isNull, TranslatorHandler handler) {
+            return handler.wrapFunctionQuery(isNull, isNull.field(), () -> translate(isNull, handler));
+        }
+
+        private static Query translate(IsNull isNull, TranslatorHandler handler) {
             Query query = null;
 
             if (isNull.field() instanceof FieldAttribute) {
@@ -244,7 +250,7 @@ public final class ExpressionTranslators {
                 query = new ScriptQuery(isNull.source(), isNull.asScript());
             }
 
-            return handler.wrapFunctionQuery(isNull, isNull.field(), query);
+            return query;
         }
     }
 
@@ -265,7 +271,7 @@ public final class ExpressionTranslators {
 
         public static Query doTranslate(BinaryComparison bc, TranslatorHandler handler) {
             checkBinaryComparison(bc);
-            return handler.wrapFunctionQuery(bc, bc.left(), translate(bc, handler));
+            return handler.wrapFunctionQuery(bc, bc.left(), () -> translate(bc, handler));
         }
 
         private static Query translate(BinaryComparison bc, TranslatorHandler handler) {
@@ -340,10 +346,11 @@ public final class ExpressionTranslators {
             return doTranslate(r, handler);
         }
 
-        public static Query doTranslate(Range r, TranslatorHandler handler) {
-            Expression val = r.value();
+        public static Query doTranslate(Range r, TranslatorHandler handler) {            
+            return handler.wrapFunctionQuery(r, r.value(), () -> translate(r, handler));
+        }
 
-            Query query = null;
+        private static RangeQuery translate(Range r, TranslatorHandler handler) {            
             Object lower = valueOf(r.lower());
             Object upper = valueOf(r.upper());
             String format = null;
@@ -368,10 +375,8 @@ public final class ExpressionTranslators {
                 }
                 format = formatter.pattern();
             }
-            query = handler.wrapFunctionQuery(r, val,
-                    new RangeQuery(r.source(), handler.nameOf(val), lower, r.includeLower(), upper, r.includeUpper(), format, r.zoneId()));
-            
-            return query;
+            return new RangeQuery(
+                r.source(), handler.nameOf(r.value()), lower, r.includeLower(), upper, r.includeUpper(), format, r.zoneId());
         }
     }
 
@@ -383,6 +388,10 @@ public final class ExpressionTranslators {
         }
 
         public static Query doTranslate(In in, TranslatorHandler handler) {
+            return handler.wrapFunctionQuery(in, in.value(), () -> translate(in, handler));
+        }
+
+        private static Query translate(In in, TranslatorHandler handler) {
             Query q;
             if (in.value() instanceof FieldAttribute) {
                 // equality should always be against an exact match (which is important for strings)
@@ -406,8 +415,9 @@ public final class ExpressionTranslators {
                         assert o instanceof ZonedDateTime : "expected a ZonedDateTime, but got: " + o.getClass().getName();
                         // see comment in Ranges#doTranslate() as to why formatting as String is required
                         String zdt = formatter.format((ZonedDateTime) o);
-                        RangeQuery right = new RangeQuery(in.source(), fa.exactAttribute().name(),
-                                zdt, true, zdt, true, formatter.pattern(), in.zoneId());
+                        RangeQuery right = new RangeQuery(
+                            in.source(), fa.exactAttribute().name(),
+                            zdt, true, zdt, true, formatter.pattern(), in.zoneId());
                         q = q == null ? right : new BoolQuery(in.source(), false, q, right);
                     }
                 } else {
@@ -416,7 +426,7 @@ public final class ExpressionTranslators {
             } else {
                 q = new ScriptQuery(in.source(), in.asScript());
             }
-            return handler.wrapFunctionQuery(in, in.value(), q);
+            return q;
         }
     }
 
@@ -432,7 +442,7 @@ public final class ExpressionTranslators {
             if (q != null) {
                 return q;
             }
-            return handler.wrapFunctionQuery(f, f, new ScriptQuery(f.source(), f.asScript()));
+            return handler.wrapFunctionQuery(f, f, () -> new ScriptQuery(f.source(), f.asScript()));
         }
 
         public static Query doKnownTranslate(ScalarFunction f, TranslatorHandler handler) {

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/QlTranslatorHandler.java

@@ -8,11 +8,8 @@
 package org.elasticsearch.xpack.ql.planner;
 
 import org.elasticsearch.xpack.ql.expression.Expression;
-import org.elasticsearch.xpack.ql.expression.FieldAttribute;
 import org.elasticsearch.xpack.ql.expression.NamedExpression;
-import org.elasticsearch.xpack.ql.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
-import org.elasticsearch.xpack.ql.querydsl.query.ScriptQuery;
 import org.elasticsearch.xpack.ql.type.DataType;
 import org.elasticsearch.xpack.ql.type.DataTypeConverter;
 
@@ -23,14 +20,6 @@ public class QlTranslatorHandler implements TranslatorHandler {
         return ExpressionTranslators.toQuery(e, this);
     }
 
-    @Override
-    public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Query q) {
-        if (field instanceof FieldAttribute) {
-            return ExpressionTranslator.wrapIfNested(q, field);
-        }
-        return new ScriptQuery(sf.source(), sf.asScript());
-    }
-
     @Override
     public String nameOf(Expression e) {
         if (e instanceof NamedExpression) {

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/TranslatorHandler.java

@@ -8,10 +8,14 @@
 package org.elasticsearch.xpack.ql.planner;
 
 import org.elasticsearch.xpack.ql.expression.Expression;
+import org.elasticsearch.xpack.ql.expression.FieldAttribute;
 import org.elasticsearch.xpack.ql.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
+import org.elasticsearch.xpack.ql.querydsl.query.ScriptQuery;
 import org.elasticsearch.xpack.ql.type.DataType;
 
+import java.util.function.Supplier;
+
 /**
  * Parameterized handler used during query translation.
  *
@@ -21,7 +25,12 @@ public interface TranslatorHandler {
 
     Query asQuery(Expression e);
 
-    Query wrapFunctionQuery(ScalarFunction sf, Expression field, Query q);
+    default Query wrapFunctionQuery(ScalarFunction sf, Expression field, Supplier<Query> querySupplier) {
+        if (field instanceof FieldAttribute) {
+            return ExpressionTranslator.wrapIfNested(querySupplier.get(), field);
+        }
+        return new ScriptQuery(sf.source(), sf.asScript());
+    }
 
     String nameOf(Expression e);
 

x-pack/plugin/ql/src/test/resources/mapping-multi-field-variation.json

@@ -48,6 +48,12 @@
         },
         "foo_type" : { "type" : "foo" },
         "point": {"type" : "geo_point"},
-        "shape": {"type" : "geo_shape"}
+        "shape": {"type" : "geo_shape"},
+        "nested": {
+            "type": "nested",
+            "properties": {
+                "point": {"type" : "geo_point"}
+            }
+        }
     }
 }

x-pack/plugin/ql/src/test/resources/mapping-multi-field-with-nested.json

@@ -88,6 +88,9 @@
                 },
                 "start_date" : {
                     "type" : "date"
+                },
+                "location" : {
+                    "type" : "geo_point"
                 }
             }
         }

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

@@ -30,6 +30,7 @@ import org.elasticsearch.xpack.ql.expression.predicate.operator.comparison.Binar
 import org.elasticsearch.xpack.ql.expression.predicate.operator.comparison.LessThan;
 import org.elasticsearch.xpack.ql.expression.predicate.operator.comparison.LessThanOrEqual;
 import org.elasticsearch.xpack.ql.expression.predicate.regex.RegexMatch;
+import org.elasticsearch.xpack.ql.planner.ExpressionTranslator;
 import org.elasticsearch.xpack.ql.planner.ExpressionTranslators;
 import org.elasticsearch.xpack.ql.planner.TranslatorHandler;
 import org.elasticsearch.xpack.ql.querydsl.query.GeoDistanceQuery;
@@ -407,8 +408,9 @@ final class QueryTranslator {
                             Geometry geometry = ((GeoShape) geoShape).toGeometry();
                             if (geometry instanceof Point) {
                                 String field = nameOf(stDistance.left());
-                                return new GeoDistanceQuery(source, field, ((Number) value).doubleValue(),
+                                Query query = new GeoDistanceQuery(source, field, ((Number) value).doubleValue(),
                                     ((Point) geometry).getY(), ((Point) geometry).getX());
+                                return ExpressionTranslator.wrapIfNested(query, stDistance.left());
                             }
                         }
                     }

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

@@ -8,15 +8,9 @@
 package org.elasticsearch.xpack.sql.planner;
 
 import org.elasticsearch.xpack.ql.expression.Expression;
-import org.elasticsearch.xpack.ql.expression.FieldAttribute;
-import org.elasticsearch.xpack.ql.expression.function.scalar.ScalarFunction;
-import org.elasticsearch.xpack.ql.planner.ExpressionTranslator;
 import org.elasticsearch.xpack.ql.planner.TranslatorHandler;
-import org.elasticsearch.xpack.ql.querydsl.query.GeoDistanceQuery;
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
-import org.elasticsearch.xpack.ql.querydsl.query.ScriptQuery;
 import org.elasticsearch.xpack.ql.type.DataType;
-import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistance;
 import org.elasticsearch.xpack.sql.type.SqlDataTypeConverter;
 
 public class SqlTranslatorHandler implements TranslatorHandler {
@@ -32,17 +26,6 @@ public class SqlTranslatorHandler implements TranslatorHandler {
         return QueryTranslator.toQuery(e, onAggs).query;
     }
 
-    @Override
-    public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Query q) {
-        if (field instanceof StDistance && q instanceof GeoDistanceQuery) {
-            return ExpressionTranslator.wrapIfNested(q, ((StDistance) field).left());
-        }
-        if (field instanceof FieldAttribute) {
-            return ExpressionTranslator.wrapIfNested(q, field);
-        }
-        return new ScriptQuery(sf.source(), sf.asScript());
-    }
-
     @Override
     public String nameOf(Expression e) {
         return QueryTranslator.nameOf(e);

x-pack/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerifierErrorMessagesTests.java

@@ -569,6 +569,10 @@ public class VerifierErrorMessagesTests extends ESTestCase {
         accept("SELECT int FROM test WHERE dep.start_date > '2020-01-30'::date AND (int > 10 OR dep.end_date IS NULL)");
         accept("SELECT int FROM test WHERE dep.start_date > '2020-01-30'::date AND (int > 10 OR dep.end_date IS NULL) " +
                "OR NOT(dep.start_date >= '2020-01-01')");
+        String operator = randomFrom("<", "<=");
+        assertEquals("1:42: WHERE isn't (yet) compatible with scalar functions on nested fields [dep.location]",
+            error("SELECT shape FROM test " +
+                "WHERE ST_Distance(dep.location, ST_WKTToSQL('point (10 20)')) " + operator + " 25"));
     }
 
     public void testOrderByOnNested() {

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

@@ -1334,6 +1334,10 @@ public class QueryTranslatorTests extends ESTestCase {
         assertEquals(20.0, gq.lat(), 0.00001);
         assertEquals(10.0, gq.lon(), 0.00001);
         assertEquals(25.0, gq.distance(), 0.00001);
+        EsQueryExec eqe = (EsQueryExec) optimizeAndPlan(p);
+        assertThat(eqe.queryContainer().toString().replaceAll("\\s+", ""),
+            containsString("{\"geo_distance\":{\"point\":[10.0,20.0],\"distance\":25.0," +
+                "\"distance_type\":\"arc\",\"validation_method\":\"STRICT\","));
     }
 
     public void testTranslateStXY() {