#46396 fix npe on ambiguous group by

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/Attribute.java

@@ -25,7 +25,7 @@
  * is a named expression (an {@code Alias} will be created automatically for it).
  * The rest are not as they are not part of the projection and thus are not part of the derived table.
  */
-public abstract class Attribute extends NamedExpression {
+public abstract class Attribute extends NamedExpression implements Comparable<Attribute>{
 
     // empty - such as a top level attribute in SELECT cause
     // present - table name or a table name alias
@@ -142,4 +142,16 @@ public String nodeString() {
     }
 
     protected abstract String label();
+
+    @Override
+    public int compareTo(Attribute o) {
+        int result = this.sourceLocation().compareTo(o.sourceLocation());
+        if (result == 0) {
+            result = this.name().compareTo(o.name());
+        }
+        if (result == 0) {
+            result = this.qualifier.compareTo(o.qualifier);
+        }
+        return result;
+    }
 }

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/expression/AttributeAlias.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.ql.expression;
+
+import java.util.Objects;
+
+public class AttributeAlias {
+    private final Attribute attribute;
+    private final Expression expression;
+
+    public AttributeAlias(Attribute attribute, Expression expression) {
+        this.attribute = attribute;
+        this.expression = expression;
+    }
+
+    public Attribute getAttribute() {
+        return attribute;
+    }
+
+    public Expression getExpression() {
+        return expression;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        AttributeAlias that = (AttributeAlias) o;
+        return Objects.equals(attribute, that.attribute) &&
+            Objects.equals(expression, that.expression);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(attribute, expression);
+    }
+
+    @Override
+    public String toString() {
+        return "Alias {" + attribute.toString() + " -> " + expression.toString() + "}";
+    }
+}

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

@@ -14,10 +14,8 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.function.Predicate;
 
@@ -160,14 +158,14 @@ public static boolean equalsAsAttribute(Expression left, Expression right) {
         return true;
     }
 
-    public static AttributeMap<Expression> aliases(List<? extends NamedExpression> named) {
-        Map<Attribute, Expression> aliasMap = new LinkedHashMap<>();
+    public static List<AttributeAlias> aliases(List<? extends NamedExpression> named) {
+        List<AttributeAlias> aliases = new ArrayList<>();
         for (NamedExpression ne : named) {
             if (ne instanceof Alias) {
-                aliasMap.put(ne.toAttribute(), ((Alias) ne).child());
+                aliases.add(new AttributeAlias(ne.toAttribute(), ((Alias) ne).child()));
             }
         }
-        return new AttributeMap<>(aliasMap);
+        return aliases;
     }
 
     public static boolean hasReferenceAttribute(Collection<Attribute> output) {

x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/tree/Location.java

@@ -7,7 +7,7 @@
 
 import java.util.Objects;
 
-public final class Location {
+public final class Location implements Comparable<Location>{
     private final int line;
     private final int charPositionInLine;
 
@@ -26,6 +26,15 @@ public int getColumnNumber() {
         return charPositionInLine + 1;
     }
 
+    @Override
+    public int compareTo(Location other) {
+        int result = this.line - other.line;
+        if (result == 0) {
+            result = this.charPositionInLine - other.charPositionInLine;
+        }
+        return result;
+    }
+
     @Override
     public String toString() {
         return "@" + getLineNumber() + ":" + getColumnNumber();

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

@@ -10,7 +10,7 @@
 import org.elasticsearch.xpack.ql.common.Failure;
 import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
-import org.elasticsearch.xpack.ql.expression.AttributeMap;
+import org.elasticsearch.xpack.ql.expression.AttributeAlias;
 import org.elasticsearch.xpack.ql.expression.AttributeSet;
 import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
@@ -41,6 +41,7 @@
 import org.elasticsearch.xpack.ql.rule.Rule;
 import org.elasticsearch.xpack.ql.rule.RuleExecutor;
 import org.elasticsearch.xpack.ql.session.Configuration;
+import org.elasticsearch.xpack.ql.tree.Location;
 import org.elasticsearch.xpack.ql.type.DataType;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 import org.elasticsearch.xpack.ql.type.InvalidMappedField;
@@ -181,7 +182,7 @@ private static Attribute resolveAgainstList(UnresolvedAttribute u, Collection<At
                              // but also if the qualifier might not be quoted and if there's any ambiguity with nested fields
                              || Objects.equals(u.name(), attribute.qualifiedName()));
                 if (match) {
-                    matches.add(attribute.withLocation(u.source()));
+                    matches.add(attribute);
                 }
             }
         }
@@ -192,16 +193,24 @@ private static Attribute resolveAgainstList(UnresolvedAttribute u, Collection<At
         }
 
         if (matches.size() == 1) {
-            return handleSpecialFields(u, matches.get(0), allowCompound);
+            Attribute match = matches.get(0).withLocation(u.source());
+            return handleSpecialFields(u, match, allowCompound);
         }
 
         return u.withUnresolvedMessage("Reference [" + u.qualifiedName()
                 + "] is ambiguous (to disambiguate use quotes or qualifiers); matches any of " +
                  matches.stream()
-                 .map(a -> "\"" + a.qualifier() + "\".\"" + a.name() + "\"")
                  .sorted()
+                 .map(Analyzer::buildUnresolvedMessagesMatchesList)
                  .collect(toList())
                 );
+
+    }
+
+    private static String buildUnresolvedMessagesMatchesList(Attribute a) {
+        Location location = a.sourceLocation();
+        String locationString = "line " + location.getLineNumber() + ":" + location.getColumnNumber();
+        return locationString + " [" + (a.qualifier() != null ? "\"" + a.qualifier() + "\"." : "") + "\"" + a.name() + "\"]";
     }
 
     private static Attribute handleSpecialFields(UnresolvedAttribute u, Attribute named, boolean allowCompound) {
@@ -333,16 +342,25 @@ else if (plan instanceof Aggregate) {
                 if (!a.expressionsResolved() && Resolvables.resolved(a.aggregates())) {
                     List<Expression> groupings = a.groupings();
                     List<Expression> newGroupings = new ArrayList<>();
-                    AttributeMap<Expression> resolved = Expressions.aliases(a.aggregates());
+                    List<AttributeAlias> resolved = Expressions.aliases(a.aggregates());
 
                     boolean changed = false;
                     for (Expression grouping : groupings) {
                         if (grouping instanceof UnresolvedAttribute) {
-                            Attribute maybeResolved = resolveAgainstList((UnresolvedAttribute) grouping, resolved.keySet());
+                            Attribute maybeResolved = resolveAgainstList((UnresolvedAttribute) grouping,
+                                resolved.stream().map(AttributeAlias::getAttribute).collect(toList()));
                             if (maybeResolved != null) {
                                 changed = true;
-                                // use the matched expression (not its attribute)
-                                grouping = resolved.get(maybeResolved);
+                                if (maybeResolved.resolved()) {
+                                    // use the matched expression (not its attribute)
+                                    grouping = resolved.stream()
+                                        .filter(attributeAlias -> attributeAlias.getAttribute().equals(maybeResolved))
+                                        .map(AttributeAlias::getExpression)
+                                        .findAny()
+                                        .orElse(maybeResolved);
+                                } else {
+                                    grouping = maybeResolved;
+                                }
                             }
                         }
                         newGroupings.add(grouping);

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

@@ -5,15 +5,19 @@
  */
 package org.elasticsearch.xpack.sql.analysis.analyzer;
 
+import java.util.stream.Collectors;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.ql.QlIllegalArgumentException;
+import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
+import org.elasticsearch.xpack.ql.expression.Expression;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.FieldAttribute;
 import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.expression.function.FunctionRegistry;
 import org.elasticsearch.xpack.ql.index.EsIndex;
 import org.elasticsearch.xpack.ql.index.IndexResolution;
+import org.elasticsearch.xpack.ql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.ql.plan.logical.Project;
 import org.elasticsearch.xpack.ql.type.EsField;
@@ -31,6 +35,7 @@
 import static org.elasticsearch.xpack.ql.type.DataTypes.TEXT;
 import static org.elasticsearch.xpack.sql.types.SqlTypesTests.loadMapping;
 import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.hasItem;
 import static org.hamcrest.Matchers.hasItems;
 import static org.hamcrest.Matchers.hasSize;
@@ -178,13 +183,13 @@ public void testFieldAmbiguity() {
         VerificationException ex = expectThrows(VerificationException.class, () -> plan("SELECT test.bar FROM test"));
         assertEquals(
                 "Found 1 problem\nline 1:8: Reference [test.bar] is ambiguous (to disambiguate use quotes or qualifiers); "
-                        + "matches any of [\"test\".\"bar\", \"test\".\"test.bar\"]",
+                        + "matches any of [line 1:22 [\"test\".\"bar\"], line 1:22 [\"test\".\"test.bar\"]]",
                 ex.getMessage());
 
         ex = expectThrows(VerificationException.class, () -> plan("SELECT test.test FROM test"));
         assertEquals(
                 "Found 1 problem\nline 1:8: Reference [test.test] is ambiguous (to disambiguate use quotes or qualifiers); "
-                        + "matches any of [\"test\".\"test\", \"test\".\"test.test\"]",
+                        + "matches any of [line 1:23 [\"test\".\"test\"], line 1:23 [\"test\".\"test.test\"]]",
                 ex.getMessage());
 
         LogicalPlan plan = plan("SELECT test.test FROM test AS x");
@@ -201,4 +206,65 @@ public void testFieldAmbiguity() {
         assertThat(attribute.qualifier(), is("test"));
         assertThat(attribute.name(), is("test.test"));
     }
+
+    public void testAggregations() {
+        Map<String, EsField> mapping = TypesTests.loadMapping("mapping-basic.json");
+        EsIndex index = new EsIndex("test", mapping);
+        getIndexResult = IndexResolution.valid(index);
+        analyzer = new Analyzer(SqlTestUtils.TEST_CFG, functionRegistry, getIndexResult, verifier);
+
+        LogicalPlan plan = plan("SELECT sum(salary) AS s FROM test");
+        assertThat(plan, instanceOf(Aggregate.class));
+
+        Aggregate aggregate = (Aggregate) plan;
+        assertThat(aggregate.aggregates(), hasSize(1));
+        NamedExpression attribute = aggregate.aggregates().get(0);
+        assertThat(attribute, instanceOf(Alias.class));
+        assertThat(attribute.name(), is("s"));
+        assertThat(aggregate.groupings(), hasSize(0));
+
+        plan = plan("SELECT gender AS g, sum(salary) AS s FROM test GROUP BY g");
+        assertThat(plan, instanceOf(Aggregate.class));
+
+        aggregate = (Aggregate) plan;
+        List<? extends NamedExpression> aggregates = aggregate.aggregates();
+        assertThat(aggregates, hasSize(2));
+        assertThat(aggregates.get(0), instanceOf(Alias.class));
+        assertThat(aggregates.get(1), instanceOf(Alias.class));
+        List<String> names = aggregate.aggregates().stream().map(NamedExpression::name).collect(Collectors.toList());
+        assertThat(names, contains("g", "s"));
+
+        List<Expression> groupings = aggregate.groupings();
+        assertThat(groupings, hasSize(1));
+        FieldAttribute grouping = (FieldAttribute) groupings.get(0);
+        assertThat(grouping.name(), is("gender"));
+    }
+
+    public void testGroupByAmbiguity() {
+        Map<String, EsField> mapping = TypesTests.loadMapping("mapping-basic.json");
+        EsIndex index = new EsIndex("test", mapping);
+        getIndexResult = IndexResolution.valid(index);
+        analyzer = new Analyzer(SqlTestUtils.TEST_CFG, functionRegistry, getIndexResult, verifier);
+
+        VerificationException ex = expectThrows(VerificationException.class,
+            () -> plan("SELECT gender AS g, sum(salary) AS g FROM test GROUP BY g"));
+        assertEquals(
+            "Found 1 problem\nline 1:57: Reference [g] is ambiguous (to disambiguate use quotes or qualifiers); " +
+                "matches any of [line 1:8 [\"g\"], line 1:21 [\"g\"]]",
+            ex.getMessage());
+
+        ex = expectThrows(VerificationException.class,
+            () -> plan("SELECT gender AS g, max(salary) AS g, min(salary) AS g FROM test GROUP BY g"));
+        assertEquals(
+            "Found 1 problem\nline 1:75: Reference [g] is ambiguous (to disambiguate use quotes or qualifiers); " +
+                "matches any of [line 1:8 [\"g\"], line 1:21 [\"g\"], line 1:39 [\"g\"]]",
+            ex.getMessage());
+
+        ex = expectThrows(VerificationException.class,
+            () -> plan("SELECT gender AS g, last_name AS g, sum(salary) AS s FROM test GROUP BY g"));
+        assertEquals(
+            "Found 1 problem\nline 1:73: Reference [g] is ambiguous (to disambiguate use quotes or qualifiers); " +
+                "matches any of [line 1:8 [\"g\"], line 1:21 [\"g\"]]",
+            ex.getMessage());
+    }
 }