Repository: prestodb/presto
Issue: Distinct sum aggregation may produce incorrect results for empty grouping sets

`AbstractTestQueries::testGroupingSetsWithMultipleDistinctAndUnion` fails for some configurations with `task_concurrency` higher then 1:

```
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), ())
not equal
Actual 8 rows:
    [3, 1, 5050]
    [null, 7, 30300]
    [7, 1, 5050]
    [1, 1, 5050]
    [4, 1, 5050]
    [6, 1, 5050]
    [5, 1, 5050]
    [2, 1, 5050]
Expected 8 rows:
    [1, 1, 5050]
    [2, 1, 5050]
    [3, 1, 5050]
    [4, 1, 5050]
    [5, 1, 5050]
    [6, 1, 5050]
    [7, 1, 5050]
    [null, 7, 5050]
```

Simplified query:

```
> set session task_concurrency=8;
> SELECT a, SUM(DISTINCT b) FROM (SELECT * FROM (VALUES (0,1), (0,2)) t(a,b) UNION ALL SELECT * FROM (VALUES (1,1), (1,2)) t(a,b)) GROUP BY GROUPING SETS ((a), ());
  a   | _col1 
------+-------
    0 |     3 
 NULL |     6 
    1 |     3 
```

To reproduce this, number of workers must be greater then 1, task_concurrency must be greater then 1, query must contain an empty grouping set AND there must be union AND aggregation function must be sum distinct. It seems like multiple partial aggregations of empty grouping set are incorrectly combined in final aggregation, effectively multiplying  result.

CC @kokosing 


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/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/PropertyDerivations.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/optimizations/PropertyDerivations.java
index 474dd52bbc..0d8aa9d9a5 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
@@ -222,9 +222,7 @@ class PropertyDerivations
         @Override
         public ActualProperties visitGroupId(GroupIdNode node, List<ActualProperties> inputProperties)
         {
-            ActualProperties properties = Iterables.getOnlyElement(inputProperties);
-
-            return properties.translate(translateGroupIdSymbols(node));
+            return Iterables.getOnlyElement(inputProperties).translate(translateGroupIdSymbols(node));
         }
 
         private Function<Symbol, Optional<Symbol>> translateGroupIdSymbols(GroupIdNode node)
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 867ed6a2af..cdcc0f0231 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
@@ -291,6 +291,28 @@ final class StreamPropertyDerivations
             return inputToOutput;
         }
 
+        @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));
+                }
+
+                if (commonGroupingColumns.contains(symbol)) {
+                    return Optional.of(symbol);
+                }
+
+                return Optional.empty();
+            };
+        }
+
         @Override
         public StreamProperties visitAggregation(AggregationNode node, List<StreamProperties> inputProperties)
         {
@@ -392,12 +414,6 @@ final class StreamPropertyDerivations
             return Iterables.getOnlyElement(inputProperties);
         }
 
-        @Override
-        public StreamProperties visitGroupId(GroupIdNode node, List<StreamProperties> inputProperties)
-        {
-            return Iterables.getOnlyElement(inputProperties);
-        }
-
         @Override
         public StreamProperties visitWindow(WindowNode node, List<StreamProperties> inputProperties)
         {
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 7fe4713525..afd7562503 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
@@ -1590,6 +1590,17 @@ public abstract class AbstractTestQueries
     @Test
     public void testGroupingSetsWithSingleDistinctAndUnion()
             throws Exception
+    {
+        assertQuery("SELECT suppkey, COUNT(DISTINCT linenumber) FROM " +
+                        "(SELECT * FROM lineitem WHERE linenumber%2 = 0 UNION ALL SELECT * FROM lineitem WHERE linenumber%2 = 1) " +
+                        "GROUP BY GROUPING SETS ((suppkey), ())",
+                "SELECT suppkey, COUNT(DISTINCT linenumber) FROM lineitem GROUP BY suppkey UNION ALL " +
+                        "SELECT NULL, COUNT(DISTINCT linenumber) FROM lineitem");
+    }
+
+    @Test
+    public void testGroupingSetsWithSingleDistinctAndUnionGroupedArguments()
+            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) " +