[calcite] branch main updated: [CALCITE-5162] RelMdUniqueKeys can return more precise unique keys for Aggregate

[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 f278efb04 [CALCITE-5162] RelMdUniqueKeys can return more precise unique keys for Aggregate
f278efb04 is described below

commit f278efb0411ba29c7bd167f5d02e566bf542acb3
Author: chunwei.lcw <ch...@gmail.com>
AuthorDate: Fri May 20 20:15:39 2022 +0800

    [CALCITE-5162] RelMdUniqueKeys can return more precise unique keys for Aggregate
---
 .../org/apache/calcite/rel/metadata/RelMdUniqueKeys.java  | 15 ++++++++++++++-
 .../java/org/apache/calcite/test/RelMetadataTest.java     |  8 ++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdUniqueKeys.java b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdUniqueKeys.java
index f8a845f7e..21bfdbed2 100644
--- a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdUniqueKeys.java
+++ b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdUniqueKeys.java
@@ -48,6 +48,7 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static java.util.Objects.requireNonNull;
 
@@ -247,7 +248,19 @@ public class RelMdUniqueKeys
 
   public Set<ImmutableBitSet> getUniqueKeys(Aggregate rel, RelMetadataQuery mq,
       boolean ignoreNulls) {
-    if (Aggregate.isSimple(rel) || ignoreNulls) {
+    if (Aggregate.isSimple(rel)) {
+      final ImmutableBitSet groupKeys = rel.getGroupSet();
+      final Set<ImmutableBitSet> inputUniqueKeys = mq
+          .getUniqueKeys(rel.getInput(), ignoreNulls);
+      if (inputUniqueKeys == null) {
+        return ImmutableSet.of(groupKeys);
+      }
+
+      // Try to find more precise unique keys.
+      final Set<ImmutableBitSet> preciseUniqueKeys = inputUniqueKeys.stream()
+          .filter(groupKeys::contains).collect(Collectors.toSet());
+      return preciseUniqueKeys.isEmpty() ? ImmutableSet.of(groupKeys) : preciseUniqueKeys;
+    } else if (ignoreNulls) {
       // group by keys form a unique key
       return ImmutableSet.of(rel.getGroupSet());
     } else {
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 cbca1576e..d7ec865c7 100644
--- a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
+++ b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -969,6 +969,14 @@ public class RelMetadataTest {
         .assertThatUniqueKeysAre(bitSetOf());
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-5162">[CALCITE-5162]
+   * RelMdUniqueKeys can return more precise unique keys for Aggregate</a>. */
+  @Test void testGroupByPreciseUniqueKeys() {
+    sql("select empno, ename from emp group by empno, ename")
+        .assertThatUniqueKeysAre(bitSetOf(0));
+  }
+
   @Test void testFullOuterJoinUniqueness1() {
     final String sql = "select e.empno, d.deptno\n"
         + "from (select cast(null as int) empno from sales.emp "