[CALCITE-3957] AggregateMergeRule should merge SUM0 into COUNT even i…

…f GROUP BY is empty

Close apache#2097

core/src/main/java/org/apache/calcite/rel/rules/AggregateMergeRule.java

@@ -22,6 +22,7 @@
 import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.Aggregate.Group;
 import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlSplittableAggFunction;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.tools.RelBuilderFactory;
@@ -99,7 +100,7 @@ private boolean isAggregateSupported(AggregateCall aggCall) {
     }
 
     boolean hasEmptyGroup = topAgg.getGroupSets()
-        .stream().anyMatch(n -> n.isEmpty());
+        .stream().anyMatch(ImmutableBitSet::isEmpty);
 
     final List<AggregateCall> finalCalls = new ArrayList<>();
     for (AggregateCall topCall : topAgg.getAggCallList()) {
@@ -120,6 +121,7 @@ private boolean isAggregateSupported(AggregateCall aggCall) {
       // 0, which is wrong.
       if (!isAggregateSupported(bottomCall)
           || (bottomCall.getAggregation() == SqlStdOperatorTable.COUNT
+               && topCall.getAggregation().getKind() != SqlKind.SUM0
                && hasEmptyGroup)) {
         return;
       }

core/src/main/java/org/apache/calcite/rel/rules/ProjectAggregateMergeRule.java

@@ -47,7 +47,9 @@
  * Planner rule that matches a {@link Project} on a {@link Aggregate}
  * and projects away aggregate calls that are not used.
  *
- * <p>Also converts {@code NVL(SUM(x), 0)} to {@code SUM0(x)}.
+ * <p>Also converts {@code COALESCE(SUM(x), 0)} to {@code SUM0(x)}.
+ * This transformation is useful because there are cases where
+ * {@link AggregateMergeRule} can merge {@code SUM0} but not {@code SUM}.
  *
  * @see CoreRules#PROJECT_AGGREGATE_MERGE
  */

core/src/main/java/org/apache/calcite/sql/SqlSplittableAggFunction.java

@@ -198,7 +198,8 @@ public RexNode singleton(RexBuilder rexBuilder, RelDataType inputRowType,
 
     public AggregateCall merge(AggregateCall top, AggregateCall bottom) {
       if (bottom.getAggregation().getKind() == SqlKind.COUNT
-          && top.getAggregation().getKind() == SqlKind.SUM) {
+          && (top.getAggregation().getKind() == SqlKind.SUM
+              || top.getAggregation().getKind() == SqlKind.SUM0)) {
         return AggregateCall.create(bottom.getAggregation(),
             bottom.isDistinct(), bottom.isApproximate(), false,
             bottom.getArgList(), bottom.filterArg, bottom.getCollation(),

core/src/test/java/org/apache/calcite/test/RelOptRulesTest.java

@@ -1825,7 +1825,7 @@ private Sql checkPushProjectPastFilter3(ProjectFilterTransposeRule rule) {
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-2343">[CALCITE-2343]
    * Should not push over whose columns are all from left child past join since
-   * join will affect row count.</a>. */
+   * join will affect row count</a>. */
   @Test void testPushProjectWithOverPastJoin1() {
     final String sql = "select e.sal + b.comm,\n"
         + "count(e.empno) over (partition by e.deptno)\n"
@@ -4731,11 +4731,11 @@ private HepProgram getTransitiveProgram() {
 
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-2249">[CALCITE-2249]
-   * AggregateJoinTransposeRule generates inequivalent nodes if Aggregate relNode contains
-   * distinct aggregate function.</a>. */
+   * AggregateJoinTransposeRule generates non-equivalent nodes if Aggregate
+   * contains DISTINCT aggregate function</a>. */
   @Test void testPushDistinctAggregateIntoJoin() {
-    final String sql =
-            "select count(distinct sal) from sales.emp join sales.dept on job = name";
+    final String sql = "select count(distinct sal) from sales.emp\n"
+        + " join sales.dept on job = name";
     sql(sql).withRule(CoreRules.AGGREGATE_JOIN_TRANSPOSE_EXTENDED)
         .checkUnchanged();
   }
@@ -4893,6 +4893,24 @@ private HepProgram getTransitiveProgram() {
         .checkUnchanged();
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-3957">[CALCITE-3957]
+   * AggregateMergeRule should merge SUM0 into COUNT even if GROUP BY is
+   * empty</a>. (It is not valid to merge a SUM onto a SUM0 if the top GROUP BY
+   * is empty.) */
+  @Test void testAggregateMergeSum0() {
+    final String sql = "select coalesce(sum(count_comm), 0)\n"
+        + "from (\n"
+        + "  select deptno, count(comm) as count_comm\n"
+        + "  from sales.emp\n"
+        + "  group by deptno, mgr) t";
+    sql(sql)
+        .withPreRule(CoreRules.PROJECT_AGGREGATE_MERGE,
+            CoreRules.AGGREGATE_PROJECT_MERGE)
+        .withRule(CoreRules.AGGREGATE_MERGE)
+        .check();
+  }
+
   /**
    * Test case for AggregateMergeRule, should not merge 2 aggregates
    * into a single aggregate, since top agg contains empty grouping set,
@@ -5027,7 +5045,7 @@ private HepProgram getTransitiveProgram() {
   /** Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-2712">[CALCITE-2712]
    * Should remove the left join since the aggregate has no call and
-   * only uses column in the left input of the bottom join as group key.</a>. */
+   * only uses column in the left input of the bottom join as group key</a>. */
   @Test void testAggregateJoinRemove1() {
     final String sql = "select distinct e.deptno from sales.emp e\n"
         + "left outer join sales.dept d on e.deptno = d.deptno";

core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml

@@ -542,6 +542,28 @@ LogicalProject(EXPR$0=[1])
       LogicalTableScan(table=[[CATALOG, SALES, EMP]])
       LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
     LogicalTableScan(table=[[CATALOG, SALES, DEPT]])
+]]>
+        </Resource>
+    </TestCase>
+    <TestCase name="testAggregateMergeSum0">
+        <Resource name="sql">
+            <![CDATA[select coalesce(sum(count_comm), 0)
+from (
+  select deptno, count(comm) as count_comm
+  from sales.emp
+  group by deptno, mgr) t]]>
+        </Resource>
+        <Resource name="planBefore">
+            <![CDATA[
+LogicalAggregate(group=[{}], agg#0=[$SUM0($2)])
+  LogicalAggregate(group=[{3, 7}], COUNT_COMM=[COUNT()])
+    LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+        </Resource>
+        <Resource name="planAfter">
+            <![CDATA[
+LogicalAggregate(group=[{}], agg#0=[COUNT()])
+  LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
         </Resource>
     </TestCase>