From 06fa8c8fb878ae2f0de3bc5b70b4fe870f21f2d8 Mon Sep 17 00:00:00 2001
From: Mihai Budiu <mbudiu@feldera.com>
Date: Mon, 8 Apr 2024 15:04:32 -0700
Subject: [PATCH] [CALCITE-6332] Optimization
 CoreRules.AGGREGATE_EXPAND_DISTINCT_AGGREGATES_TO_JOIN produces incorrect
 results for aggregates with groupSets

Signed-off-by: Mihai Budiu <mbudiu@feldera.com>
---
 .../AggregateExpandDistinctAggregatesRule.java    |  7 +++++++
 .../org/apache/calcite/test/RelOptRulesTest.java  | 12 ++++++++++++
 .../org/apache/calcite/test/RelOptRulesTest.xml   | 15 +++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/core/src/main/java/org/apache/calcite/rel/rules/AggregateExpandDistinctAggregatesRule.java b/core/src/main/java/org/apache/calcite/rel/rules/AggregateExpandDistinctAggregatesRule.java
index 7a00db8d72cf..ccf32284e74f 100644
--- a/core/src/main/java/org/apache/calcite/rel/rules/AggregateExpandDistinctAggregatesRule.java
+++ b/core/src/main/java/org/apache/calcite/rel/rules/AggregateExpandDistinctAggregatesRule.java
@@ -133,6 +133,13 @@ public AggregateExpandDistinctAggregatesRule(
       return;
     }
 
+    if (!config.isUsingGroupingSets()
+        && aggregate.groupSets.size() > 1) {
+      // Grouping sets are not handled correctly
+      // when generating joins.
+      return;
+    }
+
     // Find all of the agg expressions. We use a LinkedHashSet to ensure determinism.
     final List<AggregateCall> aggCalls = aggregate.getAggCallList();
     // Find all aggregate calls with distinct
diff --git a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
index 51a13e0bc50a..4430186326b2 100644
--- a/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java
@@ -1990,6 +1990,18 @@ private void checkSemiOrAntiJoinProjectTranspose(JoinRelType type) {
         .check();
   }
 
+  /** Test case for <a href="https://issues.apache.org/jira/projects/CALCITE/issues/CALCITE-6332">
+   * [CALCITE-6332] Optimization CoreRules.AGGREGATE_EXPAND_DISTINCT_AGGREGATES_TO_JOIN
+   * produces incorrect results for aggregates with groupSets</a>. */
+  @Test void testIssue6332() {
+    final String sql = "select count(distinct deptno) as cd, count(*) as c\n"
+        + "from emp\n"
+        + "group by cube(deptno)";
+    sql(sql)
+        .withRule(CoreRules.AGGREGATE_EXPAND_DISTINCT_AGGREGATES_TO_JOIN)
+        .checkUnchanged();
+  }
+
   @Test void testDistinctCountMixed() {
     final String sql = "select deptno, count(distinct deptno, job) as cddj,\n"
         + "  sum(sal) as s\n"
diff --git a/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml b/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
index bc940b804c8f..672c80398e5c 100644
--- a/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
+++ b/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
@@ -5163,6 +5163,21 @@ LogicalProject(EMPNO=[$0], DEPTNO=[$1], W_COUNT=[$2])
     LogicalProject(EMPNO=[$0], DEPTNO=[$7], $2=[$9])
       LogicalWindow(window#0=[window(aggs [COUNT($0)])])
         LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testIssue6332">
+    <Resource name="sql">
+      <![CDATA[select count(distinct deptno) as cd, count(*) as c
+from emp
+group by cube(deptno)]]>
+    </Resource>
+    <Resource name="planBefore">
+      <![CDATA[
+LogicalProject(CD=[$1], C=[$2])
+  LogicalAggregate(group=[{0}], groups=[[{0}, {}]], CD=[COUNT(DISTINCT $0)], C=[COUNT()])
+    LogicalProject(DEPTNO=[$7])
+      LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>