Introduce scope to semantic analyzer and update correlation in Apply node

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

@@ -82,16 +82,16 @@ class AggregationAnalyzer
     private final Metadata metadata;
     private final Set<Expression> columnReferences;
 
-    private final RelationType tupleDescriptor;
+    private final Scope scope;
 
-    public AggregationAnalyzer(List<Expression> groupByExpressions, Metadata metadata, RelationType tupleDescriptor, Set<Expression> columnReferences)
+    public AggregationAnalyzer(List<Expression> groupByExpressions, Metadata metadata, Scope scope, Set<Expression> columnReferences)
     {
         requireNonNull(groupByExpressions, "groupByExpressions is null");
         requireNonNull(metadata, "metadata is null");
-        requireNonNull(tupleDescriptor, "tupleDescriptor is null");
+        requireNonNull(scope, "scope is null");
         requireNonNull(columnReferences, "columnReferences is null");
 
-        this.tupleDescriptor = tupleDescriptor;
+        this.scope = scope;
         this.metadata = metadata;
         this.columnReferences = ImmutableSet.copyOf(columnReferences);
         this.expressions = ImmutableList.copyOf(groupByExpressions);
@@ -115,12 +115,12 @@ public AggregationAnalyzer(List<Expression> groupByExpressions, Metadata metadat
                 name = DereferenceExpression.getQualifiedName(checkType(expression, DereferenceExpression.class, "expression"));
             }
 
-            List<Field> fields = tupleDescriptor.resolveFields(name);
+            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(tupleDescriptor.indexOf(field));
+                fieldIndexes.add(scope.getRelationType().indexOf(field));
             }
         }
         this.fieldIndexes = fieldIndexes.build();
@@ -365,12 +365,12 @@ protected Boolean visitDereferenceExpression(DereferenceExpression node, Void co
 
         private Boolean isField(QualifiedName qualifiedName)
         {
-            List<Field> fields = tupleDescriptor.resolveFields(qualifiedName);
+            List<Field> fields = scope.getRelationType().resolveFields(qualifiedName);
             checkState(!fields.isEmpty(), "No fields for name '%s'", qualifiedName);
             checkState(fields.size() <= 1, "Found more than one field for name '%s': %s", qualifiedName, fields);
 
             Field field = Iterables.getOnlyElement(fields);
-            return fieldIndexes.contains(tupleDescriptor.indexOf(field));
+            return fieldIndexes.contains(scope.getRelationType().indexOf(field));
         }
 
         @Override

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

@@ -18,6 +18,7 @@
 import com.facebook.presto.metadata.TableHandle;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.sql.tree.DefaultTraversalVisitor;
 import com.facebook.presto.sql.tree.ExistsPredicate;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.FunctionCall;
@@ -35,8 +36,10 @@
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ListMultimap;
 
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.Immutable;
 
 import java.util.IdentityHashMap;
@@ -51,14 +54,13 @@
 
 public class Analysis
 {
-    private final Statement statement;
+    private final Statement root;
     private String updateType;
 
     private final IdentityHashMap<Table, Query> namedQueries = new IdentityHashMap<>();
 
-    private RelationType outputDescriptor;
-    private final IdentityHashMap<Node, RelationType> outputDescriptors = new IdentityHashMap<>();
-    private final IdentityHashMap<Expression, Integer> resolvedNames = new IdentityHashMap<>();
+    private final IdentityHashMap<Node, Scope> scopes = new IdentityHashMap<>();
+    private final Set<Expression> columnReferences = newIdentityHashSet();
 
     private final IdentityHashMap<QuerySpecification, List<FunctionCall>> aggregates = new IdentityHashMap<>();
     private final IdentityHashMap<QuerySpecification, List<List<Expression>>> groupByExpressions = new IdentityHashMap<>();
@@ -93,14 +95,14 @@ public class Analysis
 
     private Optional<Insert> insert = Optional.empty();
 
-    public Analysis(Statement statement)
+    public Analysis(Statement root)
     {
-        this.statement = statement;
+        this.root = root;
     }
 
     public Statement getStatement()
     {
-        return statement;
+        return root;
     }
 
     public String getUpdateType()
@@ -133,16 +135,6 @@ public void setCreateTableAsSelectNoOp(boolean createTableAsSelectNoOp)
         this.createTableAsSelectNoOp = createTableAsSelectNoOp;
     }
 
-    public void addResolvedNames(Map<Expression, Integer> mappings)
-    {
-        resolvedNames.putAll(mappings);
-    }
-
-    public Optional<Integer> getFieldIndex(Expression expression)
-    {
-        return Optional.ofNullable(resolvedNames.get(expression));
-    }
-
     public void setAggregates(QuerySpecification node, List<FunctionCall> aggregates)
     {
         this.aggregates.put(node, aggregates);
@@ -299,25 +291,88 @@ public List<FunctionCall> getWindowFunctions(QuerySpecification query)
         return windowFunctions.get(query);
     }
 
-    public void setOutputDescriptor(RelationType descriptor)
+    public void addColumnReferences(Set<Expression> columnReferences)
     {
-        outputDescriptor = descriptor;
+        this.columnReferences.addAll(columnReferences);
     }
 
-    public RelationType getOutputDescriptor()
+    public Scope getScope(Node node)
+    {
+        return tryGetScope(node).orElseThrow(() -> new IllegalArgumentException(String.format("Analysis does not contain information for node: %s", node)));
+    }
+
+    public Optional<Scope> tryGetScope(Node node)
+    {
+        if (scopes.containsKey(node)) {
+            return Optional.of(scopes.get(node));
+        }
+
+        if (root == null) {
+            return Optional.empty();
+        }
+
+        GetScopeVisitor visitor = new GetScopeVisitor(scopes, node);
+        visitor.process(root, null);
+        return visitor.getResult();
+    }
+
+    public Scope getRootScope()
     {
-        return outputDescriptor;
+        return getScope(root);
     }
 
-    public void setOutputDescriptor(Node node, RelationType descriptor)
+    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");
+            this.node = requireNonNull(node, "node is null");
+        }
+
+        @Override
+        public Void process(Node current, @Nullable Scope candidate)
+        {
+            if (result != null) {
+                return null;
+            }
+
+            if (scopes.containsKey(current)) {
+                candidate = scopes.get(current);
+            }
+            if (node == current) {
+                result = candidate;
+            }
+            else {
+                super.process(current, candidate);
+            }
+
+            return null;
+        }
+
+        public Optional<Scope> getResult()
+        {
+            return Optional.ofNullable(result);
+        }
+    }
+
+    public void setScope(Node node, Scope scope)
+    {
+        scopes.put(node, scope);
+    }
+
+    public RelationType getOutputDescriptor()
     {
-        outputDescriptors.put(node, descriptor);
+        return getOutputDescriptor(root);
     }
 
     public RelationType getOutputDescriptor(Node node)
     {
-        Preconditions.checkState(outputDescriptors.containsKey(node), "Output descriptor missing for %s. Broken analysis?", node);
-        return outputDescriptors.get(node);
+        return getScope(node).getRelationType();
     }
 
     public TableHandle getTableHandle(Table table)
@@ -342,7 +397,7 @@ public void addFunctionSignatures(IdentityHashMap<FunctionCall, Signature> infos
 
     public Set<Expression> getColumnReferences()
     {
-        return resolvedNames.keySet();
+        return ImmutableSet.copyOf(columnReferences);
     }
 
     public void addTypes(IdentityHashMap<Expression, Type> types)

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

@@ -59,8 +59,7 @@ public Analysis analyze(Statement statement)
         Statement rewrittenStatement = StatementRewrite.rewrite(session, metadata, sqlParser, queryExplainer, statement);
         Analysis analysis = new Analysis(rewrittenStatement);
         StatementAnalyzer analyzer = new StatementAnalyzer(analysis, metadata, sqlParser, accessControl, session, experimentalSyntaxEnabled);
-        RelationType outputDescriptor = analyzer.process(rewrittenStatement, new AnalysisContext());
-        analysis.setOutputDescriptor(outputDescriptor);
+        analyzer.process(rewrittenStatement, Scope.builder().build());
         return analysis;
     }