Repository: prestodb/presto
Issue: Presto doesn't correctly deal with complex aggregations when arguments are also group by columns

Hi,

I might found a bug related to ROLLUP function.
Could you look into this issue?
### Problem

ROLLUP function return 0 as total count if `COUNT(column)` is used.
If `COUNT(*)` is used, Presto returns a correct result.
#### PostgreSQL 9.5.2

```
$ psql
psql (9.5.2)
Type "help" for help.

postgres=# WITH test_data AS (SELECT * FROM (
postgres(# VALUES
postgres(# ('a', 1),
postgres(# ('b', 2),
postgres(# ('a', 3)
postgres(# ) as t(id, cnt)
postgres(# )
postgres-# SELECT id, COUNT(*) as s_cnt
postgres-# FROM test_data
postgres-# GROUP BY ROLLUP(id)
postgres-# ;
 id | s_cnt
----+-------
 a  |     2
 b  |     1
    |     3
(3 rows)
```

```
postgres=# WITH test_data AS (SELECT * FROM (
postgres(# VALUES
postgres(# ('a', 1),
postgres(# ('b', 2),
postgres(# ('a', 3)
postgres(# ) as t(id, cnt)
postgres(# )
postgres-# SELECT id, COUNT(id) as s_cnt
postgres-# FROM test_data
postgres-# GROUP BY ROLLUP(id);
 id | s_cnt
----+-------
 a  |     2
 b  |     1
    |     3
(3 rows)
```
### Presto 0.146

```
> WITH test_data AS (SELECT * FROM (
> VALUES
> ('a', 1),
> ('b', 2),
> ('a', 3)
> ) as t(id, cnt)
> )
> SELECT id, COUNT(*) as s_cnt
> FROM test_data
> GROUP BY ROLLUP(id);
> 
+------+-------+
| id   | s_cnt |
+------+-------+
| null | 3     |
| a    | 2     |
| b    | 1     |
+------+-------+
3 rows in set
```

_The following SQL show this problem_.
`null` returns `0`

```
> WITH test_data AS (SELECT * FROM (
> VALUES
> ('a', 1),
> ('b', 2),
> ('a', 3)
> ) as t(id, cnt)
> )
> SELECT id, COUNT(id) as s_cnt
> FROM test_data
> GROUP BY ROLLUP(id);

+------+-------+
| id   | s_cnt |
+------+-------+
| null | 0     |
| a    | 2     |
| b    | 1     |
+------+-------+
3 rows in set
```


Changed files:
presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java
presto-main/src/main/java/com/facebook/presto/sql/planner/plan/ChildReplacer.java
presto-main/src/main/java/com/facebook/presto/sql/planner/LocalExecutionPlanner.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/plan/GroupIdNode.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 d4266a86b1..18653de108 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
@@ -41,25 +41,27 @@ public class GroupIdOperator
     {
         private final int operatorId;
         private final PlanNodeId planNodeId;
-        private final List<Type> inputTypes;
         private final List<Type> outputTypes;
         private final List<List<Integer>> groupingSetChannels;
+        private final List<Integer> groupingChannels;
+        private final List<Integer> copyChannels;
 
         private boolean closed;
 
         public GroupIdOperatorFactory(
                 int operatorId,
                 PlanNodeId planNodeId,
-                List<? extends Type> inputTypes,
-                List<List<Integer>> groupingSetChannels)
+                List<? extends Type> outputTypes,
+                List<List<Integer>> groupingSetChannels,
+                List<Integer> groupingChannels,
+                List<Integer> copyChannels)
         {
             this.operatorId = operatorId;
             this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
+            this.outputTypes = ImmutableList.copyOf(requireNonNull(outputTypes));
             this.groupingSetChannels = ImmutableList.copyOf(requireNonNull(groupingSetChannels));
-            this.inputTypes = ImmutableList.copyOf(requireNonNull(inputTypes));
-
-            // add the groupId channel to the output types
-            this.outputTypes = ImmutableList.<Type>builder().addAll(inputTypes).add(BIGINT).build();
+            this.groupingChannels = ImmutableList.copyOf(requireNonNull(groupingChannels));
+            this.copyChannels = ImmutableList.copyOf(requireNonNull(copyChannels));
         }
 
         @Override
@@ -78,10 +80,11 @@ public class GroupIdOperator
                     .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(inputTypes.size());
+                groupingSetNullChannels[i] = new BitSet(groupingChannels.size() + copyChannels.size());
                 // first set all grouping columns to true
                 for (int groupingColumn : allGroupingColumns) {
                     groupingSetNullChannels[i].set(groupingColumn, true);
@@ -92,13 +95,15 @@ public class GroupIdOperator
                 }
             }
 
-            Block[] nullBlocks = new Block[inputTypes.size()];
-            for (int i = 0; i < nullBlocks.length; i++) {
-                nullBlocks[i] = inputTypes.get(i).createBlockBuilder(new BlockBuilderStatus(), 1)
+            // create null blocks for every grouping channel
+            Block[] nullBlocks = new Block[groupingChannels.size()];
+            for (int i = 0; i < groupingChannels.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++) {
                 BlockBuilder builder = BIGINT.createBlockBuilder(new BlockBuilderStatus(), 1);
@@ -106,7 +111,13 @@ public class GroupIdOperator
                 groupIdBlocks[i] = builder.build();
             }
 
-            return new GroupIdOperator(operatorContext, outputTypes, groupingSetNullChannels, nullBlocks, groupIdBlocks);
+            // 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);
         }
 
         @Override
@@ -118,7 +129,7 @@ public class GroupIdOperator
         @Override
         public OperatorFactory duplicate()
         {
-            return new GroupIdOperatorFactory(operatorId, planNodeId, inputTypes, groupingSetChannels);
+            return new GroupIdOperatorFactory(operatorId, planNodeId, outputTypes, groupingSetChannels, groupingChannels, copyChannels);
         }
     }
 
@@ -127,6 +138,8 @@ public class GroupIdOperator
     private final BitSet[] groupingSetNullChannels;
     private final Block[] nullBlocks;
     private final Block[] groupIdBlocks;
+    private final int[] groupInputs;
+    private final int[] copyInputs;
 
     private Page currentPage = null;
     private int currentGroupingSet = 0;
@@ -137,15 +150,18 @@ public class GroupIdOperator
             List<Type> types,
             BitSet[] groupingSetNullChannels,
             Block[] nullBlocks,
-            Block[] groupIdBlocks)
+            Block[] groupIdBlocks,
+            int[] groupInputs,
+            int[] copyInputs)
     {
         this.operatorContext = requireNonNull(operatorContext, "operatorContext is null");
-        this.types = requireNonNull(types, "inputTypes is null");
+        this.types = requireNonNull(types, "types is null");
         this.groupingSetNullChannels =  requireNonNull(groupingSetNullChannels, "groupingSetNullChannels is null");
         this.nullBlocks = requireNonNull(nullBlocks);
-        checkArgument(nullBlocks.length == (types.size() - 1), "length of nullBlocks must be one plus length of types");
         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);
     }
 
     @Override
@@ -200,17 +216,21 @@ public class GroupIdOperator
     private Page generateNextPage()
     {
         // generate 'n' pages for every input page, where n is the number of grouping sets
-        Block[] outputBlocks = new Block[currentPage.getChannelCount() + 1];
+        Block[] outputBlocks = new Block[types.size()];
 
-        for (int channel = 0; channel < currentPage.getChannelCount(); channel++) {
-            if (groupingSetNullChannels[currentGroupingSet].get(channel)) {
-                outputBlocks[channel] = new RunLengthEncodedBlock(nullBlocks[channel], currentPage.getPositionCount());
+        for (int i = 0; i < groupInputs.length; i++) {
+            if (groupingSetNullChannels[currentGroupingSet].get(groupInputs[i])) {
+                outputBlocks[i] = new RunLengthEncodedBlock(nullBlocks[i], currentPage.getPositionCount());
             }
             else {
-                outputBlocks[channel] = currentPage.getBlock(channel);
+                outputBlocks[i] = currentPage.getBlock(groupInputs[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;
         Page outputPage = new Page(currentPage.getPositionCount(), outputBlocks);
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 7689c44064..4b29a1e346 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
@@ -818,24 +818,42 @@ public class LocalExecutionPlanner
         public PhysicalOperation visitGroupId(GroupIdNode node, LocalExecutionPlanContext context)
         {
             PhysicalOperation source = node.getSource().accept(this, context);
+            ImmutableMap.Builder<Symbol, Integer> newLayout = ImmutableMap.builder();
+            ImmutableList.Builder<Type> outputTypes = ImmutableList.builder();
 
-            // add groupId to the layout
-            int groupIdChannel = source.getLayout().values().stream()
-                    .mapToInt(Integer::intValue)
-                    .max()
-                    .orElse(-1) + 1;
+            int outputChannel = 0;
 
-            Map<Symbol, Integer> newLayout = ImmutableMap.<Symbol, Integer>builder()
-                    .putAll(source.getLayout())
-                    .put(node.getGroupIdSymbol(), groupIdChannel)
-                    .build();
+            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);
+            }
+
+            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++);
+                outputTypes.add(source.getTypes().get(inputChannel));
+                copyChannels.add(inputChannel);
+            }
+
+            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(), source.getTypes(), groupingSetChannels);
-            return new PhysicalOperation(groupIdOperatorFactory, newLayout, source);
+            OperatorFactory groupIdOperatorFactory = new GroupIdOperator.GroupIdOperatorFactory(context.getNextOperatorId(),
+                    node.getId(),
+                    outputTypes.build(),
+                    groupingSetChannels,
+                    groupingChannels.build(),
+                    copyChannels.build());
+
+            return new PhysicalOperation(groupIdOperatorFactory, newLayout.build(), source);
         }
 
         @Override
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 c6e0f073e4..02264cd077 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.Iterables;
 
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -369,29 +370,9 @@ class QueryPlanner
         }
 
         // 2. Aggregate
-        ImmutableMap.Builder<Symbol, FunctionCall> aggregationAssignments = ImmutableMap.builder();
-        ImmutableMap.Builder<Symbol, Signature> functions = ImmutableMap.builder();
 
-        // 2.a. Rewrite aggregates in terms of pre-projected inputs
+        // 2.a. Rewrite group by expressions in terms of pre-projected inputs
         TranslationMap translations = new TranslationMap(subPlan.getRelationPlan(), analysis);
-        boolean needPostProjectionCoercion = false;
-        for (FunctionCall aggregate : analysis.getAggregates(node)) {
-            Expression rewritten = subPlan.rewrite(aggregate);
-            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.
-            // Therefore we can end up with this implicit cast, and have to move it into a post-projection
-            if (rewritten instanceof Cast) {
-                rewritten = ((Cast) rewritten).getExpression();
-                needPostProjectionCoercion = true;
-            }
-            aggregationAssignments.put(newSymbol, (FunctionCall) rewritten);
-            translations.put(aggregate, newSymbol);
-
-            functions.put(newSymbol, analysis.getFunctionSignature(aggregate));
-        }
-
-        // 2.b. Rewrite group by expressions in terms of pre-projected inputs
         ImmutableList.Builder<List<Symbol>> groupingSetsSymbolsBuilder = ImmutableList.builder();
         ImmutableSet.Builder<Symbol> distinctGroupingSymbolsBuilder = ImmutableSet.builder();
         for (List<Expression> groupingSet : groupingSets) {
@@ -405,15 +386,49 @@ class QueryPlanner
             groupingSetsSymbolsBuilder.add(groupingColumns.build());
         }
 
+        // 2.b. Add a groupIdNode and groupIdSymbol if there are multiple grouping sets
         List<List<Symbol>> groupingSetsSymbols = groupingSetsSymbolsBuilder.build();
-        // 2.c. Add a groupIdNode and groupIdSymbol if there are multiple grouping sets
         if (groupingSets.size() > 1) {
+            Set<Symbol> groupIdInputs = new HashSet<>();
+            distinctGroupingColumns.stream()
+                    .map(subPlan::translate)
+                    .forEach(groupIdInputs::add);
+
+            ImmutableMap.Builder<Symbol, Symbol> identityMapping = ImmutableMap.builder();
+            for (Expression argument : arguments) {
+                Symbol output = symbolAllocator.newSymbol(argument, analysis.getTypeWithCoercions(argument), "id");
+                identityMapping.put(subPlan.translate(argument), output);
+                groupIdInputs.add(subPlan.translate(argument));
+
+                // 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);
+            }
+
             Symbol groupIdSymbol = symbolAllocator.newSymbol("groupId", BIGINT);
-            GroupIdNode groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), subPlan.getRoot().getOutputSymbols(), groupingSetsSymbols, groupIdSymbol);
+            GroupIdNode groupId = new GroupIdNode(idAllocator.getNextId(), subPlan.getRoot(), groupingSetsSymbols, identityMapping.build(), groupIdSymbol);
             subPlan = subPlan.withNewRoot(groupId);
             distinctGroupingSymbolsBuilder.add(groupIdSymbol);
         }
-        List<Symbol> distinctGroupingSymbols = distinctGroupingSymbolsBuilder.build().asList();
+
+        // 2.c. Rewrite aggregates in terms of pre-projected inputs
+        ImmutableMap.Builder<Symbol, FunctionCall> aggregationAssignments = ImmutableMap.builder();
+        ImmutableMap.Builder<Symbol, Signature> functions = ImmutableMap.builder();
+        boolean needPostProjectionCoercion = false;
+        for (FunctionCall aggregate : analysis.getAggregates(node)) {
+            Expression rewritten = subPlan.rewrite(aggregate);
+            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.
+            // Therefore we can end up with this implicit cast, and have to move it into a post-projection
+            if (rewritten instanceof Cast) {
+                rewritten = ((Cast) rewritten).getExpression();
+                needPostProjectionCoercion = true;
+            }
+            aggregationAssignments.put(newSymbol, (FunctionCall) rewritten);
+            translations.put(aggregate, newSymbol);
+
+            functions.put(newSymbol, analysis.getFunctionSignature(aggregate));
+        }
 
         // 2.d. 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
@@ -437,6 +452,7 @@ class QueryPlanner
             masks.put(aggregateSymbol, marker);
         }
 
+        List<Symbol> distinctGroupingSymbols = distinctGroupingSymbolsBuilder.build().asList();
         for (Map.Entry<Set<Expression>, Symbol> entry : argumentMarkers.entrySet()) {
             ImmutableList.Builder<Symbol> builder = ImmutableList.builder();
             builder.addAll(distinctGroupingSymbols);
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 d93456e715..367d4deb5c 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,6 +113,7 @@ public final class SymbolExtractor
             node.getSource().accept(this, context);
 
             builder.add(node.getGroupIdSymbol());
+            builder.addAll(node.getIdentityMappings().values());
 
             return null;
         }
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 f1fb28c362..dcf584f27e 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
@@ -405,12 +405,21 @@ public class PruneUnreferencedOutputs
         {
             checkState(node.getDistinctGroupingColumns().stream().allMatch(column -> context.get().contains(column)));
 
-            PlanNode source = context.rewrite(node.getSource(), ImmutableSet.copyOf(context.get()));
-            List<Symbol> requiredSymbols = context.get().stream()
-                    .filter(symbol -> !symbol.equals(node.getGroupIdSymbol()))
-                    .collect(toImmutableList());
+            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> identityMapping = identityMappingBuilder.build();
+
+            PlanNode source = context.rewrite(node.getSource(), ImmutableSet.<Symbol>builder()
+                    .addAll(identityMapping.keySet())
+                    .addAll(node.getDistinctGroupingColumns())
+                    .build());
 
-            return new GroupIdNode(node.getId(), source, requiredSymbols, node.getGroupingSets(), node.getGroupIdSymbol());
+            return new GroupIdNode(node.getId(), source, node.getGroupingSets(), identityMapping, node.getGroupIdSymbol());
         }
 
         @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 b61ef06f94..bde6ac893d 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
@@ -162,7 +162,12 @@ public class UnaliasSymbolReferences
                     .map(this::canonicalize)
                     .collect(Collectors.toList());
 
-            return new GroupIdNode(node.getId(), source, canonicalize(node.getInputSymbols()), groupingSetsSymbols, canonicalize(node.getGroupIdSymbol()));
+            ImmutableMap.Builder<Symbol, Symbol> newPassthroughMap = ImmutableMap.builder();
+            for (Symbol inputSymbol : node.getIdentityMappings().keySet()) {
+                newPassthroughMap.put(canonicalize(inputSymbol), canonicalize(node.getIdentityMappings().get(inputSymbol)));
+            }
+
+            return new GroupIdNode(node.getId(), source, groupingSetsSymbols, newPassthroughMap.build(), 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 40c98bebf0..960396df1b 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
@@ -174,7 +174,7 @@ public class ChildReplacer
     @Override
     public PlanNode visitGroupId(GroupIdNode node, List<PlanNode> newChildren)
     {
-        return new GroupIdNode(node.getId(), Iterables.getOnlyElement(newChildren), node.getInputSymbols(), node.getGroupingSets(), node.getGroupIdSymbol());
+        return new GroupIdNode(node.getId(), Iterables.getOnlyElement(newChildren), node.getGroupingSets(), node.getIdentityMappings(), 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 2b3c7b7616..daf3560ced 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
@@ -17,12 +17,15 @@ import com.facebook.presto.sql.planner.Symbol;
 import com.fasterxml.jackson.annotation.JsonCreator;
 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 javax.annotation.concurrent.Immutable;
 
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import static java.util.Objects.requireNonNull;
@@ -33,21 +36,21 @@ public class GroupIdNode
         extends PlanNode
 {
     private final PlanNode source;
-    private final List<Symbol> inputSymbols;
     private final List<List<Symbol>> groupingSets;
+    private final Map<Symbol, Symbol> identityMappings;
     private final Symbol groupIdSymbol;
 
     @JsonCreator
     public GroupIdNode(@JsonProperty("id") PlanNodeId id,
             @JsonProperty("source") PlanNode source,
-            @JsonProperty("inputSymbols") List<Symbol> inputSymbols,
             @JsonProperty("groupingSets") List<List<Symbol>> groupingSets,
+            @JsonProperty("identityMappings") Map<Symbol, Symbol> identityMappings,
             @JsonProperty("groupIdSymbol") Symbol groupIdSymbol)
     {
         super(id);
         this.source = requireNonNull(source);
-        this.inputSymbols = ImmutableList.copyOf(requireNonNull(inputSymbols));
         this.groupingSets = ImmutableList.copyOf(requireNonNull(groupingSets));
+        this.identityMappings = ImmutableMap.copyOf(requireNonNull(identityMappings));
         this.groupIdSymbol = requireNonNull(groupIdSymbol);
     }
 
@@ -55,7 +58,8 @@ public class GroupIdNode
     public List<Symbol> getOutputSymbols()
     {
         return ImmutableList.<Symbol>builder()
-                .addAll(source.getOutputSymbols())
+                .addAll(getDistinctGroupingColumns())
+                .addAll(identityMappings.values())
                 .add(groupIdSymbol)
                 .build();
     }
@@ -72,10 +76,12 @@ public class GroupIdNode
         return source;
     }
 
-    @JsonProperty
-    public List<Symbol> getInputSymbols()
+    public Set<Symbol> getInputSymbols()
     {
-        return inputSymbols;
+        return ImmutableSet.<Symbol>builder()
+                .addAll(identityMappings.keySet())
+                .addAll(getDistinctGroupingColumns())
+                .build();
     }
 
     @JsonProperty
@@ -84,6 +90,12 @@ public class GroupIdNode
         return groupingSets;
     }
 
+    @JsonProperty
+    public Map<Symbol, Symbol> getIdentityMappings()
+    {
+        return identityMappings;
+    }
+
     public List<Symbol> getDistinctGroupingColumns()
     {
         return groupingSets.stream()
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 125bb02472..738a825da6 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
@@ -71,23 +71,28 @@ public class TestGroupIdOperator
                 .build();
 
         GroupIdOperatorFactory operatorFactory =
-                new GroupIdOperatorFactory(0, new PlanNodeId("test"), ImmutableList.of(BIGINT, VARCHAR, BOOLEAN, BIGINT), ImmutableList.of(ImmutableList.of(1, 2), ImmutableList.of(3)));
+                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));
 
         Operator operator = operatorFactory.createOperator(driverContext);
 
-        MaterializedResult expected = resultBuilder(driverContext.getSession(), BIGINT, VARCHAR, BOOLEAN, BIGINT, BIGINT)
-                .row(100L, "400", true, null, 0L)
-                .row(101L, "401", false, null, 0L)
-                .row(102L, "402", true, null, 0L)
-                .row(200L, "500", true, null, 0L)
-                .row(201L, "501", false, null, 0L)
-                .row(202L, "502", true, null, 0L)
-                .row(100L, null, null, 1000L, 1L)
-                .row(101L, null, null, 1001L, 1L)
-                .row(102L, null, null, 1002L, 1L)
-                .row(200L, null, null, 1100L, 1L)
-                .row(201L, null, null, 1101L, 1L)
-                .row(202L, null, null, 1102L, 1L)
+        MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR, BOOLEAN, BIGINT, BIGINT, BIGINT)
+                .row("400", true, null, 100L, 0L)
+                .row("401", false, null, 101L, 0L)
+                .row("402", true, null, 102L, 0L)
+                .row("500", true, null, 200L, 0L)
+                .row("501", false, null, 201L, 0L)
+                .row("502", true, null, 202L, 0L)
+                .row(null, null, 1000L, 100L, 1L)
+                .row(null, null, 1001L, 101L, 1L)
+                .row(null, null, 1002L, 102L, 1L)
+                .row(null, null, 1100L, 200L, 1L)
+                .row(null, null, 1101L, 201L, 1L)
+                .row(null, null, 1102L, 202L, 1L)
                 .build();
 
         List<Page> pages = toPages(operator, input.iterator());
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 1d43de6a41..bbad900d5d 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
@@ -1420,6 +1420,33 @@ public abstract class AbstractTestQueries
                 "SELECT orderpriority, -1 FROM orders GROUP BY orderpriority");
     }
 
+    @Test
+    public void testGroupingSetsAggregateOnGroupedColumn()
+            throws Exception
+    {
+        assertQuery("SELECT orderpriority, COUNT(orderpriority) FROM orders GROUP BY ROLLUP (orderpriority)",
+                "SELECT orderpriority, COUNT(orderpriority) FROM orders GROUP BY orderpriority UNION " +
+                        "SELECT NULL, COUNT(orderpriority) FROM orders");
+    }
+
+    @Test
+    public void testGroupingSetsMultipleAggregatesOnGroupedColumn()
+            throws Exception
+    {
+        assertQuery("SELECT linenumber, suppkey, SUM(suppkey), COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY GROUPING SETS ((linenumber, suppkey), ())",
+                "SELECT linenumber, suppkey, SUM(suppkey), COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY linenumber, suppkey UNION " +
+                        "SELECT NULL, NULL, SUM(suppkey), COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem");
+    }
+
+    @Test
+    public void testGroupingSetsMultipleAggregatesWithGroupedColumns()
+            throws Exception
+    {
+        assertQuery("SELECT linenumber, suppkey, COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY GROUPING SETS ((linenumber, suppkey), ())",
+                "SELECT linenumber, suppkey, COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem GROUP BY linenumber, suppkey UNION " +
+                        "SELECT NULL, NULL, COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem");
+    }
+
     @Test
     public void testRollup()
             throws Exception