[calcite] branch main updated: [CALCITE-5149] Refine RelMdColumnUniqueness for Aggregate by considering intersect keys between target keys and group keys

[ch...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

chunwei pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/calcite.git


The following commit(s) were added to refs/heads/main by this push:
     new c5d45cb094 [CALCITE-5149] Refine RelMdColumnUniqueness for Aggregate by considering intersect keys between target keys and group keys
c5d45cb094 is described below

commit c5d45cb094407658f5f933240bc4492b6239e45e
Author: chunwei.lcw <ch...@gmail.com>
AuthorDate: Thu May 12 17:43:07 2022 +0800

    [CALCITE-5149] Refine RelMdColumnUniqueness for Aggregate by considering intersect keys between target keys and group keys
---
 .../calcite/rel/metadata/RelMdColumnUniqueness.java  | 20 ++++++++++++++++++--
 .../org/apache/calcite/test/RelMetadataTest.java     | 18 ++++++++++++++++++
 .../org/apache/calcite/test/RelOptRulesTest.xml      |  9 ++++-----
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdColumnUniqueness.java b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdColumnUniqueness.java
index 7e912f05c9..aef93c3a55 100644
--- a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdColumnUniqueness.java
+++ b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdColumnUniqueness.java
@@ -351,8 +351,24 @@ public class RelMdColumnUniqueness
     if (Aggregate.isSimple(rel) || ignoreNulls) {
       columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
       // group by keys form a unique key
-      ImmutableBitSet groupKey = ImmutableBitSet.range(rel.getGroupCount());
-      return columns.contains(groupKey);
+      final ImmutableBitSet groupKey = ImmutableBitSet.range(rel.getGroupCount());
+      final boolean contained = columns.contains(groupKey);
+      if (contained) {
+        return true;
+      } else if (!Aggregate.isSimple(rel)) {
+        return false;
+      }
+
+      final ImmutableBitSet commonKeys = columns.intersect(groupKey);
+      if (commonKeys.isEmpty()) {
+        return false;
+      }
+      final ImmutableBitSet.Builder targetColumns = ImmutableBitSet.builder();
+      for (int key: commonKeys) {
+        targetColumns.set(rel.getGroupSet().nth(key));
+      }
+
+      return mq.areColumnsUnique(rel.getInput(), targetColumns.build(), ignoreNulls);
     }
     return null;
   }
diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
index 247cf870cb..cbca1576e6 100644
--- a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -1116,6 +1116,24 @@ public class RelMetadataTest {
         .assertThatAreColumnsUnique(bitSetOf(0), is(true));
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-5149">[CALCITE-5149]
+   * Refine RelMdColumnUniqueness for Aggregate by considering intersect keys
+   * between target keys and group keys</a>. */
+  @Test void testColumnUniquenessForAggregate() {
+    sql("select empno, ename, count(1) as cnt from emp group by empno, ename")
+        .assertThatAreColumnsUnique(bitSetOf(0, 1), is(true));
+
+    sql("select empno, ename, count(1) as cnt from emp group by empno, ename")
+        .assertThatAreColumnsUnique(bitSetOf(0), is(true));
+
+    sql("select ename, empno, count(1) as cnt from emp group by ename, empno")
+        .assertThatAreColumnsUnique(bitSetOf(1), is(true));
+
+    sql("select empno, ename, count(1) as cnt from emp group by empno, ename")
+        .assertThatAreColumnsUnique(bitSetOf(2), is(false));
+  }
+
   @Test void testGroupBy() {
     sql("select deptno, count(*), sum(sal) from emp group by deptno")
         .assertThatUniqueKeysAre(bitSetOf(0));
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 c30ab69e9a..908829af8c 100644
--- a/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
+++ b/core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml
@@ -1835,11 +1835,10 @@ LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$
 LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
   LogicalJoin(condition=[AND(=($3, $10), =($5, $9))], joinType=[inner])
     LogicalTableScan(table=[[CATALOG, SALES, EMP]])
-    LogicalAggregate(group=[{0, 1}])
-      LogicalProject(EXPR$0=[$2], EMPNO=[$1])
-        LogicalAggregate(group=[{0, 1}], EXPR$0=[MAX($2)])
-          LogicalProject(DEPTNO=[$7], EMPNO=[$0], SAL=[$5])
-            LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+    LogicalProject(EXPR$0=[$2], EMPNO=[$1])
+      LogicalAggregate(group=[{0, 1}], EXPR$0=[MAX($2)])
+        LogicalProject(DEPTNO=[$7], EMPNO=[$0], SAL=[$5])
+          LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>