SQL: Enhance detection of scalars on nested fields

Enhance detection and validation of scalars on top of nested fields
used in WHERE, HAVING and ORDER BY. Previously, in case of aliases
or ordinal references the usage of scalars on nested fields was not
detected.

Follows: ea4504c

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Verifier.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.xpack.sql.expression.AttributeSet;
 import org.elasticsearch.xpack.sql.expression.Exists;
 import org.elasticsearch.xpack.sql.expression.Expression;
+import org.elasticsearch.xpack.sql.expression.ExpressionId;
 import org.elasticsearch.xpack.sql.expression.Expressions;
 import org.elasticsearch.xpack.sql.expression.FieldAttribute;
 import org.elasticsearch.xpack.sql.expression.NamedExpression;
@@ -699,13 +700,33 @@ private static void checkForScoreInsideFunctions(LogicalPlan p, Set<Failure> loc
     }
 
     private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan p, Set<Failure> localFailures) {
+        // collect Attribute sources
+        // only Aliases are interesting since these are the only ones that hide expressions
+        // FieldAttribute for example are self replicating.
+        final Map<ExpressionId, Expression> collectRefs = new LinkedHashMap<>();
+        p.forEachUp(plan -> plan.forEachExpressionsUp(e -> {
+            if (e instanceof Alias) {
+                Alias a = (Alias) e;
+                collectRefs.put(a.id(), a.child());
+            }
+        }));
+        p.forEachDown(project -> project.projections().forEach(e -> collectRefs.putIfAbsent(e.id(), e)), Project.class);
+
         List<FieldAttribute> nested = new ArrayList<>();
         Consumer<FieldAttribute> matchNested = fa -> {
             if (fa.isNested()) {
                 nested.add(fa);
             }
         };
-        Consumer<Expression> checkForNested = e -> e.forEachUp(matchNested, FieldAttribute.class);
+        Consumer<Expression> checkForNested = e -> {
+            Expression expr = e;
+            if (e instanceof NamedExpression) {
+                expr = collectRefs.getOrDefault(((NamedExpression) e).id(), e);
+
+            }
+            expr.forEachUp(matchNested, FieldAttribute.class);
+        };
+
         Consumer<ScalarFunction> checkForNestedInFunction =  f -> f.forEachDown(
                 arg -> arg.forEachUp(matchNested, FieldAttribute.class));
 
@@ -744,8 +765,14 @@ private static void checkNestedUsedInGroupByOrHavingOrWhereOrOrderBy(LogicalPlan
         }
 
         // check in order by (scalars not allowed)
-        p.forEachDown(ob -> ob.order().forEach(o -> o.forEachUp(e ->
-                e.forEachUp(checkForNestedInFunction, ScalarFunction.class)
+        p.forEachDown(ob -> ob.order().forEach(o -> o.forEachUp(e -> {
+                    Expression expr = e;
+                    if (e instanceof NamedExpression) {
+                        expr = collectRefs.getOrDefault(((NamedExpression) e).id(), e);
+
+                    }
+                    expr.forEachUp(checkForNestedInFunction, ScalarFunction.class);
+                }
         )), OrderBy.class);
         if (!nested.isEmpty()) {
             localFailures.add(

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

@@ -342,22 +342,33 @@ public void testGroupByOnNested() {
     public void testHavingOnNested() {
         assertEquals("1:51: HAVING isn't (yet) compatible with nested fields [dep.start_date]",
                 error("SELECT int FROM test GROUP BY int HAVING AVG(YEAR(dep.start_date)) > 1980"));
+        assertEquals("1:22: HAVING isn't (yet) compatible with nested fields [dep.start_date]",
+                error("SELECT int, AVG(YEAR(dep.start_date)) AS average FROM test GROUP BY int HAVING average > 1980"));
+        assertEquals("1:22: HAVING isn't (yet) compatible with nested fields [dep.start_date, dep.end_date]",
+                error("SELECT int, AVG(YEAR(dep.start_date)) AS a, MAX(MONTH(dep.end_date)) AS b " +
+                        "FROM test GROUP BY int " +
+                        "HAVING a > 1980 AND b < 10"));
     }
 
     public void testWhereOnNested() {
         assertEquals("1:33: WHERE isn't (yet) compatible with scalar functions on nested fields [dep.start_date]",
                 error("SELECT int FROM test WHERE YEAR(dep.start_date) + 10 > 0"));
         assertEquals("1:13: WHERE isn't (yet) compatible with scalar functions on nested fields [dep.start_date]",
                 error("SELECT YEAR(dep.start_date) + 10 AS a FROM test WHERE int > 10 AND (int < 3 OR NOT (a > 5))"));
-        accept("SELECT int FROM test WHERE dep.start_date > CAST('2020-01-30' AS datetime) AND " +
-                "(int > 10 OR dep.end_date IS NULL)");
-        accept("SELECT int FROM test WHERE dep.start_date > CAST('2020-01-30'as datetime) AND " +
+        accept("SELECT int FROM test WHERE dep.start_date > CAST('2020-01-30'AS date) " +
+                "AND (int > 10 OR dep.end_date IS NULL)");
+        accept("SELECT int FROM test WHERE dep.start_date > CAST('2020-01-30' AS date) AND " +
                 "(int > 10 OR dep.end_date IS NULL) OR NOT(dep.start_date >= '2020-01-01')");
     }
 
     public void testOrderByOnNested() {
         assertEquals("1:36: ORDER BY isn't (yet) compatible with scalar functions on nested fields [dep.start_date]",
                 error("SELECT int FROM test ORDER BY YEAR(dep.start_date) + 10"));
+        assertEquals("1:13: ORDER BY isn't (yet) compatible with scalar functions on nested fields [dep.start_date]",
+                error("SELECT YEAR(dep.start_date) + 10  FROM test ORDER BY 1"));
+        assertEquals("1:13: ORDER BY isn't (yet) compatible with scalar functions on nested fields " +
+                        "[dep.start_date, dep.end_date]",
+                error("SELECT YEAR(dep.start_date) + 10 AS a, MONTH(dep.end_date) - 10 as b FROM test ORDER BY 1, 2"));
         accept("SELECT int FROM test ORDER BY dep.start_date, dep.end_date");
     }