Add VariableReferenceExpression to TopNRowNumberNode

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java

@@ -418,7 +418,7 @@ public PlanWithProperties visitTopNRowNumber(TopNRowNumberNode node, PreferredPr
                                 idAllocator.getNextId(),
                                 child.getNode(),
                                 node.getSpecification(),
-                                node.getRowNumberSymbol(),
+                                node.getRowNumberVariable(),
                                 node.getMaxRowCountPerPartition(),
                                 true,
                                 node.getHashSymbol()),

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/HashGenerationOptimizer.java

@@ -290,7 +290,7 @@ public PlanWithProperties visitTopNRowNumber(TopNRowNumberNode node, HashComputa
                             node.getId(),
                             child.getNode(),
                             node.getSpecification(),
-                            node.getRowNumberSymbol(),
+                            node.getRowNumberVariable(),
                             node.getMaxRowCountPerPartition(),
                             node.isPartial(),
                             Optional.of(hashSymbol)),

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java

@@ -623,7 +623,7 @@ public PlanNode visitTopNRowNumber(TopNRowNumberNode node, RewriteContext<Set<Sy
             return new TopNRowNumberNode(node.getId(),
                     source,
                     node.getSpecification(),
-                    node.getRowNumberSymbol(),
+                    node.getRowNumberVariable(),
                     node.getMaxRowCountPerPartition(),
                     node.isPartial(),
                     node.getHashSymbol());

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java

@@ -422,7 +422,7 @@ public PlanNode visitTopNRowNumber(TopNRowNumberNode node, RewriteContext<Void>
                     node.getId(),
                     context.rewrite(node.getSource()),
                     canonicalizeAndDistinct(node.getSpecification()),
-                    canonicalize(node.getRowNumberSymbol()),
+                    canonicalize(node.getRowNumberVariable()),
                     node.getMaxRowCountPerPartition(),
                     node.isPartial(),
                     canonicalize(node.getHashSymbol()));

presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/WindowFilterPushDown.java

@@ -265,7 +265,7 @@ private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limi
             return new TopNRowNumberNode(idAllocator.getNextId(),
                     windowNode.getSource(),
                     windowNode.getSpecification(),
-                    getOnlyElement(windowNode.getCreatedSymbols()),
+                    getOnlyElement(windowNode.getCreatedVariable()),
                     limit,
                     false,
                     Optional.empty());

presto-main/src/main/java/com/facebook/presto/sql/planner/plan/TopNRowNumberNode.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.sql.planner.plan;
 
 import com.facebook.presto.spi.plan.PlanNodeId;
+import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.sql.planner.OrderingScheme;
 import com.facebook.presto.sql.planner.Symbol;
 import com.facebook.presto.sql.planner.plan.WindowNode.Specification;
@@ -28,7 +29,6 @@
 import java.util.Optional;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static com.google.common.collect.Iterables.concat;
 import static java.util.Objects.requireNonNull;
 
 @Immutable
@@ -37,7 +37,7 @@ public final class TopNRowNumberNode
 {
     private final PlanNode source;
     private final Specification specification;
-    private final Symbol rowNumberSymbol;
+    private final VariableReferenceExpression rowNumberVariable;
     private final int maxRowCountPerPartition;
     private final boolean partial;
     private final Optional<Symbol> hashSymbol;
@@ -47,7 +47,7 @@ public TopNRowNumberNode(
             @JsonProperty("id") PlanNodeId id,
             @JsonProperty("source") PlanNode source,
             @JsonProperty("specification") Specification specification,
-            @JsonProperty("rowNumberSymbol") Symbol rowNumberSymbol,
+            @JsonProperty("rowNumberVariable") VariableReferenceExpression rowNumberVariable,
             @JsonProperty("maxRowCountPerPartition") int maxRowCountPerPartition,
             @JsonProperty("partial") boolean partial,
             @JsonProperty("hashSymbol") Optional<Symbol> hashSymbol)
@@ -57,13 +57,13 @@ public TopNRowNumberNode(
         requireNonNull(source, "source is null");
         requireNonNull(specification, "specification is null");
         checkArgument(specification.getOrderingScheme().isPresent(), "specification orderingScheme is absent");
-        requireNonNull(rowNumberSymbol, "rowNumberSymbol is null");
+        requireNonNull(rowNumberVariable, "rowNumberVariable is null");
         checkArgument(maxRowCountPerPartition > 0, "maxRowCountPerPartition must be > 0");
         requireNonNull(hashSymbol, "hashSymbol is null");
 
         this.source = source;
         this.specification = specification;
-        this.rowNumberSymbol = rowNumberSymbol;
+        this.rowNumberVariable = rowNumberVariable;
         this.maxRowCountPerPartition = maxRowCountPerPartition;
         this.partial = partial;
         this.hashSymbol = hashSymbol;
@@ -78,10 +78,13 @@ public List<PlanNode> getSources()
     @Override
     public List<Symbol> getOutputSymbols()
     {
+        ImmutableList.Builder<Symbol> builder = ImmutableList.<Symbol>builder().addAll(source.getOutputSymbols());
+
         if (!partial) {
-            return ImmutableList.copyOf(concat(source.getOutputSymbols(), ImmutableList.of(rowNumberSymbol)));
+            builder.add(new Symbol(rowNumberVariable.getName()));
         }
-        return ImmutableList.copyOf(source.getOutputSymbols());
+
+        return builder.build();
     }
 
     @JsonProperty
@@ -106,10 +109,15 @@ public OrderingScheme getOrderingScheme()
         return specification.getOrderingScheme().get();
     }
 
-    @JsonProperty
     public Symbol getRowNumberSymbol()
     {
-        return rowNumberSymbol;
+        return new Symbol(rowNumberVariable.getName());
+    }
+
+    @JsonProperty
+    public VariableReferenceExpression getRowNumberVariable()
+    {
+        return rowNumberVariable;
     }
 
     @JsonProperty
@@ -139,6 +147,6 @@ public <R, C> R accept(InternalPlanVisitor<R, C> visitor, C context)
     @Override
     public PlanNode replaceChildren(List<PlanNode> newChildren)
     {
-        return new TopNRowNumberNode(getId(), Iterables.getOnlyElement(newChildren), specification, rowNumberSymbol, maxRowCountPerPartition, partial, hashSymbol);
+        return new TopNRowNumberNode(getId(), Iterables.getOnlyElement(newChildren), specification, rowNumberVariable, maxRowCountPerPartition, partial, hashSymbol);
     }
 }

presto-main/src/main/java/com/facebook/presto/sql/planner/plan/WindowNode.java

@@ -108,6 +108,11 @@ public Set<Symbol> getCreatedSymbols()
         return windowFunctions.keySet().stream().map(variable -> new Symbol(variable.getName())).collect(toImmutableSet());
     }
 
+    public Set<VariableReferenceExpression> getCreatedVariable()
+    {
+        return windowFunctions.keySet();
+    }
+
     @JsonProperty
     public PlanNode getSource()
     {

presto-main/src/main/java/com/facebook/presto/sql/planner/planPrinter/PlanPrinter.java

@@ -625,7 +625,7 @@ public Void visitTopNRowNumber(TopNRowNumberNode node, Void context)
                     "TopNRowNumber",
                     format("[%s limit %s]%s", Joiner.on(", ").join(args), node.getMaxRowCountPerPartition(), formatHash(node.getHashSymbol())));
 
-            nodeOutput.appendDetailsLine("%s := %s", node.getRowNumberSymbol(), "row_number()");
+            nodeOutput.appendDetailsLine("%s := %s", node.getRowNumberVariable(), "row_number()");
 
             return processChildren(node, context);
         }