From 6035c31642d02d1b643941a925f1d07ef53032e1 Mon Sep 17 00:00:00 2001
From: Andras Palinkas <andras.palinkas@elastic.co>
Date: Mon, 8 Feb 2021 20:30:56 -0500
Subject: [PATCH] QL: Turn eager query translations lazy

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
---
 .../xpack/eql/planner/QueryTranslator.java    |  4 +-
 .../ql/planner/ExpressionTranslators.java     | 42 ++++++++++++-------
 .../xpack/ql/planner/QlTranslatorHandler.java |  6 ++-
 .../xpack/ql/planner/TranslatorHandler.java   |  4 +-
 .../sql/planner/SqlTranslatorHandler.java     | 10 +++--
 5 files changed, 41 insertions(+), 25 deletions(-)

diff --git a/x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/planner/QueryTranslator.java b/x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/planner/QueryTranslator.java
index 8f2d66eb4aeb8..b5b422a71e684 100644
--- a/x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/planner/QueryTranslator.java
+++ b/x-pack/plugin/eql/src/main/java/org/elasticsearch/xpack/eql/planner/QueryTranslator.java
@@ -84,7 +84,7 @@ protected Query asQuery(InsensitiveBinaryComparison bc, TranslatorHandler handle
 
         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 @@ public static Query doTranslate(ScalarFunction f, TranslatorHandler handler) {
                 }
             }
 
-            return handler.wrapFunctionQuery(f, f, new ScriptQuery(f.source(), f.asScript()));
+            return handler.wrapFunctionQuery(f, f, () -> new ScriptQuery(f.source(), f.asScript()));
         }
     }
 
diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/ExpressionTranslators.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/ExpressionTranslators.java
index 8dcf398a55801..e301619dc0fda 100644
--- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/ExpressionTranslators.java
+++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/ExpressionTranslators.java
@@ -216,15 +216,17 @@ protected Query asQuery(IsNotNull isNotNull, TranslatorHandler handler) {
         }
 
         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 @@ protected Query asQuery(IsNull isNull, TranslatorHandler handler) {
         }
 
         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 static Query doTranslate(IsNull isNull, TranslatorHandler handler) {
                 query = new ScriptQuery(isNull.source(), isNull.asScript());
             }
 
-            return handler.wrapFunctionQuery(isNull, isNull.field(), query);
+            return query;
         }
     }
 
@@ -265,7 +271,7 @@ public static void checkBinaryComparison(BinaryComparison bc) {
 
         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 @@ protected Query asQuery(Range r, TranslatorHandler handler) {
             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 static Query doTranslate(Range r, TranslatorHandler handler) {
                 }
                 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 @@ protected Query asQuery(In in, TranslatorHandler handler) {
         }
 
         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 static Query doTranslate(In in, TranslatorHandler handler) {
                         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 static Query doTranslate(In in, TranslatorHandler handler) {
             } else {
                 q = new ScriptQuery(in.source(), in.asScript());
             }
-            return handler.wrapFunctionQuery(in, in.value(), q);
+            return q;
         }
     }
 
@@ -432,7 +442,7 @@ public static Query doTranslate(ScalarFunction f, TranslatorHandler handler) {
             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) {
diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/QlTranslatorHandler.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/QlTranslatorHandler.java
index 4430e7b0662bc..97bd5b7fc35df 100644
--- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/QlTranslatorHandler.java
+++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/QlTranslatorHandler.java
@@ -16,6 +16,8 @@
 import org.elasticsearch.xpack.ql.type.DataType;
 import org.elasticsearch.xpack.ql.type.DataTypeConverter;
 
+import java.util.function.Supplier;
+
 public class QlTranslatorHandler implements TranslatorHandler {
 
     @Override
@@ -24,9 +26,9 @@ public Query asQuery(Expression e) {
     }
 
     @Override
-    public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Query q) {
+    public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Supplier<Query> querySupplier) {
         if (field instanceof FieldAttribute) {
-            return ExpressionTranslator.wrapIfNested(q, field);
+            return ExpressionTranslator.wrapIfNested(querySupplier.get(), field);
         }
         return new ScriptQuery(sf.source(), sf.asScript());
     }
diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/TranslatorHandler.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/TranslatorHandler.java
index 1bcbbbd307eed..e3dff17553f8c 100644
--- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/TranslatorHandler.java
+++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/planner/TranslatorHandler.java
@@ -12,6 +12,8 @@
 import org.elasticsearch.xpack.ql.querydsl.query.Query;
 import org.elasticsearch.xpack.ql.type.DataType;
 
+import java.util.function.Supplier;
+
 /**
  * Parameterized handler used during query translation.
  *
@@ -21,7 +23,7 @@ public interface TranslatorHandler {
 
     Query asQuery(Expression e);
 
-    Query wrapFunctionQuery(ScalarFunction sf, Expression field, Query q);
+    Query wrapFunctionQuery(ScalarFunction sf, Expression field, Supplier<Query> querySupplier);
 
     String nameOf(Expression e);
 
diff --git a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/SqlTranslatorHandler.java b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/SqlTranslatorHandler.java
index d6386ebd84b9b..277e9da2bc06c 100644
--- a/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/SqlTranslatorHandler.java
+++ b/x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/planner/SqlTranslatorHandler.java
@@ -19,6 +19,8 @@
 import org.elasticsearch.xpack.sql.expression.function.scalar.geo.StDistance;
 import org.elasticsearch.xpack.sql.type.SqlDataTypeConverter;
 
+import java.util.function.Supplier;
+
 public class SqlTranslatorHandler implements TranslatorHandler {
 
     private final boolean onAggs;
@@ -33,12 +35,12 @@ public Query asQuery(Expression e) {
     }
 
     @Override
-    public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Query q) {
-        if (field instanceof StDistance && q instanceof GeoDistanceQuery) {
-            return ExpressionTranslator.wrapIfNested(q, ((StDistance) field).left());
+    public Query wrapFunctionQuery(ScalarFunction sf, Expression field, Supplier<Query> querySupplier) {
+        if (field instanceof StDistance && querySupplier.get() instanceof GeoDistanceQuery) {
+            return ExpressionTranslator.wrapIfNested(querySupplier.get(), ((StDistance) field).left());
         }
         if (field instanceof FieldAttribute) {
-            return ExpressionTranslator.wrapIfNested(q, field);
+            return ExpressionTranslator.wrapIfNested(querySupplier.get(), field);
         }
         return new ScriptQuery(sf.source(), sf.asScript());
     }