Fix ORDER BY expression to use output columns

Fixes prestodb#4698.

presto-main/src/main/java/com/facebook/presto/sql/analyzer/AggregationAnalyzer.java

@@ -106,21 +106,7 @@ public AggregationAnalyzer(List<Expression> groupByExpressions, Metadata metadat
         // and the '*' will be expanded to Field references. Therefore we translate all simple name expressions
         // in the group by clause to fields they reference so that the expansion from '*' can be matched against them
         for (Expression expression : Iterables.filter(expressions, columnReferences::contains)) {
-            QualifiedName name;
-            if (expression instanceof QualifiedNameReference) {
-                name = ((QualifiedNameReference) expression).getName();
-            }
-            else {
-                name = DereferenceExpression.getQualifiedName(checkType(expression, DereferenceExpression.class, "expression"));
-            }
-
-            List<Field> fields = scope.getRelationType().resolveFields(name);
-            checkState(fields.size() <= 1, "Found more than one field for name '%s': %s", name, fields);
-
-            if (fields.size() == 1) {
-                Field field = Iterables.getOnlyElement(fields);
-                fieldIndexes.add(scope.getRelationType().indexOf(field));
-            }
+            fieldIndexes.add(scope.resolveField(expression).getFieldIndex());
         }
         this.fieldIndexes = fieldIndexes.build();
     }

presto-main/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java

@@ -313,10 +313,10 @@ public Scope getRootScope()
     private static class GetScopeVisitor
             extends DefaultTraversalVisitor<Void, Scope>
     {
+
         private final IdentityHashMap<Node, Scope> scopes;
         private final Node node;
         private Scope result;
-
         public GetScopeVisitor(IdentityHashMap<Node, Scope> scopes, Node node)
         {
             this.scopes = requireNonNull(scopes, "scopes is null");
@@ -347,8 +347,8 @@ public Optional<Scope> getResult()
         {
             return Optional.ofNullable(result);
         }
-    }
 
+    }
     public void setScope(Node node, Scope scope)
     {
         scopes.put(node, scope);
@@ -389,6 +389,11 @@ public Set<Expression> getColumnReferences()
         return ImmutableSet.copyOf(columnReferences);
     }
 
+    public void addType(Expression expression, Type type)
+    {
+        types.put(expression, type);
+    }
+
     public void addTypes(IdentityHashMap<Expression, Type> types)
     {
         this.types.putAll(types);
@@ -477,6 +482,15 @@ public double getSampleRatio(SampledRelation relation)
         return sampleRatios.get(relation);
     }
 
+    public void copyExpressionAnalysis(Expression expression, Expression analyzedExpression)
+    {
+        types.put(expression, getType(analyzedExpression));
+        coercions.put(expression, getCoercion(analyzedExpression));
+        if (isTypeOnlyCoercion(analyzedExpression)) {
+            typeOnlyCoercions.add(expression);
+        }
+    }
+
     @Immutable
     public static final class Insert
     {

presto-main/src/main/java/com/facebook/presto/sql/analyzer/ExpressionAnalyzer.java

@@ -302,7 +302,7 @@ protected Type visitDereferenceExpression(DereferenceExpression node, StackableA
 
             // If this Dereference looks like column reference, try match it to column first.
             if (qualifiedName != null) {
-                Optional<ResolvedField> resolvedField = scope.tryResolveField(node, qualifiedName);
+                Optional<ResolvedField> resolvedField = scope.tryResolveField(qualifiedName);
                 if (resolvedField.isPresent()) {
                     return handleResolvedField(node, resolvedField.get());
                 }

presto-main/src/main/java/com/facebook/presto/sql/analyzer/Field.java

@@ -15,11 +15,15 @@
 
 import com.facebook.presto.spi.type.Type;
 import com.facebook.presto.sql.tree.QualifiedName;
+import com.google.common.collect.ImmutableList;
+
+import javax.annotation.concurrent.Immutable;
 
 import java.util.Optional;
 
 import static java.util.Objects.requireNonNull;
 
+@Immutable
 public class Field
 {
     public enum Kind
@@ -34,18 +38,21 @@ public enum Kind
 
     public static Field newUnqualified(String name, Type type)
     {
-        requireNonNull(name, "name is null");
-        requireNonNull(type, "type is null");
-
-        return new Field(Optional.empty(), Optional.of(name), type, Kind.VISIBLE);
+        return newUnqualified(Optional.of(name), type);
     }
 
     public static Field newUnqualified(Optional<String> name, Type type)
+    {
+        return newUnqualified(name, type, Kind.VISIBLE);
+    }
+
+    public static Field newUnqualified(Optional<String> name, Type type, Kind kind)
     {
         requireNonNull(name, "name is null");
         requireNonNull(type, "type is null");
+        requireNonNull(kind, "kind is null");
 
-        return new Field(Optional.empty(), name, type, Kind.VISIBLE);
+        return new Field(Optional.empty(), name, type, kind);
     }
 
     public static Field newQualified(QualifiedName relationAlias, Optional<String> name, Type type, Kind kind)
@@ -76,6 +83,19 @@ public Optional<QualifiedName> getRelationAlias()
         return relationAlias;
     }
 
+    public Optional<QualifiedName> getQualifiedName()
+    {
+        if (name.isPresent()) {
+            ImmutableList.Builder<String> parts = ImmutableList.builder();
+            relationAlias.ifPresent(alias -> parts.addAll(alias.getParts()));
+            parts.add(name.get());
+            return Optional.of(new QualifiedName(parts.build()));
+        }
+        else {
+            return Optional.empty();
+        }
+    }
+
     public Optional<String> getName()
     {
         return name;

presto-main/src/main/java/com/facebook/presto/sql/analyzer/Scope.java

@@ -18,19 +18,21 @@
 import com.facebook.presto.sql.tree.QualifiedName;
 import com.facebook.presto.sql.tree.QualifiedNameReference;
 import com.facebook.presto.sql.tree.WithQuery;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 
 import javax.annotation.concurrent.Immutable;
 
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
 import static com.facebook.presto.sql.analyzer.SemanticExceptions.createAmbiguousAttributeException;
 import static com.facebook.presto.sql.analyzer.SemanticExceptions.createMissingAttributeException;
+import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.Iterables.getOnlyElement;
 import static java.util.Objects.requireNonNull;
 
 @Immutable
@@ -61,16 +63,41 @@ public RelationType getRelationType()
         return relation;
     }
 
+    public ResolvedField resolveField(Expression expression)
+    {
+        return resolveField(expression, asQualifiedName(expression));
+    }
+
     public ResolvedField resolveField(Expression expression, QualifiedName name)
     {
-        return tryResolveField(expression, name).orElseThrow(() -> createMissingAttributeException(expression, name));
+        List<ResolvedField> resolvedFields = resolveField(name, true);
+        if (resolvedFields.size() == 0) {
+            throw createMissingAttributeException(expression, name);
+        }
+        else if (resolvedFields.size() == 1) {
+            return resolvedFields.get(0);
+        }
+        else {
+            resolvedFields = filterVisible(resolvedFields);
+            if (resolvedFields.size() == 1) {
+                return resolvedFields.get(0);
+            }
+            throw createAmbiguousAttributeException(expression, name);
+        }
+    }
+
+    private List<ResolvedField> filterVisible(List<ResolvedField> resolvedFields)
+    {
+        return resolvedFields.stream()
+                .filter(resolvedField -> resolvedField.getField().isVisible())
+                .collect(toImmutableList());
     }
 
     public Optional<ResolvedField> tryResolveField(Expression expression)
     {
         QualifiedName qualifiedName = asQualifiedName(expression);
         if (qualifiedName != null) {
-            return tryResolveField(expression, qualifiedName);
+            return tryResolveField(qualifiedName);
         }
         return Optional.empty();
     }
@@ -87,40 +114,52 @@ else if (expression instanceof DereferenceExpression) {
         return name;
     }
 
-    public Optional<ResolvedField> tryResolveField(Expression node, QualifiedName name)
+    public Optional<ResolvedField> tryResolveField(QualifiedName name)
     {
-        return resolveField(node, name, true);
+        List<ResolvedField> resolvedFields = resolveField(name, true);
+        if (resolvedFields.size() > 1) {
+            resolvedFields = filterVisible(resolvedFields);
+        }
+        if (resolvedFields.size() == 1) {
+            return Optional.of(resolvedFields.get(0));
+        }
+        return Optional.empty();
     }
 
-    private Optional<ResolvedField> resolveField(Expression node, QualifiedName name, boolean local)
+    private List<ResolvedField> resolveField(QualifiedName name, boolean local)
     {
         List<Field> matches = relation.resolveFields(name);
-        if (matches.size() > 1) {
-            throw createAmbiguousAttributeException(node, name);
-        }
-
         if (matches.isEmpty()) {
             if (isColumnReference(name, relation)) {
-                return Optional.empty();
+                return ImmutableList.of();
             }
             Scope boundary = this;
             while (!boundary.queryBoundary) {
                 if (boundary.parent.isPresent()) {
                     boundary = boundary.parent.get();
                 }
                 else {
-                    return Optional.empty();
+                    return ImmutableList.of();
                 }
             }
             if (boundary.parent.isPresent()) {
                 // jump over the query boundary
-                return boundary.parent.get().resolveField(node, name, false);
+                return boundary.parent.get().resolveField(name, false);
             }
-            return Optional.empty();
-        }
-        else {
-            return Optional.of(asResolvedField(getOnlyElement(matches), local));
+            return ImmutableList.of();
         }
+
+        List<Optional<QualifiedName>> names = new ArrayList<>();
+        return matches.stream()
+                .filter(field -> {
+                    if (!names.contains(field.getQualifiedName())) {
+                        names.add(field.getQualifiedName());
+                        return true;
+                    }
+                    return false;
+                })
+                .map(field -> asResolvedField(field, local))
+                .collect(toImmutableList());
     }
 
     private ResolvedField asResolvedField(Field field, boolean local)
@@ -172,6 +211,14 @@ public boolean isApproximate()
         return approximate;
     }
 
+    public Scope copyWithPrunedHiddenFields()
+    {
+        ImmutableList<Field> nonHiddenFields = relation.getAllFields().stream()
+                .filter(field -> field.getKind() != Field.Kind.HIDDEN)
+                .collect(toImmutableList());
+        return new Scope(parent, new RelationType(nonHiddenFields), namedQueries, approximate, queryBoundary);
+    }
+
     public static final class Builder
     {
         private RelationType relationType = new RelationType();