Add VariableReferenceExpression to RowNumberNode

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

@@ -264,7 +264,7 @@ public PlanWithProperties visitRowNumber(RowNumberNode node, HashComputationSet
                             node.getId(),
                             child.getNode(),
                             node.getPartitionBy(),
-                            node.getRowNumberSymbol(),
+                            node.getRowNumberVariable(),
                             node.getMaxRowCountPerPartition(),
                             Optional.of(hashSymbol)),
                     child.getHashSymbols());

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

@@ -604,7 +604,7 @@ public PlanNode visitRowNumber(RowNumberNode node, RewriteContext<Set<Symbol>> c
             }
             PlanNode source = context.rewrite(node.getSource(), expectedInputs.build());
 
-            return new RowNumberNode(node.getId(), source, node.getPartitionBy(), node.getRowNumberSymbol(), node.getMaxRowCountPerPartition(), node.getHashSymbol());
+            return new RowNumberNode(node.getId(), source, node.getPartitionBy(), node.getRowNumberVariable(), node.getMaxRowCountPerPartition(), node.getHashSymbol());
         }
 
         @Override

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

@@ -412,7 +412,7 @@ public PlanNode visitTableFinish(TableFinishNode node, RewriteContext<Void> cont
         @Override
         public PlanNode visitRowNumber(RowNumberNode node, RewriteContext<Void> context)
         {
-            return new RowNumberNode(node.getId(), context.rewrite(node.getSource()), canonicalizeAndDistinct(node.getPartitionBy()), canonicalize(node.getRowNumberSymbol()), node.getMaxRowCountPerPartition(), canonicalize(node.getHashSymbol()));
+            return new RowNumberNode(node.getId(), context.rewrite(node.getSource()), canonicalizeAndDistinct(node.getPartitionBy()), canonicalize(node.getRowNumberVariable()), node.getMaxRowCountPerPartition(), canonicalize(node.getHashSymbol()));
         }
 
         @Override

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

@@ -118,7 +118,7 @@ public PlanNode visitWindow(WindowNode node, RewriteContext<Void> context)
                 return new RowNumberNode(idAllocator.getNextId(),
                         rewrittenSource,
                         node.getPartitionBy(),
-                        getOnlyElement(node.getCreatedSymbols()),
+                        getOnlyElement(node.getWindowFunctions().keySet()),
                         Optional.empty(),
                         Optional.empty());
             }
@@ -257,7 +257,7 @@ private static RowNumberNode mergeLimit(RowNumberNode node, int newRowCountPerPa
             if (node.getMaxRowCountPerPartition().isPresent()) {
                 newRowCountPerPartition = Math.min(node.getMaxRowCountPerPartition().get(), newRowCountPerPartition);
             }
-            return new RowNumberNode(node.getId(), node.getSource(), node.getPartitionBy(), node.getRowNumberSymbol(), Optional.of(newRowCountPerPartition), node.getHashSymbol());
+            return new RowNumberNode(node.getId(), node.getSource(), node.getPartitionBy(), node.getRowNumberVariable(), Optional.of(newRowCountPerPartition), node.getHashSymbol());
         }
 
         private TopNRowNumberNode convertToTopNRowNumber(WindowNode windowNode, int limit)

presto-main/src/main/java/com/facebook/presto/sql/planner/plan/RowNumberNode.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.Symbol;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
@@ -25,7 +26,6 @@
 import java.util.List;
 import java.util.Optional;
 
-import static com.google.common.collect.Iterables.concat;
 import static java.util.Objects.requireNonNull;
 
 @Immutable
@@ -35,29 +35,29 @@ public final class RowNumberNode
     private final PlanNode source;
     private final List<Symbol> partitionBy;
     private final Optional<Integer> maxRowCountPerPartition;
-    private final Symbol rowNumberSymbol;
+    private final VariableReferenceExpression rowNumberVariable;
     private final Optional<Symbol> hashSymbol;
 
     @JsonCreator
     public RowNumberNode(
             @JsonProperty("id") PlanNodeId id,
             @JsonProperty("source") PlanNode source,
             @JsonProperty("partitionBy") List<Symbol> partitionBy,
-            @JsonProperty("rowNumberSymbol") Symbol rowNumberSymbol,
+            @JsonProperty("rowNumberVariable") VariableReferenceExpression rowNumberVariable,
             @JsonProperty("maxRowCountPerPartition") Optional<Integer> maxRowCountPerPartition,
             @JsonProperty("hashSymbol") Optional<Symbol> hashSymbol)
     {
         super(id);
 
         requireNonNull(source, "source is null");
         requireNonNull(partitionBy, "partitionBy is null");
-        requireNonNull(rowNumberSymbol, "rowNumberSymbol is null");
+        requireNonNull(rowNumberVariable, "rowNumberVariable is null");
         requireNonNull(maxRowCountPerPartition, "maxRowCountPerPartition is null");
         requireNonNull(hashSymbol, "hashSymbol is null");
 
         this.source = source;
         this.partitionBy = ImmutableList.copyOf(partitionBy);
-        this.rowNumberSymbol = rowNumberSymbol;
+        this.rowNumberVariable = rowNumberVariable;
         this.maxRowCountPerPartition = maxRowCountPerPartition;
         this.hashSymbol = hashSymbol;
     }
@@ -71,7 +71,19 @@ public List<PlanNode> getSources()
     @Override
     public List<Symbol> getOutputSymbols()
     {
-        return ImmutableList.copyOf(concat(source.getOutputSymbols(), ImmutableList.of(rowNumberSymbol)));
+        return ImmutableList.<Symbol>builder()
+                .addAll(source.getOutputSymbols())
+                .add(new Symbol(rowNumberVariable.getName()))
+                .build();
+    }
+
+    @Override
+    public List<VariableReferenceExpression> getOutputVariables()
+    {
+        return ImmutableList.<VariableReferenceExpression>builder()
+                .addAll(source.getOutputVariables())
+                .add(rowNumberVariable)
+                .build();
     }
 
     @JsonProperty
@@ -86,10 +98,15 @@ public List<Symbol> getPartitionBy()
         return partitionBy;
     }
 
-    @JsonProperty
     public Symbol getRowNumberSymbol()
     {
-        return rowNumberSymbol;
+        return new Symbol(rowNumberVariable.getName());
+    }
+
+    @JsonProperty
+    public VariableReferenceExpression getRowNumberVariable()
+    {
+        return rowNumberVariable;
     }
 
     @JsonProperty
@@ -113,6 +130,6 @@ public <R, C> R accept(InternalPlanVisitor<R, C> visitor, C context)
     @Override
     public PlanNode replaceChildren(List<PlanNode> newChildren)
     {
-        return new RowNumberNode(getId(), Iterables.getOnlyElement(newChildren), partitionBy, rowNumberSymbol, maxRowCountPerPartition, hashSymbol);
+        return new RowNumberNode(getId(), Iterables.getOnlyElement(newChildren), partitionBy, rowNumberVariable, maxRowCountPerPartition, hashSymbol);
     }
 }

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

@@ -646,7 +646,7 @@ public Void visitRowNumber(RowNumberNode node, Void context)
             NodeRepresentation nodeOutput = addNode(node,
                     "RowNumber",
                     format("[%s]%s", Joiner.on(", ").join(args), formatHash(node.getHashSymbol())));
-            nodeOutput.appendDetailsLine("%s := %s", node.getRowNumberSymbol(), "row_number()");
+            nodeOutput.appendDetailsLine("%s := %s", node.getRowNumberVariable(), "row_number()");
 
             return processChildren(node, context);
         }

presto-main/src/test/java/com/facebook/presto/cost/TestRowNumberStatsRule.java

@@ -42,6 +42,7 @@ public void testSingleGroupingKey()
                         ImmutableList.of(pb.symbol("x", BIGINT)),
                         Optional.empty(),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .setOutputRowCount(10)
@@ -64,6 +65,7 @@ public void testSingleGroupingKey()
                         ImmutableList.of(pb.symbol("x", BIGINT)),
                         Optional.of(1),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .setOutputRowCount(10)
@@ -84,6 +86,7 @@ public void testSingleGroupingKey()
                         ImmutableList.of(pb.symbol("y", BIGINT)),
                         Optional.empty(),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .setOutputRowCount(60)
@@ -104,6 +107,7 @@ public void testSingleGroupingKey()
                         ImmutableList.of(pb.symbol("x", BIGINT)),
                         Optional.of(1),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .addSymbolStatistics(new Symbol("x"), xStats)
@@ -121,6 +125,7 @@ public void testMultipleGroupingKeys()
                         ImmutableList.of(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT)),
                         Optional.empty(),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .setOutputRowCount(60)
@@ -141,6 +146,7 @@ public void testMultipleGroupingKeys()
                         ImmutableList.of(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT)),
                         Optional.empty(),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .setOutputRowCount(20)
@@ -161,6 +167,7 @@ public void testMultipleGroupingKeys()
                         ImmutableList.of(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT)),
                         Optional.empty(),
                         pb.symbol("z", BIGINT),
+                        pb.variable(pb.symbol("z", BIGINT)),
                         pb.values(pb.symbol("x", BIGINT), pb.symbol("y", BIGINT))))
                 .withSourceStats(0, PlanNodeStatsEstimate.builder()
                         .setOutputRowCount(20)

presto-main/src/test/java/com/facebook/presto/sql/planner/iterative/rule/test/PlanBuilder.java

@@ -746,13 +746,13 @@ public WindowNode window(WindowNode.Specification specification, Map<VariableRef
                 0);
     }
 
-    public RowNumberNode rowNumber(List<Symbol> partitionBy, Optional<Integer> maxRowCountPerPartition, Symbol rowNumberSymbol, PlanNode source)
+    public RowNumberNode rowNumber(List<Symbol> partitionBy, Optional<Integer> maxRowCountPerPartition, Symbol rowNumberSymbol, VariableReferenceExpression rownNumberVariable, PlanNode source)
     {
         return new RowNumberNode(
                 idAllocator.getNextId(),
                 source,
                 partitionBy,
-                rowNumberSymbol,
+                rownNumberVariable,
                 maxRowCountPerPartition,
                 Optional.empty());
     }