ESQL: continue resolving attributes for Eval (elastic#99601)

nik9000 · Sep 19, 2023 · f04af72 · f04af72
1 parent 53f2eef
commit f04af72
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 4 deletions.
diff --git a/docs/changelog/99601.yaml b/docs/changelog/99601.yaml
@@ -0,0 +1,6 @@
+pr: 99601
+summary: "ESQL: continue resolving attributes for Eval"
+area: ES|QL
+type: bug
+issues:
+ - 99576
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/eval.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/eval.csv-spec
@@ -166,3 +166,36 @@ emp_no:i | sum:i
 10002 | 3230647
 10003 | 3230970
 ; 
+
+chainedEvalReusingPreviousValue
+from employees | sort emp_no | eval x1 = concat(first_name, "."), x2 = concat(x1, "."), x3 = concat(x2, ".") | keep x*, first_name | limit 5;
+
+   x1:keyword  |   x2:keyword  |  x3:keyword   |first_name:keyword   
+Georgi.        |Georgi..       |Georgi...      |Georgi         
+Bezalel.       |Bezalel..      |Bezalel...     |Bezalel        
+Parto.         |Parto..        |Parto...       |Parto          
+Chirstian.     |Chirstian..    |Chirstian...   |Chirstian      
+Kyoichi.       |Kyoichi..      |Kyoichi...     |Kyoichi
+;
+
+chainedEvalReusingPreviousValue2
+from employees | sort emp_no | eval x1 = concat(first_name, "."), x2 = concat(x1, last_name), x3 = concat(x2, gender) | keep x*, first_name, gender | limit 5;
+
+   x1:keyword  |   x2:keyword    |  x3:keyword      |first_name:keyword|gender:keyword
+Georgi.        |Georgi.Facello   |Georgi.FacelloM   |Georgi            |M              
+Bezalel.       |Bezalel.Simmel   |Bezalel.SimmelF   |Bezalel           |F              
+Parto.         |Parto.Bamford    |Parto.BamfordM    |Parto             |M              
+Chirstian.     |Chirstian.Koblick|Chirstian.KoblickM|Chirstian         |M              
+Kyoichi.       |Kyoichi.Maliniak |Kyoichi.MaliniakM |Kyoichi           |M
+;
+
+chainedEvalReusingPreviousValue3
+from employees | sort emp_no | eval x1 = concat(first_name, "."), x2 = concat(x1, last_name), x3 = concat(x2, x1) | keep x*, first_name | limit 5;
+
+   x1:keyword  |   x2:keyword    |        x3:keyword         |first_name:keyword
+Georgi.        |Georgi.Facello   |Georgi.FacelloGeorgi.      |Georgi         
+Bezalel.       |Bezalel.Simmel   |Bezalel.SimmelBezalel.     |Bezalel        
+Parto.         |Parto.Bamford    |Parto.BamfordParto.        |Parto          
+Chirstian.     |Chirstian.Koblick|Chirstian.KoblickChirstian.|Chirstian      
+Kyoichi.       |Kyoichi.Maliniak |Kyoichi.MaliniakKyoichi.   |Kyoichi
+;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java
@@ -323,13 +323,17 @@ protected LogicalPlan doRule(LogicalPlan plan) {
                 return resolveEnrich(p, childrenOutput);
             }
 
-            return plan.transformExpressionsUp(UnresolvedAttribute.class, ua -> resolveAttribute(ua, childrenOutput));
+            return plan.transformExpressionsUp(UnresolvedAttribute.class, ua -> maybeResolveAttribute(ua, childrenOutput));
         }
 
-        private Attribute resolveAttribute(UnresolvedAttribute ua, List<Attribute> childrenOutput) {
+        private Attribute maybeResolveAttribute(UnresolvedAttribute ua, List<Attribute> childrenOutput) {
             if (ua.customMessage()) {
                 return ua;
             }
+            return resolveAttribute(ua, childrenOutput);
+        }
+
+        private Attribute resolveAttribute(UnresolvedAttribute ua, List<Attribute> childrenOutput) {
             Attribute resolved = ua;
             var named = resolveAgainstList(ua, childrenOutput);
             // if resolved, return it; otherwise keep it in place to be resolved later
@@ -441,7 +445,7 @@ private LogicalPlan resolveRename(Rename rename, List<Attribute> childrenOutput)
                     // remove attributes overwritten by a renaming: `| keep a, b, c | rename a as b`
                     projections.removeIf(x -> x.name().equals(alias.name()));
 
-                    var resolved = resolveAttribute(ua, childrenOutput);
+                    var resolved = maybeResolveAttribute(ua, childrenOutput);
                     if (resolved instanceof UnsupportedAttribute || resolved.resolved()) {
                         var realiased = (NamedExpression) alias.replaceChildren(List.of(resolved));
                         projections.replaceAll(x -> x.equals(resolved) ? realiased : x);
@@ -490,7 +494,7 @@ private LogicalPlan resolveRename(Rename rename, List<Attribute> childrenOutput)
         private LogicalPlan resolveEnrich(Enrich enrich, List<Attribute> childrenOutput) {
 
             if (enrich.matchField().toAttribute() instanceof UnresolvedAttribute ua) {
-                Attribute resolved = resolveAttribute(ua, childrenOutput);
+                Attribute resolved = maybeResolveAttribute(ua, childrenOutput);
                 if (resolved.equals(ua)) {
                     return enrich;
                 }

diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java
@@ -33,6 +33,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.stream.IntStream;
 
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
 import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyze;
@@ -1254,6 +1255,24 @@ public void testEnrichExcludesPolicyKey() {
         assertThat(e.getMessage(), containsString("Unknown column [id]"));
     }
 
+    public void testChainedEvalFieldsUse() {
+        var query = "from test | eval x0 = pow(salary, 1), x1 = pow(x0, 2), x2 = pow(x1, 3)";
+        int additionalEvals = randomIntBetween(0, 5);
+        for (int i = 0, j = 3; i < additionalEvals; i++, j++) {
+            query += ", x" + j + " = pow(x" + (j - 1) + ", " + i + ")";
+        }
+        assertProjection(query + " | keep x*", IntStream.range(0, additionalEvals + 3).mapToObj(v -> "x" + v).toArray(String[]::new));
+    }
+
+    public void testMissingAttributeException_InChainedEval() {
+        var e = expectThrows(VerificationException.class, () -> analyze("""
+            from test
+            | eval x1 = concat(first_name, "."), x2 = concat(x1, last_name), x3 = concat(x2, x1), x4 = concat(x3, x5)
+            | keep x*
+            """));
+        assertThat(e.getMessage(), containsString("Unknown column [x5], did you mean any of [x1, x2, x3]?"));
+    }
+
     private void verifyUnsupported(String query, String errorMessage) {
         verifyUnsupported(query, errorMessage, "mapping-multi-field-variation.json");
     }