[CALCITE-6214] Add primary key unit test

JiajunBernoulli · Feb 2, 2024 · 4671d96 · 4671d96
1 parent a4419bb
commit 4671d96
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 1 deletion.
diff --git a/core/src/main/java/org/apache/calcite/tools/RelBuilder.java b/core/src/main/java/org/apache/calcite/tools/RelBuilder.java
@@ -2431,7 +2431,9 @@ private RelBuilder aggregate_(GroupKeyImpl groupKey,
     final List<AggregateCall> aggregateCalls = new ArrayList<>();
     for (AggCallPlus aggCall : aggCalls) {
       AggregateCall aggregateCall = aggCall.aggregateCall(registrar, groupSet, r);
-      aggregateCall = removeRedundantDistinct(aggregateCall, groupSet, r);
+      if (groupSets.size() <= 1) {
+        aggregateCall = removeRedundantDistinct(aggregateCall, groupSet, r);
+      }
       aggregateCalls.add(aggregateCall);
     }
 

diff --git a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java
@@ -5103,6 +5103,50 @@ void checkUserDefinedOrderByOver(NullCollation nullCollation) {
     sql(sql).ok();
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6214">[CALCITE-6214]
+   * Remove DISTINCT in COUNT if field is unique</a>. */
+  @Test void testRemoveDistinctIfUnique5() {
+    // empno is unique key
+    final String sql = "SELECT COUNT(DISTINCT empno)\n"
+        + "FROM emp\n";
+    // Default save redundant distinct
+    sql(sql)
+        .withConfig(c ->
+            c.addRelBuilderConfigTransform(c2 ->
+                c2.withRedundantDistinct(false))).ok();
+  }
+
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6214">[CALCITE-6214]
+   * Remove DISTINCT in COUNT if field is unique</a>. */
+  @Test void testRemoveDistinctIfUnique6() {
+    // empno is unique key
+    final String sql = "SELECT deptno, COUNT(DISTINCT empno)\n"
+        + "FROM emp\n"
+        + "GROUP BY deptno";
+    // Default save redundant distinct
+    sql(sql)
+        .withConfig(c ->
+        c.addRelBuilderConfigTransform(c2 ->
+            c2.withRedundantDistinct(false))).ok();
+  }
+
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6214">[CALCITE-6214]
+   * Remove DISTINCT in COUNT if field is unique</a>. */
+  @Test void testRemoveDistinctIfUnique7() {
+    // empno is unique key
+    final String sql = "SELECT deptno, COUNT(DISTINCT empno)\n"
+        + "FROM emp\n"
+        + "GROUP BY ROLLUP(deptno)";
+    // Default save redundant distinct
+    sql(sql)
+        .withConfig(c ->
+            c.addRelBuilderConfigTransform(c2 ->
+                c2.withRedundantDistinct(false))).ok();
+  }
+
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-5089">[CALCITE-5089]
    * Allow GROUP BY ALL or DISTINCT set quantifier on GROUPING SETS</a>. */

diff --git a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
@@ -6456,6 +6456,48 @@ LogicalProject(CDS=[$1], CS=[$2], SDS=[$3], SS=[$4])
     LogicalAggregate(group=[{0, 1}])
       LogicalProject(DEPTNO=[$7], SAL=[$5])
         LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testRemoveDistinctIfUnique5">
+    <Resource name="sql">
+      <![CDATA[SELECT COUNT(DISTINCT empno)
+FROM emp
+]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalAggregate(group=[{}], EXPR$0=[COUNT($0)])
+  LogicalProject(EMPNO=[$0])
+    LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testRemoveDistinctIfUnique6">
+    <Resource name="sql">
+      <![CDATA[SELECT deptno, COUNT(DISTINCT empno)
+FROM emp
+GROUP BY deptno]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalAggregate(group=[{0}], EXPR$1=[COUNT($1)])
+  LogicalProject(DEPTNO=[$7], EMPNO=[$0])
+    LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testRemoveDistinctIfUnique7">
+    <Resource name="sql">
+      <![CDATA[SELECT deptno, COUNT(DISTINCT empno)
+FROM emp
+GROUP BY ROLLUP(deptno)]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalAggregate(group=[{0}], groups=[[{0}, {}]], EXPR$1=[COUNT(DISTINCT $1)])
+  LogicalProject(DEPTNO=[$7], EMPNO=[$0])
+    LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>