Prevent removal of aggregation nodes when grouping sets are present

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

@@ -45,7 +45,8 @@ public class RemoveRedundantAggregateDistinct
     private static boolean hasAggregations(AggregationNode node)
     {
         return ((GroupReference) node.getSource()).getLogicalProperties().isPresent() &&
-                !node.getAggregations().isEmpty();
+                !node.getAggregations().isEmpty() &&
+                node.getGroupingSetCount() == 1;
     }
 
     @Override

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

@@ -35,7 +35,7 @@ public class RemoveRedundantDistinct
 
     private static boolean distinctOfUniqueKey(AggregationNode node)
     {
-        return node.hasNonEmptyGroupingSet() &&
+        return node.getGroupingSetCount() == 1 &&
                 node.getAggregations().isEmpty() &&
                 ((GroupReference) node.getSource()).getLogicalProperties().isPresent() &&
                 ((GroupReference) node.getSource()).getLogicalProperties().get().isDistinct(node.getGroupingKeys().stream().collect(Collectors.toSet()));

presto-main/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java

@@ -22,6 +22,7 @@
 import com.facebook.presto.spi.plan.FilterNode;
 import com.facebook.presto.spi.plan.LimitNode;
 import com.facebook.presto.spi.plan.PlanNode;
+import com.facebook.presto.spi.plan.ProjectNode;
 import com.facebook.presto.spi.plan.TableScanNode;
 import com.facebook.presto.spi.plan.TopNNode;
 import com.facebook.presto.spi.plan.ValuesNode;
@@ -1637,6 +1638,33 @@ public void testRedundantDistinctLimitNodeRemoval()
                                         tableScan("orders")))));
     }
 
+    @Test
+    public void testRedundantDistinctRemovalWithMultipleGroupingSets()
+    {
+        Session exploitConstraints = Session.builder(this.getQueryRunner().getDefaultSession())
+                .setSystemProperty(EXPLOIT_CONSTRAINTS, Boolean.toString(true))
+                .build();
+
+        assertPlan("SELECT orderkey FROM (SELECT * FROM (SELECT * FROM orders WHERE 1=0)) GROUP BY ROLLUP (orderkey)",
+                exploitConstraints,
+                anyTree(
+                        node(AggregationNode.class,
+                                node(ProjectNode.class,
+                                        node(ValuesNode.class)))));
+
+        assertPlan("with t as (select orderkey, count(1) cnt from (select * from (select * from orders where 1=0) left join (select partkey, suppkey from lineitem where 1=0) on partkey=10 where suppkey is not null) group by rollup(orderkey)) select t1.orderkey, t1.cnt from t t1 cross join t t2",
+                exploitConstraints,
+                anyTree(
+                        node(JoinNode.class,
+                                anyTree(
+                                        node(AggregationNode.class,
+                                                node(ProjectNode.class,
+                                                        node(ValuesNode.class)))),
+                                anyTree(node(AggregationNode.class,
+                                        node(ProjectNode.class,
+                                                node(ValuesNode.class)))))));
+    }
+
     @Test
     public void testLeafNodeInPlanExceedException()
     {

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

@@ -131,6 +131,17 @@ public void singleTableTests()
                         node(ProjectNode.class,
                                 aggregation(aggregations,
                                         tableScan("lineitem")))));
+
+        aggregations = ImmutableMap.of(
+                "count", functionCall("count", true, ImmutableList.of(anySymbol())),
+                "avg", functionCall("avg", true, ImmutableList.of(anySymbol())));
+
+        // Multiple grouping sets, so distinct's cannot be removed, even for unique groups
+        tester().assertThat(ImmutableSet.of(new RemoveRedundantAggregateDistinct()), logicalPropertiesProvider)
+                .on("SELECT count(distinct linenumber), avg(distinct tax) FROM lineitem GROUP BY GROUPING SETS ((orderkey), (partkey))")
+                .matches(anyTree(
+                        aggregation(aggregations,
+                                anyTree(tableScan("lineitem")))));
     }
 
     @Test

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

@@ -22,13 +22,15 @@
 import com.facebook.presto.sql.planner.iterative.rule.test.BaseRuleTest;
 import com.facebook.presto.sql.planner.iterative.rule.test.RuleTester;
 import com.facebook.presto.sql.relational.FunctionResolution;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 import java.util.Optional;
 
+import static com.facebook.presto.spi.plan.AggregationNode.groupingSets;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.any;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.assignUniqueId;
 import static com.facebook.presto.sql.planner.assertions.PlanMatchPattern.node;
@@ -156,6 +158,17 @@ public void doesNotFire()
                                     p.values(5, c))));
                 })
                 .doesNotFire();
+
+        tester().assertThat(new RemoveRedundantDistinct(), logicalPropertiesProvider)
+                .on(p -> {
+                    VariableReferenceExpression c = p.variable("c");
+                    return p.aggregation(builder -> builder
+                            .groupingSets(groupingSets(ImmutableList.of(c), 2, ImmutableSet.of(0)))
+                            .source(p.limit(
+                                    1, // Forces MaxCard to be 1, so a unique key does exist on 'c' (one of the grouping keys)
+                                    p.values(5, c))));
+                })
+                .doesNotFire(); // Since multiple grouping sets are used however, the rule does not fire
     }
 
     @Test

presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java

@@ -7511,5 +7511,7 @@ public void testGuardConstraintFramework()
     {
         assertQuery("with t as (select orderkey, count(1) cnt from (select * from (select * from orders where 1=0) left join (select partkey, suppkey from lineitem where 1=0) on partkey=10 where suppkey is not null) group by rollup(orderkey)) select t1.orderkey, t1.cnt from t t1 cross join t t2",
                 "values (null, 0)");
+        assertQuery("select orderkey from (select * from (select * from orders where 1=0)) group by rollup(orderkey)",
+                "values (null)");
     }
 }