Repository: prestodb/presto
Issue: Incorrect results when GROUPING SETS refer to columns derived from same column

This query:

``` sql
WITH
   t (v) AS (VALUES 1, 2),
   u (v1, v2) AS (SELECT v v1, v v2 FROM t)
SELECT v1, v2
FROM u
GROUP BY GROUPING SETS ((v1), (v2))
```

incorrectly produces:

```
 v1 | v2
----+----
  1 |  1
  2 |  2
  1 |  1
  2 |  2
(4 rows)
```

The expected result is:

```
  v1  |  v2
------+------
    1 | NULL
    2 | NULL
 NULL |    1
 NULL |    2
(4 rows)
```


Changed files:
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ChildReplacer.java
presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java
presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
presto-main/src/test/java/com/facebook/presto/operator/TestGroupIdOperator.java
presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java
presto-main/src/main/java/com/facebook/presto/sql/planner/SymbolExtractor.java
presto-main/src/main/java/com/facebook/presto/operator/GroupIdOperator.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PredicatePushDown.java
presto-main/src/main/java/com/facebook/presto/sql/planner/plan/GroupIdNode.java
presto-main/src/main/java/com/facebook/presto/sql/planner/PlanPrinter.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/GroupIdOperator.java b/presto-main/src/main/java/com/facebook/presto/operator/GroupIdOperator.java
index 18653de108..30588fe289 100644
--- a/presto-main/src/main/java/com/facebook/presto/operator/GroupIdOperator.java
+++ b/presto-main/src/main/java/com/facebook/presto/operator/GroupIdOperator.java
@@ -22,14 +22,11 @@ import com.facebook.presto.spi.type.Type;
 import com.facebook.presto.sql.planner.plan.PlanNodeId;
 import com.google.common.collect.ImmutableList;
 
-import java.util.BitSet;
-import java.util.Collection;
+import java.util.Arrays;
 import java.util.List;
-import java.util.Set;
+import java.util.Map;
 
 import static com.facebook.presto.spi.type.BigintType.BIGINT;
-import static com.facebook.presto.util.ImmutableCollectors.toImmutableSet;
-import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
 import static java.util.Objects.requireNonNull;
 
@@ -42,9 +39,7 @@ public class GroupIdOperator
         private final int operatorId;
         private final PlanNodeId planNodeId;
         private final List<Type> outputTypes;
-        private final List<List<Integer>> groupingSetChannels;
-        private final List<Integer> groupingChannels;
-        private final List<Integer> copyChannels;
+        private final List<Map<Integer, Integer>> groupingSetMappings;
 
         private boolean closed;
 
@@ -52,16 +47,12 @@ public class GroupIdOperator
                 int operatorId,
                 PlanNodeId planNodeId,
                 List<? extends Type> outputTypes,
-                List<List<Integer>> groupingSetChannels,
-                List<Integer> groupingChannels,
-                List<Integer> copyChannels)
+                List<Map<Integer, Integer>> groupingSetMappings)
         {
             this.operatorId = operatorId;
             this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
             this.outputTypes = ImmutableList.copyOf(requireNonNull(outputTypes));
-            this.groupingSetChannels = ImmutableList.copyOf(requireNonNull(groupingSetChannels));
-            this.groupingChannels = ImmutableList.copyOf(requireNonNull(groupingChannels));
-            this.copyChannels = ImmutableList.copyOf(requireNonNull(copyChannels));
+            this.groupingSetMappings = ImmutableList.copyOf(requireNonNull(groupingSetMappings));
         }
 
         @Override
@@ -76,48 +67,35 @@ public class GroupIdOperator
             checkState(!closed, "Factory is already closed");
             OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, GroupIdOperator.class.getSimpleName());
 
-            Set<Integer> allGroupingColumns = groupingSetChannels.stream()
-                    .flatMap(Collection::stream)
-                    .collect(toImmutableSet());
-
-            // create an array of bitset for fast lookup of which columns are part of a given grouping set
-            // will have a 'true' for every channel that should be set to null for each grouping set
-            BitSet[] groupingSetNullChannels = new BitSet[groupingSetChannels.size()];
-            for (int i = 0; i < groupingSetChannels.size(); i++) {
-                groupingSetNullChannels[i] = new BitSet(groupingChannels.size() + copyChannels.size());
-                // first set all grouping columns to true
-                for (int groupingColumn : allGroupingColumns) {
-                    groupingSetNullChannels[i].set(groupingColumn, true);
-                }
-                // then set all the columns in this grouping set to false
-                for (int nonNullGroupingColumn : groupingSetChannels.get(i)) {
-                    groupingSetNullChannels[i].set(nonNullGroupingColumn, false);
+            // create an int array for fast lookup of input columns for every grouping set
+            int[][] groupingSetInputs = new int[groupingSetMappings.size()][outputTypes.size() - 1];
+            for (int i = 0; i < groupingSetMappings.size(); i++) {
+                // -1 means the output column is null
+                Arrays.fill(groupingSetInputs[i], -1);
+
+                // anything else is an input column to copy
+                for (int outputChannel : groupingSetMappings.get(i).keySet()) {
+                    groupingSetInputs[i][outputChannel] = groupingSetMappings.get(i).get(outputChannel);
                 }
             }
 
-            // create null blocks for every grouping channel
-            Block[] nullBlocks = new Block[groupingChannels.size()];
-            for (int i = 0; i < groupingChannels.size(); i++) {
+            // it's easier to create null blocks for every output column even though we only null out some grouping column outputs
+            Block[] nullBlocks = new Block[outputTypes.size()];
+            for (int i = 0; i < outputTypes.size(); i++) {
                 nullBlocks[i] = outputTypes.get(i).createBlockBuilder(new BlockBuilderStatus(), 1)
                         .appendNull()
                         .build();
             }
 
             // create groupid blocks for every group
-            Block[] groupIdBlocks = new Block[groupingSetNullChannels.length];
-            for (int i = 0; i < groupingSetNullChannels.length; i++) {
+            Block[] groupIdBlocks = new Block[groupingSetMappings.size()];
+            for (int i = 0; i < groupingSetMappings.size(); i++) {
                 BlockBuilder builder = BIGINT.createBlockBuilder(new BlockBuilderStatus(), 1);
                 BIGINT.writeLong(builder, i);
                 groupIdBlocks[i] = builder.build();
             }
 
-            // create array of input channels for every grouping channel
-            int[] groupInputs = groupingChannels.stream().mapToInt(Integer::intValue).toArray();
-
-            // create array of input channels for every copy channel
-            int[] copyInputs = copyChannels.stream().mapToInt(Integer::intValue).toArray();
-
-            return new GroupIdOperator(operatorContext, outputTypes, groupingSetNullChannels, nullBlocks, groupIdBlocks, groupInputs, copyInputs);
+            return new GroupIdOperator(operatorContext, outputTypes, groupingSetInputs, nullBlocks, groupIdBlocks);
         }
 
         @Override
@@ -129,17 +107,15 @@ public class GroupIdOperator
         @Override
         public OperatorFactory duplicate()
         {
-            return new GroupIdOperatorFactory(operatorId, planNodeId, outputTypes, groupingSetChannels, groupingChannels, copyChannels);
+            return new GroupIdOperatorFactory(operatorId, planNodeId, outputTypes, groupingSetMappings);
         }
     }
 
     private final OperatorContext operatorContext;
     private final List<Type> types;
-    private final BitSet[] groupingSetNullChannels;
+    private final int[][] groupingSetInputs;
     private final Block[] nullBlocks;
     private final Block[] groupIdBlocks;
-    private final int[] groupInputs;
-    private final int[] copyInputs;
 
     private Page currentPage = null;
     private int currentGroupingSet = 0;
@@ -148,20 +124,15 @@ public class GroupIdOperator
     public GroupIdOperator(
             OperatorContext operatorContext,
             List<Type> types,
-            BitSet[] groupingSetNullChannels,
+            int[][] groupingSetInputs,
             Block[] nullBlocks,
-            Block[] groupIdBlocks,
-            int[] groupInputs,
-            int[] copyInputs)
+            Block[] groupIdBlocks)
     {
         this.operatorContext = requireNonNull(operatorContext, "operatorContext is null");
-        this.types = requireNonNull(types, "types is null");
-        this.groupingSetNullChannels =  requireNonNull(groupingSetNullChannels, "groupingSetNullChannels is null");
-        this.nullBlocks = requireNonNull(nullBlocks);
-        this.groupIdBlocks = requireNonNull(groupIdBlocks);
-        checkArgument(groupIdBlocks.length == groupingSetNullChannels.length, "groupIdBlocks and groupingSetNullChannels must have the same length");
-        this.groupInputs = requireNonNull(groupInputs);
-        this.copyInputs = requireNonNull(copyInputs);
+        this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
+        this.groupingSetInputs =  requireNonNull(groupingSetInputs, "groupingSetInputs is null");
+        this.nullBlocks = requireNonNull(nullBlocks, "nullBlocks is null");
+        this.groupIdBlocks = requireNonNull(groupIdBlocks,  "groupIdBlocks is null");
     }
 
     @Override
@@ -218,21 +189,17 @@ public class GroupIdOperator
         // generate 'n' pages for every input page, where n is the number of grouping sets
         Block[] outputBlocks = new Block[types.size()];
 
-        for (int i = 0; i < groupInputs.length; i++) {
-            if (groupingSetNullChannels[currentGroupingSet].get(groupInputs[i])) {
+        for (int i = 0; i < groupingSetInputs[currentGroupingSet].length; i++) {
+            if (groupingSetInputs[currentGroupingSet][i] == -1) {
                 outputBlocks[i] = new RunLengthEncodedBlock(nullBlocks[i], currentPage.getPositionCount());
             }
             else {
-                outputBlocks[i] = currentPage.getBlock(groupInputs[i]);
+                outputBlocks[i] = currentPage.getBlock(groupingSetInputs[currentGroupingSet][i]);
             }
         }
 
-        for (int i = 0; i < copyInputs.length; i++) {
-            outputBlocks[groupInputs.length + i] = currentPage.getBlock(copyInputs[i]);
-        }
-
         outputBlocks[outputBlocks.length - 1] = new RunLengthEncodedBlock(groupIdBlocks[currentGroupingSet], currentPage.getPositionCount());
-        currentGroupingSet = (currentGroupingSet + 1) % groupingSetNullChannels.length;
+        currentGroupingSet = (currentGroupingSet + 1) % groupingSetInputs.length;
         Page outputPage = new Page(currentPage.getPositionCount(), outputBlocks);
 
         if (currentGroupingSet == 0) {
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java
index 24acf28fe5..b3a4d5344a 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.java
@@ -832,42 +832,50 @@ public class LocalExecutionPlanner
         public PhysicalOperation visitGroupId(GroupIdNode node, LocalExecutionPlanContext context)
         {
             PhysicalOperation source = node.getSource().accept(this, context);
-            ImmutableMap.Builder<Symbol, Integer> newLayout = ImmutableMap.builder();
+            Map<Symbol, Integer> newLayout = new HashMap<>();
             ImmutableList.Builder<Type> outputTypes = ImmutableList.builder();
 
             int outputChannel = 0;
 
-            ImmutableList.Builder<Integer> groupingChannels = ImmutableList.builder();
-            for (Symbol inputSymbol : node.getDistinctGroupingColumns()) {
-                int inputChannel = source.getLayout().get(inputSymbol);
-                newLayout.put(inputSymbol, outputChannel++);
-                outputTypes.add(source.getTypes().get(inputChannel));
-                groupingChannels.add(inputChannel);
+            for (Symbol output : node.getGroupingSets().stream().flatMap(Collection::stream).collect(Collectors.toSet())) {
+                newLayout.put(output, outputChannel++);
+                outputTypes.add(source.getTypes().get(source.getLayout().get(node.getGroupingSetMappings().get(output))));
             }
 
-            ImmutableList.Builder<Integer> copyChannels = ImmutableList.builder();
-            for (Symbol inputSymbol : node.getIdentityMappings().keySet()) {
-                int inputChannel = source.getLayout().get(inputSymbol);
-                newLayout.put(node.getIdentityMappings().get(inputSymbol), outputChannel++);
+            Map<Symbol, Integer> argumentMappings = new HashMap<>();
+            for (Symbol output : node.getArgumentMappings().keySet()) {
+                int inputChannel = source.getLayout().get(node.getArgumentMappings().get(output));
+
+                newLayout.put(output, outputChannel++);
                 outputTypes.add(source.getTypes().get(inputChannel));
-                copyChannels.add(inputChannel);
+                argumentMappings.put(output, inputChannel);
+            }
+
+            // for every grouping set, create a mapping of all output to input channels (including arguments)
+            ImmutableList.Builder<Map<Integer, Integer>> mappings = ImmutableList.builder();
+            for (List<Symbol> groupingSet : node.getGroupingSets()) {
+                ImmutableMap.Builder<Integer, Integer> setMapping = ImmutableMap.builder();
+
+                for (Symbol output : groupingSet) {
+                    setMapping.put(newLayout.get(output), source.getLayout().get(node.getGroupingSetMappings().get(output)));
+                }
+
+                for (Symbol output : argumentMappings.keySet()) {
+                    setMapping.put(newLayout.get(output), argumentMappings.get(output));
+                }
+
+                mappings.add(setMapping.build());
             }
 
             newLayout.put(node.getGroupIdSymbol(), outputChannel);
             outputTypes.add(BIGINT);
 
-            List<List<Integer>> groupingSetChannels = node.getGroupingSets().stream()
-                    .map(groupingSet -> getChannelsForSymbols(groupingSet, source.getLayout()))
-                    .collect(toImmutableList());
-
             OperatorFactory groupIdOperatorFactory = new GroupIdOperator.GroupIdOperatorFactory(context.getNextOperatorId(),
                     node.getId(),
                     outputTypes.build(),
-                    groupingSetChannels,
-                    groupingChannels.build(),
-                    copyChannels.build());
+                    mappings.build());
 
-            return new PhysicalOperation(groupIdOperatorFactory, newLayout.build(), source);
+            return new PhysicalOperation(groupIdOperatorFactory, newLayout, source);
         }
 
         @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanPrinter.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanPrinter.java
index b22f3d5665..84092bcff7 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanPrinter.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanPrinter.java
@@ -523,7 +523,23 @@ public class PlanPrinter
         @Override
         public Void visitGroupId(GroupIdNode node, Integer indent)
         {
-            print(indent, "- GroupId%s => [%s]", node.getGroupingSets(), formatOutputs(node.getOutputSymbols()));
+            // grouping sets are easier to understand in terms of inputs
+            List<List<Symbol>> inputGroupingSetSymbols = node.getGroupingSets().stream()
+                    .map(set -> set.stream()
+                            .map(symbol -> node.getGroupingSetMappings().get(symbol))
+                            .collect(Collectors.toList()))
+                    .collect(Collectors.toList());
+
+            print(indent, "- GroupId%s => [%s]", inputGroupingSetSymbols, formatOutputs(node.getOutputSymbols()));
+            printStats(indent + 2, node.getId());
+
+            for (Map.Entry<Symbol, Symbol> mapping : node.getGroupingSetMappings().entrySet()) {
+                print(indent + 2, "%s := %s", mapping.getKey(), mapping.getValue());
+            }
+            for (Map.Entry<Symbol, Symbol> argument : node.getArgumentMappings().entrySet()) {
+                print(indent + 2, "%s := %s", argument.getKey(), argument.getValue());
+            }
+
             return processChildren(node, indent + 1);
         }
 
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java
index 3c5978bedf..323048bc9f 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java
@@ -60,6 +60,7 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -383,47 +384,81 @@ class QueryPlanner
 
         // 2. Aggregate
 
-        // 2.a. Rewrite group by expressions in terms of pre-projected inputs
-        TranslationMap translations = new TranslationMap(subPlan.getRelationPlan(), analysis);
-        ImmutableList.Builder<List<Symbol>> groupingSetsSymbolsBuilder = ImmutableList.builder();
+        // 2.a. Rewrite aggregate arguments
+        TranslationMap argumentTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis);
+        ImmutableMap.Builder<Symbol, Symbol> argumentMappingBuilder = ImmutableMap.builder();
+        for (Expression argument : arguments) {
+            Expression parametersReplaced = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), argument);
+            argumentTranslations.addIntermediateMapping(argument, parametersReplaced);
+            Symbol input = subPlan.translate(parametersReplaced);
+
+            if (!argumentTranslations.containsSymbol(parametersReplaced)) {
+                Symbol output = symbolAllocator.newSymbol(parametersReplaced, analysis.getTypeWithCoercions(parametersReplaced), "arg");
+                argumentMappingBuilder.put(output, input);
+                argumentTranslations.put(parametersReplaced, output);
+            }
+        }
+        Map<Symbol, Symbol> argumentMappings  = argumentMappingBuilder.build();
+
+        // 2.b. Rewrite grouping columns
+        TranslationMap groupingTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis);
+        Map<Symbol, Symbol> groupingSetMappings = new HashMap<>();
+        List<List<Symbol>> groupingSymbols = new ArrayList<>();
+
         for (List<Expression> groupingSet : groupingSets) {
-            ImmutableList.Builder<Symbol> groupingColumns = ImmutableList.builder();
+            ImmutableList.Builder<Symbol> symbols = ImmutableList.builder();
             for (Expression expression : groupingSet) {
-                Expression rewritten = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), expression);
-                translations.addIntermediateMapping(expression, rewritten);
-                Symbol symbol = subPlan.translate(rewritten);
-                groupingColumns.add(symbol);
-                translations.put(rewritten, symbol);
+                Expression parametersReplaced = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), expression);
+                groupingTranslations.addIntermediateMapping(expression, parametersReplaced);
+                Symbol input = subPlan.translate(expression);
+
+                Symbol output;
+                if (!groupingTranslations.containsSymbol(parametersReplaced)) {
+                    output = symbolAllocator.newSymbol(parametersReplaced, analysis.getTypeWithCoercions(expression), "gid");
+                    groupingTranslations.put(parametersReplaced, output);
+                }
+                else {
+                    output = groupingTranslations.get(parametersReplaced);
+                }
+
+                groupingSetMappings.put(output, input);
+                symbols.add(output);
             }
-            groupingSetsSymbolsBuilder.add(groupingColumns.build());
+            groupingSymbols.add(symbols.build());
         }
 
-        // 2.b. Add a groupIdNode and groupIdSymbol if there are multiple grouping sets
+        // 2.c. Generate GroupIdNode (multiple grouping sets) or ProjectNode (single grouping set)
         Optional<Symbol> groupIdSymbol = Optional.empty();
-        List<List<Symbol>> groupingSetsSymbols = groupingSetsSymbolsBuilder.build();
         if (groupingSets.size() > 1) {
-            ImmutableMap.Builder<Symbol, Symbol> identityMapping = ImmutableMap.builder();
-            for (Expression argument : ImmutableSet.copyOf(arguments)) {
-                Symbol output = symbolAllocator.newSymbol(argument, analysis.getTypeWithCoercions(argument), "id");
-                identityMapping.put(subPlan.translate(argument), output);
+            groupIdSymbol = Optional.of(symbolAllocator.newSymbol("groupId", BIGINT));
+            GroupIdNode groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSymbols, groupingSetMappings, argumentMappings, groupIdSymbol.get());
+            subPlan = new PlanBuilder(groupingTranslations, groupId, analysis.getParameters());
+        }
+        else {
+            Map<Symbol, Expression> projections = new HashMap<>();
+            for (Symbol output : argumentMappings.keySet()) {
+                projections.putIfAbsent(output, argumentMappings.get(output).toSymbolReference());
+            }
 
-                // relies on the fact that group by expressions have already been re-written, and will not be affected by this mapping change
-                subPlan.getTranslations().put(argument, output);
+            for (Symbol output : groupingSetMappings.keySet()) {
+                projections.putIfAbsent(output, groupingSetMappings.get(output).toSymbolReference());
             }
 
-            groupIdSymbol = Optional.of(symbolAllocator.newSymbol("groupId", BIGINT));
-            GroupIdNode groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSetsSymbols, identityMapping.build(), groupIdSymbol.get());
-            subPlan = subPlan.withNewRoot(groupId);
+            ProjectNode project = new ProjectNode(idAllocator.getNextId(), subPlan.getRoot(), projections);
+            subPlan = new PlanBuilder(groupingTranslations, project, analysis.getParameters());
         }
 
-        // 2.c. Rewrite aggregates in terms of pre-projected inputs
+        TranslationMap aggregationTranslations = new TranslationMap(subPlan.getRelationPlan(), analysis);
+        aggregationTranslations.copyMappingsFrom(groupingTranslations);
+
+        // 2.d. Rewrite aggregates
         ImmutableMap.Builder<Symbol, FunctionCall> aggregationAssignments = ImmutableMap.builder();
         ImmutableMap.Builder<Symbol, Signature> functions = ImmutableMap.builder();
         boolean needPostProjectionCoercion = false;
         for (FunctionCall aggregate : analysis.getAggregates(node)) {
             Expression parametersReplaced = ExpressionTreeRewriter.rewriteWith(new ParameterRewriter(analysis.getParameters(), analysis), aggregate);
-            translations.addIntermediateMapping(aggregate, parametersReplaced);
-            Expression rewritten = subPlan.rewrite(parametersReplaced);
+            aggregationTranslations.addIntermediateMapping(aggregate, parametersReplaced);
+            Expression rewritten = argumentTranslations.rewrite(parametersReplaced);
             Symbol newSymbol = symbolAllocator.newSymbol(rewritten, analysis.getType(aggregate));
 
             // TODO: this is a hack, because we apply coercions to the output of expressions, rather than the arguments to expressions.
@@ -433,12 +468,12 @@ class QueryPlanner
                 needPostProjectionCoercion = true;
             }
             aggregationAssignments.put(newSymbol, (FunctionCall) rewritten);
-            translations.put(parametersReplaced, newSymbol);
+            aggregationTranslations.put(parametersReplaced, newSymbol);
 
             functions.put(newSymbol, analysis.getFunctionSignature(aggregate));
         }
 
-        // 2.d. Mark distinct rows for each aggregate that has DISTINCT
+        // 2.e. Mark distinct rows for each aggregate that has DISTINCT
         // Map from aggregate function arguments to marker symbols, so that we can reuse the markers, if two aggregates have the same argument
         Map<Set<Expression>, Symbol> argumentMarkers = new HashMap<>();
         // Map from aggregate functions to marker symbols
@@ -446,7 +481,7 @@ class QueryPlanner
         for (FunctionCall aggregate : Iterables.filter(analysis.getAggregates(node), FunctionCall::isDistinct)) {
             Set<Expression> args = ImmutableSet.copyOf(aggregate.getArguments());
             Symbol marker = argumentMarkers.get(args);
-            Symbol aggregateSymbol = translations.get(aggregate);
+            Symbol aggregateSymbol = aggregationTranslations.get(aggregate);
             if (marker == null) {
                 if (args.size() == 1) {
                     marker = symbolAllocator.newSymbol(getOnlyElement(args), BOOLEAN, "distinct");
@@ -462,14 +497,14 @@ class QueryPlanner
 
         for (Map.Entry<Set<Expression>, Symbol> entry : argumentMarkers.entrySet()) {
             ImmutableList.Builder<Symbol> builder = ImmutableList.builder();
-            builder.addAll(groupingSetsSymbols.stream()
+            builder.addAll(groupingSymbols.stream()
                     .flatMap(Collection::stream)
                     .distinct()
                     .collect(Collectors.toList()));
             groupIdSymbol.ifPresent(builder::add);
 
             for (Expression expression : entry.getKey()) {
-                builder.add(subPlan.translate(expression));
+                builder.add(argumentTranslations.get(expression));
             }
             subPlan = subPlan.withNewRoot(
                     new MarkDistinctNode(
@@ -486,12 +521,12 @@ class QueryPlanner
                 aggregationAssignments.build(),
                 functions.build(),
                 masks,
-                groupingSetsSymbols,
+                groupingSymbols,
                 AggregationNode.Step.SINGLE,
                 Optional.empty(),
                 groupIdSymbol);
 
-        subPlan = new PlanBuilder(translations, aggregationNode, analysis.getParameters());
+        subPlan = new PlanBuilder(aggregationTranslations, aggregationNode, analysis.getParameters());
 
         // 3. Post-projection
         // Add back the implicit casts that we removed in 2.a
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/SymbolExtractor.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/SymbolExtractor.java
index 61c18e2263..129fee725e 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/SymbolExtractor.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/SymbolExtractor.java
@@ -113,7 +113,8 @@ public final class SymbolExtractor
             node.getSource().accept(this, context);
 
             builder.add(node.getGroupIdSymbol());
-            builder.addAll(node.getIdentityMappings().values());
+            builder.addAll(node.getGroupingSetMappings().keySet());
+            builder.addAll(node.getArgumentMappings().keySet());
 
             return null;
         }
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
index a444ff1d20..86bac59425 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
@@ -312,15 +312,13 @@ public class AddExchanges
 
         private Function<Symbol, Optional<Symbol>> translateGroupIdSymbols(GroupIdNode node)
         {
-            Map<Symbol, Symbol> invertedMappings = ImmutableBiMap.copyOf(node.getIdentityMappings()).inverse();
-            List<Symbol> commonGroupingColumns = node.getCommonGroupingColumns();
             return symbol -> {
-                if (invertedMappings.containsKey(symbol)) {
-                    return Optional.of(invertedMappings.get(symbol));
+                if (node.getArgumentMappings().containsKey(symbol)) {
+                    return Optional.of(node.getArgumentMappings().get(symbol));
                 }
 
-                if (commonGroupingColumns.contains(symbol)) {
-                    return Optional.of(symbol);
+                if (node.getCommonGroupingColumns().contains(symbol)) {
+                    return Optional.of(node.getGroupingSetMappings().get(symbol));
                 }
 
                 return Optional.empty();
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PredicatePushDown.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PredicatePushDown.java
index 2753a999b2..3495784790 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PredicatePushDown.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PredicatePushDown.java
@@ -30,7 +30,6 @@ import com.facebook.presto.sql.planner.Symbol;
 import com.facebook.presto.sql.planner.SymbolAllocator;
 import com.facebook.presto.sql.planner.plan.AggregationNode;
 import com.facebook.presto.sql.planner.plan.AssignUniqueId;
-import com.facebook.presto.sql.planner.plan.ChildReplacer;
 import com.facebook.presto.sql.planner.plan.ExchangeNode;
 import com.facebook.presto.sql.planner.plan.FilterNode;
 import com.facebook.presto.sql.planner.plan.GroupIdNode;
@@ -221,26 +220,25 @@ public class PredicatePushDown
         {
             checkState(!DependencyExtractor.extractUnique(context.get()).contains(node.getGroupIdSymbol()), "groupId symbol cannot be referenced in predicate");
 
-            List<Symbol> commonGroupingSymbols = node.getCommonGroupingColumns();
+            Map<Symbol, SymbolReference> commonGroupingSymbolMapping = node.getGroupingSetMappings().entrySet().stream()
+                    .filter(entry -> node.getCommonGroupingColumns().contains(entry.getKey()))
+                    .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().toSymbolReference()));
+
             Predicate<Expression> pushdownEligiblePredicate = conjunct -> DependencyExtractor.extractUnique(conjunct).stream()
-                    .allMatch(commonGroupingSymbols::contains);
+                    .allMatch(commonGroupingSymbolMapping.keySet()::contains);
 
             Map<Boolean, List<Expression>> conjuncts = extractConjuncts(context.get()).stream().collect(Collectors.partitioningBy(pushdownEligiblePredicate));
 
-            // Push down conjuncts from the inherited predicate that apply to the common grouping columns, or don't apply to any grouping columns
-            PlanNode rewrittenSource = context.rewrite(node.getSource(), combineConjuncts(conjuncts.get(true)));
-
-            PlanNode output = node;
-            if (rewrittenSource != node.getSource()) {
-                output = ChildReplacer.replaceChildren(node, ImmutableList.of(rewrittenSource));
-            }
+            // Push down conjuncts from the inherited predicate that apply to common grouping symbols
+            PlanNode rewrittenNode = context.defaultRewrite(node,
+                    ExpressionTreeRewriter.rewriteWith(new ExpressionSymbolInliner(commonGroupingSymbolMapping), combineConjuncts(conjuncts.get(true))));
 
             // All other conjuncts, if any, will be in the filter node.
             if (!conjuncts.get(false).isEmpty()) {
-                output = new FilterNode(idAllocator.getNextId(), output, combineConjuncts(conjuncts.get(false)));
+                rewrittenNode = new FilterNode(idAllocator.getNextId(), rewrittenNode, combineConjuncts(conjuncts.get(false)));
             }
 
-            return output;
+            return rewrittenNode;
         }
 
         @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
index 0d8aa9d9a5..8c8f3537df 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
@@ -76,7 +76,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.Function;
 
 import static com.facebook.presto.SystemSessionProperties.planWithTableNodePartitioning;
 import static com.facebook.presto.spi.predicate.TupleDomain.extractFixedValues;
@@ -222,23 +221,22 @@ class PropertyDerivations
         @Override
         public ActualProperties visitGroupId(GroupIdNode node, List<ActualProperties> inputProperties)
         {
-            return Iterables.getOnlyElement(inputProperties).translate(translateGroupIdSymbols(node));
-        }
-
-        private Function<Symbol, Optional<Symbol>> translateGroupIdSymbols(GroupIdNode node)
-        {
-            List<Symbol> commonGroupingColumns = node.getCommonGroupingColumns();
-            return symbol -> {
-                if (node.getIdentityMappings().containsKey(symbol)) {
-                    return Optional.of(node.getIdentityMappings().get(symbol));
+            Map<Symbol, Symbol> inputToOutputMappings = new HashMap<>();
+            for (Map.Entry<Symbol, Symbol> setMapping : node.getGroupingSetMappings().entrySet()) {
+                if (node.getCommonGroupingColumns().contains(setMapping.getKey())) {
+                    // TODO: Add support for translating a property on a single column to multiple columns
+                    // when GroupIdNode is copying a single input grouping column into multiple output grouping columns (i.e. aliases), this is basically picking one arbitrarily
+                    inputToOutputMappings.putIfAbsent(setMapping.getValue(), setMapping.getKey());
                 }
+            }
 
-                if (commonGroupingColumns.contains(symbol)) {
-                    return Optional.of(symbol);
-                }
+            // TODO: Add support for translating a property on a single column to multiple columns
+            // this is deliberately placed after the grouping columns, because preserving properties has a bigger perf impact
+            for (Map.Entry<Symbol, Symbol> argumentMapping : node.getArgumentMappings().entrySet()) {
+                inputToOutputMappings.putIfAbsent(argumentMapping.getValue(), argumentMapping.getKey());
+            }
 
-                return Optional.empty();
-            };
+            return Iterables.getOnlyElement(inputProperties).translate(column -> Optional.ofNullable(inputToOutputMappings.get(column)));
         }
 
         @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java
index 9658c59f7f..41eb36742b 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PruneUnreferencedOutputs.java
@@ -68,15 +68,16 @@ import com.google.common.collect.Sets;
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableSet;
-import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Predicates.in;
 import static com.google.common.collect.Iterables.concat;
 import static java.util.Objects.requireNonNull;
@@ -408,23 +409,31 @@ public class PruneUnreferencedOutputs
         @Override
         public PlanNode visitGroupId(GroupIdNode node, RewriteContext<Set<Symbol>> context)
         {
-            checkState(node.getDistinctGroupingColumns().stream().allMatch(column -> context.get().contains(column)));
+            ImmutableSet.Builder<Symbol> expectedInputs = ImmutableSet.builder();
 
-            ImmutableMap.Builder<Symbol, Symbol> identityMappingBuilder = ImmutableMap.builder();
-            for (Map.Entry<Symbol, Symbol> entry : node.getIdentityMappings().entrySet()) {
-                if (context.get().contains(entry.getValue())) {
-                    identityMappingBuilder.put(entry);
-                }
-            }
+            Map<Symbol, Symbol> newArgumentMappings = node.getArgumentMappings().entrySet().stream()
+                    .filter(entry -> context.get().contains(entry.getKey()))
+                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+            expectedInputs.addAll(newArgumentMappings.values());
+
+            ImmutableList.Builder<List<Symbol>> newGroupingSets = ImmutableList.builder();
+            Map<Symbol, Symbol> newGroupingMapping = new HashMap<>();
 
-            Map<Symbol, Symbol> identityMapping = identityMappingBuilder.build();
+            for (List<Symbol> groupingSet : node.getGroupingSets()) {
+                ImmutableList.Builder<Symbol> newGroupingSet = ImmutableList.builder();
 
-            PlanNode source = context.rewrite(node.getSource(), ImmutableSet.<Symbol>builder()
-                    .addAll(identityMapping.keySet())
-                    .addAll(node.getDistinctGroupingColumns())
-                    .build());
+                for (Symbol output : groupingSet) {
+                    if (context.get().contains(output)) {
+                        newGroupingSet.add(output);
+                        newGroupingMapping.putIfAbsent(output, node.getGroupingSetMappings().get(output));
+                        expectedInputs.add(node.getGroupingSetMappings().get(output));
+                    }
+                }
+                newGroupingSets.add(newGroupingSet.build());
+            }
 
-            return new GroupIdNode(node.getId(), source, node.getGroupingSets(), identityMapping, node.getGroupIdSymbol());
+            PlanNode source = context.rewrite(node.getSource(), expectedInputs.build());
+            return new GroupIdNode(node.getId(), source, newGroupingSets.build(), newGroupingMapping, newArgumentMappings, node.getGroupIdSymbol());
         }
 
         @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
index 62c639e66f..c0977978f1 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/StreamPropertyDerivations.java
@@ -307,23 +307,22 @@ final class StreamPropertyDerivations
         @Override
         public StreamProperties visitGroupId(GroupIdNode node, List<StreamProperties> inputProperties)
         {
-            return Iterables.getOnlyElement(inputProperties).translate(translateGroupIdSymbols(node));
-        }
-
-        private Function<Symbol, Optional<Symbol>> translateGroupIdSymbols(GroupIdNode node)
-        {
-            List<Symbol> commonGroupingColumns = node.getCommonGroupingColumns();
-            return symbol -> {
-                if (node.getIdentityMappings().containsKey(symbol)) {
-                    return Optional.of(node.getIdentityMappings().get(symbol));
+            Map<Symbol, Symbol> inputToOutputMappings = new HashMap<>();
+            for (Map.Entry<Symbol, Symbol> setMapping : node.getGroupingSetMappings().entrySet()) {
+                if (node.getCommonGroupingColumns().contains(setMapping.getKey())) {
+                    // TODO: Add support for translating a property on a single column to multiple columns
+                    // when GroupIdNode is copying a single input grouping column into multiple output grouping columns (i.e. aliases), this is basically picking one arbitrarily
+                    inputToOutputMappings.putIfAbsent(setMapping.getValue(), setMapping.getKey());
                 }
+            }
 
-                if (commonGroupingColumns.contains(symbol)) {
-                    return Optional.of(symbol);
-                }
+            // TODO: Add support for translating a property on a single column to multiple columns
+            // this is deliberately placed after the grouping columns, because preserving properties has a bigger perf impact
+            for (Map.Entry<Symbol, Symbol> argumentMapping : node.getArgumentMappings().entrySet()) {
+                inputToOutputMappings.putIfAbsent(argumentMapping.getValue(), argumentMapping.getKey());
+            }
 
-                return Optional.empty();
-            };
+            return Iterables.getOnlyElement(inputProperties).translate(column -> Optional.ofNullable(inputToOutputMappings.get(column)));
         }
 
         @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
index 3761439cae..153bd13229 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/UnaliasSymbolReferences.java
@@ -81,7 +81,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableSet;
@@ -157,16 +156,31 @@ public class UnaliasSymbolReferences
         public PlanNode visitGroupId(GroupIdNode node, RewriteContext<Void> context)
         {
             PlanNode source = context.rewrite(node.getSource());
-            List<List<Symbol>> groupingSetsSymbols = node.getGroupingSets().stream()
-                    .map(this::canonicalize)
-                    .collect(Collectors.toList());
 
-            ImmutableMap.Builder<Symbol, Symbol> newPassthroughMap = ImmutableMap.builder();
-            for (Symbol inputSymbol : node.getIdentityMappings().keySet()) {
-                newPassthroughMap.put(canonicalize(inputSymbol), canonicalize(node.getIdentityMappings().get(inputSymbol)));
+            Map<Symbol, Symbol> newGroupingMappings = new HashMap<>();
+            ImmutableList.Builder<List<Symbol>> newGroupingSets = ImmutableList.builder();
+
+            for (List<Symbol> groupingSet : node.getGroupingSets()) {
+                ImmutableList.Builder<Symbol> newGroupingSet = ImmutableList.builder();
+                for (Symbol output : groupingSet) {
+                    newGroupingMappings.putIfAbsent(canonicalize(output), canonicalize(node.getGroupingSetMappings().get(output)));
+                    newGroupingSet.add(canonicalize(output));
+                }
+                newGroupingSets.add(newGroupingSet.build());
+            }
+
+            Map<Symbol, Symbol> newArgumentMappings = new HashMap<>();
+            for (Symbol output : node.getArgumentMappings().keySet()) {
+                Symbol canonicalOutput = canonicalize(output);
+                if (newArgumentMappings.containsKey(canonicalOutput)) {
+                    map(output, canonicalOutput);
+                }
+                else {
+                    newArgumentMappings.put(canonicalOutput, canonicalize(node.getArgumentMappings().get(output)));
+                }
             }
 
-            return new GroupIdNode(node.getId(), source, groupingSetsSymbols, newPassthroughMap.build(), canonicalize(node.getGroupIdSymbol()));
+            return new GroupIdNode(node.getId(), source, newGroupingSets.build(), newGroupingMappings, newArgumentMappings, canonicalize(node.getGroupIdSymbol()));
         }
 
         @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ChildReplacer.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ChildReplacer.java
index 41f8c140f5..a760186471 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ChildReplacer.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ChildReplacer.java
@@ -169,7 +169,7 @@ public class ChildReplacer
     @Override
     public PlanNode visitGroupId(GroupIdNode node, List<PlanNode> newChildren)
     {
-        return new GroupIdNode(node.getId(), Iterables.getOnlyElement(newChildren), node.getGroupingSets(), node.getIdentityMappings(), node.getGroupIdSymbol());
+        return new GroupIdNode(node.getId(), Iterables.getOnlyElement(newChildren), node.getGroupingSets(), node.getGroupingSetMappings(), node.getArgumentMappings(), node.getGroupIdSymbol());
     }
 
     @Override
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/GroupIdNode.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/GroupIdNode.java
index daf3560ced..1dc5bb48d4 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/GroupIdNode.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/plan/GroupIdNode.java
@@ -19,6 +19,7 @@ import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Sets;
 
 import javax.annotation.concurrent.Immutable;
 
@@ -27,39 +28,53 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
-import static java.util.stream.Collectors.toList;
+import static java.util.stream.Collectors.toSet;
 
 @Immutable
 public class GroupIdNode
         extends PlanNode
 {
     private final PlanNode source;
+
+    // in terms of output symbols
     private final List<List<Symbol>> groupingSets;
-    private final Map<Symbol, Symbol> identityMappings;
+
+    // from output to input symbols
+    private final Map<Symbol, Symbol> groupingSetMappings;
+    private final Map<Symbol, Symbol> argumentMappings;
+
     private final Symbol groupIdSymbol;
 
     @JsonCreator
     public GroupIdNode(@JsonProperty("id") PlanNodeId id,
             @JsonProperty("source") PlanNode source,
             @JsonProperty("groupingSets") List<List<Symbol>> groupingSets,
-            @JsonProperty("identityMappings") Map<Symbol, Symbol> identityMappings,
+            @JsonProperty("groupingSetMappings") Map<Symbol, Symbol> groupingSetMappings,
+            @JsonProperty("argumentMappings") Map<Symbol, Symbol> argumentMappings,
             @JsonProperty("groupIdSymbol") Symbol groupIdSymbol)
     {
         super(id);
         this.source = requireNonNull(source);
         this.groupingSets = ImmutableList.copyOf(requireNonNull(groupingSets));
-        this.identityMappings = ImmutableMap.copyOf(requireNonNull(identityMappings));
+        this.groupingSetMappings = ImmutableMap.copyOf(requireNonNull(groupingSetMappings));
+        this.argumentMappings = ImmutableMap.copyOf(requireNonNull(argumentMappings));
         this.groupIdSymbol = requireNonNull(groupIdSymbol);
+
+        checkArgument(Sets.intersection(groupingSetMappings.keySet(), argumentMappings.keySet()).isEmpty(), "argument outputs and grouping outputs must be a disjoint set");
     }
 
     @Override
     public List<Symbol> getOutputSymbols()
     {
         return ImmutableList.<Symbol>builder()
-                .addAll(getDistinctGroupingColumns())
-                .addAll(identityMappings.values())
+                .addAll(groupingSets.stream()
+                        .flatMap(Collection::stream)
+                        .collect(toSet()))
+                .addAll(argumentMappings.keySet())
                 .add(groupIdSymbol)
                 .build();
     }
@@ -76,14 +91,6 @@ public class GroupIdNode
         return source;
     }
 
-    public Set<Symbol> getInputSymbols()
-    {
-        return ImmutableSet.<Symbol>builder()
-                .addAll(identityMappings.keySet())
-                .addAll(getDistinctGroupingColumns())
-                .build();
-    }
-
     @JsonProperty
     public List<List<Symbol>> getGroupingSets()
     {
@@ -91,26 +98,15 @@ public class GroupIdNode
     }
 
     @JsonProperty
-    public Map<Symbol, Symbol> getIdentityMappings()
+    public Map<Symbol, Symbol> getGroupingSetMappings()
     {
-        return identityMappings;
+        return groupingSetMappings;
     }
 
-    public List<Symbol> getDistinctGroupingColumns()
+    @JsonProperty
+    public Map<Symbol, Symbol> getArgumentMappings()
     {
-        return groupingSets.stream()
-                .flatMap(Collection::stream)
-                .distinct()
-                .collect(toList());
-    }
-
-    public List<Symbol> getCommonGroupingColumns()
-    {
-        Set<Symbol> intersection = new HashSet<>(groupingSets.get(0));
-        for (int i = 1; i < getGroupingSets().size(); i++) {
-            intersection.retainAll(groupingSets.get(i));
-        }
-        return ImmutableList.copyOf(intersection);
+        return argumentMappings;
     }
 
     @JsonProperty
@@ -124,4 +120,26 @@ public class GroupIdNode
     {
         return visitor.visitGroupId(this, context);
     }
+
+    public Set<Symbol> getInputSymbols()
+    {
+        return ImmutableSet.<Symbol>builder()
+                .addAll(argumentMappings.values())
+                .addAll(groupingSets.stream()
+                        .map(set -> set.stream()
+                                .map(groupingSetMappings::get).collect(Collectors.toList()))
+                        .flatMap(Collection::stream)
+                        .collect(toSet()))
+                .build();
+    }
+
+    // returns the common grouping columns in terms of output symbols
+    public Set<Symbol> getCommonGroupingColumns()
+    {
+        Set<Symbol> intersection = new HashSet<>(groupingSets.get(0));
+        for (int i = 1; i < groupingSets.size(); i++) {
+            intersection.retainAll(groupingSets.get(i));
+        }
+        return ImmutableSet.copyOf(intersection);
+    }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java b/presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java
index eb7703171e..c5b7677248 100644
--- a/presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java
+++ b/presto-main/src/main/java/com/facebook/presto/util/GraphvizPrinter.java
@@ -332,11 +332,14 @@ public final class GraphvizPrinter
         @Override
         public Void visitGroupId(GroupIdNode node, Void context)
         {
-            List<String> groupingSets = node.getGroupingSets().stream()
-                    .map(groupingSet -> "(" + Joiner.on(", ").join(groupingSet) + ")")
+            // grouping sets are easier to understand in terms of inputs
+            List<String> inputGroupingSetSymbols = node.getGroupingSets().stream()
+                    .map(set -> "(" + Joiner.on(", ").join(set.stream()
+                            .map(symbol -> node.getGroupingSetMappings().get(symbol))
+                            .collect(Collectors.toList())) + ")")
                     .collect(Collectors.toList());
 
-            printNode(node, "GroupId", Joiner.on(", ").join(groupingSets), NODE_COLORS.get(NodeType.AGGREGATE));
+            printNode(node, "GroupId", Joiner.on(", ").join(inputGroupingSetSymbols), NODE_COLORS.get(NodeType.AGGREGATE));
             return node.getSource().accept(this, context);
         }
 
diff --git a/presto-main/src/test/java/com/facebook/presto/operator/TestGroupIdOperator.java b/presto-main/src/test/java/com/facebook/presto/operator/TestGroupIdOperator.java
index e366afa8bc..53efa7593c 100644
--- a/presto-main/src/test/java/com/facebook/presto/operator/TestGroupIdOperator.java
+++ b/presto-main/src/test/java/com/facebook/presto/operator/TestGroupIdOperator.java
@@ -19,6 +19,7 @@ import com.facebook.presto.spi.Page;
 import com.facebook.presto.sql.planner.plan.PlanNodeId;
 import com.facebook.presto.testing.MaterializedResult;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -72,9 +73,7 @@ public class TestGroupIdOperator
                 new GroupIdOperatorFactory(0,
                         new PlanNodeId("test"),
                         ImmutableList.of(VARCHAR, BOOLEAN, BIGINT, BIGINT, BIGINT),
-                        ImmutableList.of(ImmutableList.of(1, 2), ImmutableList.of(3)),
-                        ImmutableList.of(1, 2, 3),
-                        ImmutableList.of(0));
+                        ImmutableList.of(ImmutableMap.of(0, 1, 1, 2, 3, 0), ImmutableMap.of(2, 3, 3, 0)));
 
         MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR, BOOLEAN, BIGINT, BIGINT, BIGINT)
                 .row("400", true, null, 100L, 0L)
diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
index bfbd609fb5..adbff45529 100644
--- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
+++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
@@ -1516,6 +1516,19 @@ public abstract class AbstractTestQueries
                         "SELECT SUM(CAST(quantity AS BIGINT)) FROM lineitem WHERE quantity < 0");
     }
 
+    @Test
+    public void testGroupingSetsAliasedGroupingColumns()
+            throws Exception
+    {
+        assertQuery("SELECT lna, lnb, SUM(quantity) " +
+                        "FROM (SELECT linenumber lna, linenumber lnb, CAST(quantity AS BIGINT) quantity FROM lineitem) " +
+                        "GROUP BY GROUPING SETS ((lna, lnb), (lna), (lnb), ())",
+                "SELECT linenumber, linenumber, SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY linenumber UNION ALL " +
+                        "SELECT linenumber, NULL, SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY linenumber UNION ALL " +
+                        "SELECT NULL, linenumber, SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY linenumber UNION ALL " +
+                        "SELECT NULL, NULL, SUM(CAST(quantity AS BIGINT)) FROM lineitem");
+    }
+
     @Test
     public void testGroupingSetMixedExpressionAndColumn()
             throws Exception