SQL: Failing Group By queries due to different ExpressionIds (#43072)

Fix an issue that arises from the use of ExpressionIds as keys in a lookup map
that helps the QueryTranslator to identify the grouping columns. The issue is
that the same expression in different parts of the query (SELECT clause and GROUP BY clause)
ends up with different ExpressionIds so the lookup fails. So, instead of ExpressionIds
use the hashCode() of NamedExpression.

Fixes: #41159
Fixes: #40001
Fixes: #40240
Fixes: #33361
Fixes: #46316
Fixes: #36074
Fixes: #34543
Fixes: #37044

Fixes: #42041
(cherry picked from commit 3c38ea5)

Author: emasab

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

@@ -244,6 +244,194 @@ TRUNCATE(YEAR("birth_date"), -2)
 null
 1900
 ;
+// Fails for H2
+groupByCastScalarWithNumericRef
+SELECT CAST(ABS(EXTRACT(YEAR FROM "birth_date")) AS BIGINT) FROM test_emp GROUP BY 1 ORDER BY 1 NULLS FIRST;
+
+CAST(ABS(EXTRACT(YEAR FROM "birth_date")) AS BIGINT):l
+------------------------------------------------------
+null
+1952
+1953
+1954
+1955
+1956
+1957
+1958
+1959
+1960
+1961
+1962
+1963
+1964
+1965
+;
+
+groupByConvertScalar
+SELECT CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT) FROM test_emp GROUP BY CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT) ORDER BY CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT) NULLS FIRST;
+
+
+CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT):l
+-----------------------------------------------------------
+null
+1952
+1953
+1954
+1955
+1956
+1957
+1958
+1959
+1960
+1961
+1962
+1963
+1964
+1965
+;
+
+
+groupByConvertScalarWithAlias
+SELECT CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT) as "convert" FROM test_emp GROUP BY "convert" ORDER BY "convert" NULLS FIRST;
+
+convert:l
+---------
+null
+1952
+1953
+1954
+1955
+1956
+1957
+1958
+1959
+1960
+1961
+1962
+1963
+1964
+1965
+;
+
+groupByConvertScalarWithNumericRef
+SELECT CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT) FROM test_emp GROUP BY 1 ORDER BY 1 NULLS FIRST;
+
+CONVERT(ABS(EXTRACT(YEAR FROM "birth_date")), SQL_BIGINT):l
+-----------------------------------------------------------
+null
+1952
+1953
+1954
+1955
+1956
+1957
+1958
+1959
+1960
+1961
+1962
+1963
+1964
+1965
+;
+
+groupByConstantScalar
+SELECT PI() * emp_no FROM test_emp GROUP BY PI() * emp_no ORDER BY PI() * emp_no LIMIT 10;
+
+PI() * emp_no:d
+---------------
+31419.0681285515
+31422.2097212051
+31425.3513138587
+31428.4929065123
+31431.6344991659
+31434.7760918195
+31437.9176844731
+31441.0592771266
+31444.2008697802
+31447.3424624338
+;
+
+groupByConstantScalarWithOrderByDesc
+SELECT PI() * emp_no FROM test_emp GROUP BY PI() * emp_no ORDER BY PI() * emp_no DESC LIMIT 10;
+
+PI() * emp_no:d
+-------
+31730.0858012569
+31726.9442086033
+31723.8026159497
+31720.6610232961
+31717.5194306425
+31714.3778379889
+31711.2362453353
+31708.0946526817
+31704.9530600281
+31701.8114673746
+;
+
+groupByConstantScalarWithAlias
+SELECT PI() * emp_no AS "value" FROM test_emp GROUP BY value ORDER BY value LIMIT 10;
+
+value:d
+-------
+31419.0681285515
+31422.2097212051
+31425.3513138587
+31428.4929065123
+31431.6344991659
+31434.7760918195
+31437.9176844731
+31441.0592771266
+31444.2008697802
+31447.3424624338
+;
+
+groupByConstantScalarWithNumericRef
+SELECT PI() * emp_no FROM test_emp GROUP BY 1 ORDER BY 1 DESC LIMIT 10;
+
+PI() * emp_no:d
+-------
+31730.0858012569
+31726.9442086033
+31723.8026159497
+31720.6610232961
+31717.5194306425
+31714.3778379889
+31711.2362453353
+31708.0946526817
+31704.9530600281
+31701.8114673746
+;
+
+groupByFieldAndConstantScalarWithMultipleOrderBy
+SELECT gender, emp_no % 3 + PI() FROM test_emp GROUP BY gender, emp_no % 3 + PI() ORDER BY gender, emp_no % 3 + PI() DESC LIMIT 8;
+
+gender:s    |emp_no % 3 + PI():d
+------------+------------------
+null        |5.1415926535
+null        |4.1415926535
+null        |3.1415926535
+F           |5.1415926535
+F           |4.1415926535
+F           |3.1415926535
+M           |5.1415926535
+M           |4.1415926535
+;
+
+groupByFieldAndConstantScalarWithAliasWithOrderByDesc
+SELECT gender, emp_no % 3 + PI() as p FROM test_emp GROUP BY gender, emp_no % 3 + PI() ORDER BY gender DESC, p DESC LIMIT 8;
+
+gender:s    |p:d
+------------+------------------
+M           |5.1415926535
+M           |4.1415926535
+M           |3.1415926535
+F           |5.1415926535
+F           |4.1415926535
+F           |3.1415926535
+null        |5.1415926535
+null        |4.1415926535
+;
 
 //
 // Grouping functions

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

@@ -49,6 +49,10 @@ groupByMulScalar
 SELECT emp_no * 2 AS e FROM test_emp GROUP BY e ORDER BY e;
 groupByModScalar
 SELECT (emp_no % 3) + 1 AS e FROM test_emp GROUP BY e ORDER BY e;
+groupByCastScalar
+SELECT CAST(ABS(EXTRACT(YEAR FROM "birth_date")) AS BIGINT) FROM test_emp GROUP BY CAST(ABS(EXTRACT(YEAR FROM "birth_date")) AS BIGINT) ORDER BY CAST(ABS(EXTRACT(YEAR FROM "birth_date")) AS BIGINT) NULLS FIRST;
+groupByCastScalarWithAlias
+SELECT CAST(ABS(EXTRACT(YEAR FROM "birth_date")) AS BIGINT) as "cast" FROM test_emp GROUP BY "cast" ORDER BY "cast" NULLS FIRST;
 
 // group by nested functions with no alias
 groupByTruncate

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

@@ -65,6 +65,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
@@ -625,12 +626,15 @@ protected LogicalPlan rule(LogicalPlan plan) {
                         .map(or -> tryResolveExpression(or, o.child()))
                         .collect(toList());
 
-                AttributeSet resolvedRefs = Expressions.references(maybeResolved.stream()
-                        .filter(Expression::resolved)
-                        .collect(toList()));
 
+                Set<Expression> resolvedRefs = maybeResolved.stream()
+                    .filter(Expression::resolved)
+                    .collect(Collectors.toSet());
 
-                AttributeSet missing = resolvedRefs.subtract(o.child().outputSet());
+                AttributeSet missing = Expressions.filterReferences(
+                    resolvedRefs,
+                    o.child().outputSet()
+                );
 
                 if (!missing.isEmpty()) {
                     // Add missing attributes but project them away afterwards

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

@@ -27,7 +27,7 @@ protected VerificationException(Collection<Failure> sources) {
     public String getMessage() {
         return failures.stream()
                 .map(f -> {
-                    Location l = f.source().source().source();
+                    Location l = f.node().source().source();
                     return "line " + l.getLineNumber() + ":" + l.getColumnNumber() + ": " + f.message();
                 })
                 .collect(Collectors.joining(StringUtils.NEW_LINE, "Found " + failures.size() + " problem(s)\n", StringUtils.EMPTY));

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

@@ -77,16 +77,16 @@ public Verifier(Metrics metrics) {
     }
 
     static class Failure {
-        private final Node<?> source;
+        private final Node<?> node;
         private final String message;
 
-        Failure(Node<?> source, String message) {
-            this.source = source;
+        Failure(Node<?> node, String message) {
+            this.node = node;
             this.message = message;
         }
 
-        Node<?> source() {
-            return source;
+        Node<?> node() {
+            return node;
         }
 
         String message() {
@@ -95,7 +95,7 @@ String message() {
 
         @Override
         public int hashCode() {
-            return source.hashCode();
+            return Objects.hash(node);
         }
 
         @Override
@@ -109,7 +109,7 @@ public boolean equals(Object obj) {
             }
 
             Verifier.Failure other = (Verifier.Failure) obj;
-            return Objects.equals(source, other.source);
+            return Objects.equals(node, other.node);
         }
 
         @Override
@@ -124,7 +124,7 @@ private static Failure fail(Node<?> source, String message, Object... args) {
 
     public Map<Node<?>, String> verifyFailures(LogicalPlan plan) {
         Collection<Failure> failures = verify(plan);
-        return failures.stream().collect(toMap(Failure::source, Failure::message));
+        return failures.stream().collect(toMap(Failure::node, Failure::message));
     }
 
     Collection<Failure> verify(LogicalPlan plan) {

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

@@ -32,14 +32,14 @@ static class AttributeWrapper {
 
         @Override
         public int hashCode() {
-            return attr.semanticHash();
+            return attr.hashCode();
         }
 
         @Override
         public boolean equals(Object obj) {
             if (obj instanceof AttributeWrapper) {
                 AttributeWrapper aw = (AttributeWrapper) obj;
-                return attr.semanticEquals(aw.attr);
+                return attr.equals(aw.attr);
             }
 
             return false;
@@ -368,4 +368,4 @@ public boolean equals(Object obj) {
     public String toString() {
         return delegate.toString();
     }
-}
+}

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

@@ -126,9 +126,6 @@ public boolean resolved() {
 
     public abstract DataType dataType();
 
-    @Override
-    public abstract int hashCode();
-
     @Override
     public String toString() {
         return nodeName() + "[" + propertiesToString(false) + "]";

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

@@ -11,8 +11,11 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
 import static java.util.Collections.emptyMap;
@@ -99,6 +102,31 @@ public static AttributeSet references(List<? extends Expression> exps) {
         return set;
     }
 
+    public static AttributeSet filterReferences(Set<? extends Expression> exps, AttributeSet excluded) {
+        AttributeSet ret = new AttributeSet();
+        while (exps.size() > 0) {
+
+            Set<Expression> filteredExps = new LinkedHashSet<>();
+            for (Expression exp : exps) {
+                Expression attr = Expressions.attribute(exp);
+                if (attr == null || (excluded.contains(attr) == false)) {
+                    filteredExps.add(exp);
+                }
+            }
+
+            ret.addAll(new AttributeSet(
+                filteredExps.stream().filter(c->c.children().isEmpty())
+                .flatMap(exp->exp.references().stream())
+                .collect(Collectors.toSet())
+            ));
+
+            exps = filteredExps.stream()
+                .flatMap((Expression exp)->exp.children().stream())
+                .collect(Collectors.toSet());
+        }
+        return ret;
+    }
+
     public static String name(Expression e) {
         return e instanceof NamedExpression ? ((NamedExpression) e).name() : e.nodeName();
     }

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

@@ -97,11 +97,6 @@ private FieldAttribute innerField(EsField type) {
         return new FieldAttribute(source(), this, name() + "." + type.getName(), type, qualifier(), nullable(), id(), synthetic());
     }
 
-    @Override
-    protected Expression canonicalize() {
-        return new FieldAttribute(source(), null, "<none>", field, null, Nullability.TRUE, id(), false);
-    }
-
     @Override
     protected Attribute clone(Source source, String name, String qualifier, Nullability nullability,
                               ExpressionId id, boolean synthetic) {