Repository: prestodb/presto
Issue: GROUPING SETS over UNION fails

The following query:

```
SELECT x, COUNT(DISTINCT x) AS c
   FROM
     (
     SELECT *
     FROM (values (1)) t(x)
     UNION ALL
     SELECT *
     FROM (values (1)) t(x)
     ) t2
   GROUP BY GROUPING SETS ((x), ())
```

fails with:

```
Query 20160622_013115_00003_u5amn failed: index (0) must be less than size (0)
java.lang.IndexOutOfBoundsException: index (0) must be less than size (0)
    at com.google.common.base.Preconditions.checkElementIndex(Preconditions.java:310)
    at com.google.common.base.Preconditions.checkElementIndex(Preconditions.java:292)
    at com.google.common.collect.RegularImmutableList.get(RegularImmutableList.java:65)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitUnion(AddExchanges.java:1104)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitUnion(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.UnionNode.accept(UnionNode.java:139)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1143)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitPlan(AddExchanges.java:215)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitPlan(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.PlanVisitor.visitGroupId(PlanVisitor.java:150)
    at com.facebook.presto.sql.planner.plan.GroupIdNode.accept(GroupIdNode.java:125)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1143)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitAggregation(AddExchanges.java:283)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitAggregation(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.AggregationNode.accept(AggregationNode.java:189)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1143)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitProject(AddExchanges.java:231)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitProject(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.ProjectNode.accept(ProjectNode.java:76)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1143)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitAggregation(AddExchanges.java:283)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitAggregation(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.AggregationNode.accept(AggregationNode.java:189)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1143)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitProject(AddExchanges.java:231)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitProject(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.ProjectNode.accept(ProjectNode.java:76)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.planChild(AddExchanges.java:1143)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitOutput(AddExchanges.java:237)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges$Rewriter.visitOutput(AddExchanges.java:188)
    at com.facebook.presto.sql.planner.plan.OutputNode.accept(OutputNode.java:81)
    at com.facebook.presto.sql.planner.optimizations.AddExchanges.optimize(AddExchanges.java:150)
    at com.facebook.presto.sql.planner.LogicalPlanner.plan(LogicalPlanner.java:100)
    at com.facebook.presto.execution.SqlQueryExecution.doAnalyzeQuery(SqlQueryExecution.java:277)
    at com.facebook.presto.execution.SqlQueryExecution.analyzeQuery(SqlQueryExecution.java:256)
    at com.facebook.presto.execution.SqlQueryExecution.start(SqlQueryExecution.java:220)
    at com.facebook.presto.execution.QueuedExecution.lambda$start$1(QueuedExecution.java:68)
    at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1142)
    at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:617)
    at java.lang.Thread.run(Thread.java:745)
```


Changed files:
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/AddExchanges.java
presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java

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 f9c62657d0..6f4837b207 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
@@ -45,6 +45,7 @@ import com.facebook.presto.sql.planner.plan.EnforceSingleRowNode;
 import com.facebook.presto.sql.planner.plan.ExchangeNode;
 import com.facebook.presto.sql.planner.plan.ExplainAnalyzeNode;
 import com.facebook.presto.sql.planner.plan.FilterNode;
+import com.facebook.presto.sql.planner.plan.GroupIdNode;
 import com.facebook.presto.sql.planner.plan.IndexJoinNode;
 import com.facebook.presto.sql.planner.plan.IndexSourceNode;
 import com.facebook.presto.sql.planner.plan.JoinNode;
@@ -383,6 +384,31 @@ public class AddExchanges
                     deriveProperties(source, partial.getProperties()));
         }
 
+        @Override
+        public PlanWithProperties visitGroupId(GroupIdNode node, Context context)
+        {
+            PreferredProperties childPreference = context.getPreferredProperties().translate(translateGroupIdSymbols(node));
+            PlanWithProperties child = planChild(node, context.withPreferredProperties(childPreference));
+            return rebaseAndDeriveProperties(node, child);
+        }
+
+        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 (commonGroupingColumns.contains(symbol)) {
+                    return Optional.of(symbol);
+                }
+
+                return Optional.empty();
+            };
+        }
+
         @Override
         public PlanWithProperties visitMarkDistinct(MarkDistinctNode node, Context context)
         {
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 b586aea570..d9c01d0f7a 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
@@ -75,6 +75,7 @@ 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;
@@ -213,7 +214,25 @@ class PropertyDerivations
         @Override
         public ActualProperties visitGroupId(GroupIdNode node, List<ActualProperties> inputProperties)
         {
-            return Iterables.getOnlyElement(inputProperties);
+            ActualProperties properties = Iterables.getOnlyElement(inputProperties);
+
+            return properties.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));
+                }
+
+                if (commonGroupingColumns.contains(symbol)) {
+                    return Optional.of(symbol);
+                }
+
+                return Optional.empty();
+            };
         }
 
         @Override
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 bda81a8b9f..f8225e2e2d 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
@@ -1481,6 +1481,28 @@ public abstract class AbstractTestQueries
                         "SELECT NULL, NULL, COUNT(linenumber), SUM(CAST(quantity AS BIGINT)) FROM lineitem");
     }
 
+    @Test
+    public void testGroupingSetsWithSingleDistinctAndUnion()
+            throws Exception
+    {
+        assertQuery("SELECT linenumber, COUNT(DISTINCT linenumber) FROM " +
+                        "(SELECT * FROM lineitem WHERE linenumber%2 = 0 UNION ALL SELECT * FROM lineitem WHERE linenumber%2 = 1) " +
+                        "GROUP BY GROUPING SETS ((linenumber), ())",
+                "SELECT DISTINCT linenumber, 1 FROM lineitem UNION ALL " +
+                        "SELECT NULL, COUNT(DISTINCT linenumber) FROM lineitem");
+    }
+
+    @Test
+    public void testGroupingSetsWithMultipleDistinctAndUnion()
+            throws Exception
+    {
+        assertQuery("SELECT linenumber, COUNT(DISTINCT linenumber), SUM(DISTINCT suppkey) FROM " +
+                        "(SELECT * FROM lineitem WHERE linenumber%2 = 0 UNION ALL SELECT * FROM lineitem WHERE linenumber%2 = 1) " +
+                        "GROUP BY GROUPING SETS ((linenumber), ())",
+                "SELECT linenumber, 1, SUM(DISTINCT suppkey) FROM lineitem GROUP BY linenumber UNION ALL " +
+                        "SELECT NULL, COUNT(DISTINCT linenumber), SUM(DISTINCT suppkey) FROM lineitem");
+    }
+
     @Test
     public void testRollup()
             throws Exception