[CALCITE-6640] RelMdUniqueKeys grows exponentially when key columns are repeated in projections

core/src/main/java/org/apache/calcite/rel/metadata/RelMdUniqueKeys.java

@@ -46,15 +46,14 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Maps;
+import com.google.common.collect.Multimap;
 import com.google.common.collect.Sets;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -63,8 +62,6 @@
 import static org.apache.calcite.rel.metadata.RelMdColumnUniqueness.PASSTHROUGH_AGGREGATIONS;
 import static org.apache.calcite.rel.metadata.RelMdColumnUniqueness.getConstantColumnSet;
 
-import static java.util.Objects.requireNonNull;
-
 /**
  * RelMdUniqueKeys supplies a default implementation of
  * {@link RelMetadataQuery#getUniqueKeys} for the standard logical algebra.
@@ -172,8 +169,7 @@ private static Set<ImmutableBitSet> getProjectUniqueKeys(SingleRel rel, RelMetad
       return ImmutableSet.of();
     }
 
-    Map<Integer, ImmutableBitSet> mapInToOutPos =
-        Maps.transformValues(inToOutPosBuilder.build().asMap(), ImmutableBitSet::of);
+    Multimap<Integer, Integer> mapInToOutPos = inToOutPosBuilder.build();
 
     ImmutableSet.Builder<ImmutableBitSet> resultBuilder = ImmutableSet.builder();
     // Now add to the projUniqueKeySet the child keys that are fully
@@ -184,19 +180,12 @@ private static Set<ImmutableBitSet> getProjectUniqueKeys(SingleRel rel, RelMetad
         continue;
       }
       // colMask is mapped to output project, however, the column can be mapped more than once:
-      // select id, id, id, unique2, unique2
-      // the resulting unique keys would be {{0},{3}}, {{0},{4}}, {{0},{1},{4}}, ...
-
-      Iterable<List<ImmutableBitSet>> product =
-          Linq4j.product(
-              Util.transform(colMask, in ->
-                  Util.filter(
-                      requireNonNull(mapInToOutPos.get(in),
-                          () -> "no entry for column " + in
-                              + " in mapInToOutPos: " + mapInToOutPos).powerSet(),
-                      bs -> !bs.isEmpty())));
-
-      resultBuilder.addAll(Util.transform(product, ImmutableBitSet::union));
+      // select key1, key1, val1, val2, key2 from ...
+      // the resulting unique keys would be {{0},{4}}, {{1},{4}}
+
+      Iterable<List<Integer>> product = Linq4j.product(Util.transform(colMask, mapInToOutPos::get));
+
+      resultBuilder.addAll(Util.transform(product, ImmutableBitSet::of));
     }
     return resultBuilder.build();
   }

core/src/main/java/org/apache/calcite/util/ImmutableBitSet.java

@@ -978,6 +978,18 @@ public static boolean allContain(Collection<ImmutableBitSet> bitSets, int bit) {
     return true;
   }
 
+  /**
+   * Returns whether this is a minimal set with respect to the specified collection of bitSets.
+   */
+  public boolean isMinimal(Collection<ImmutableBitSet> bitSets) {
+    for (ImmutableBitSet other : bitSets) {
+      if (!this.equals(other) && this.contains(other)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   /** Returns whether a given predicate evaluates to true for all bits in this
    * set. */
   public boolean allMatch(IntPredicate predicate) {

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

@@ -1448,6 +1448,17 @@ private void checkColumnUniquenessForFilterWithConstantColumns(String sql) {
         .withCatalogReaderFactory(factory)
         .assertThatUniqueKeysAre();
 
+    sql("select key1, key1, key2, value1 from s.composite_keys_table")
+        .withCatalogReaderFactory(factory)
+        .assertThatUniqueKeysAre(bitSetOf(0, 2), bitSetOf(1, 2));
+    sql("select key1, key2, key2, value1 from s.composite_keys_table")
+        .withCatalogReaderFactory(factory)
+        .assertThatUniqueKeysAre(bitSetOf(0, 1), bitSetOf(0, 2));
+
+    sql("select key1, key1, key2, key2, value1 from s.composite_keys_table")
+        .withCatalogReaderFactory(factory)
+        .assertThatUniqueKeysAre(bitSetOf(0, 2), bitSetOf(0, 3), bitSetOf(1, 2), bitSetOf(1, 3));
+
     // no column of composite keys
     sql("select value1 from s.composite_keys_table")
         .withCatalogReaderFactory(factory)
@@ -1640,7 +1651,7 @@ private static ImmutableBitSet bitSetOf(int... bits) {
   @Test void calcMultipleColumnsAreUniqueCalc() {
     sql("select empno, empno from emp")
         .convertingProjectAsCalc()
-        .assertThatUniqueKeysAre(bitSetOf(0), bitSetOf(1), bitSetOf(0, 1));
+        .assertThatUniqueKeysAre(bitSetOf(0), bitSetOf(1));
   }
 
   @Test void calcMultipleColumnsAreUniqueCalc2() {
@@ -1655,7 +1666,7 @@ private static ImmutableBitSet bitSetOf(int... bits) {
         + " from emp a1 join emp a2\n"
         + " on (a1.empno=a2.empno)")
         .convertingProjectAsCalc()
-        .assertThatUniqueKeysAre(bitSetOf(0), bitSetOf(1), bitSetOf(1, 2), bitSetOf(2));
+        .assertThatUniqueKeysAre(bitSetOf(0), bitSetOf(1), bitSetOf(2));
   }
 
   @Test void calcColumnsAreNonUniqueCalc() {

core/src/test/java/org/apache/calcite/util/ImmutableBitSetTest.java

@@ -668,6 +668,38 @@ private List<ImmutableBitSet> getSortedList() {
     assertFalse(ImmutableBitSet.allContain(collection2, 4));
   }
 
+  @Test void testIsMinimal() {
+    ImmutableBitSet se = ImmutableBitSet.of();
+    ImmutableBitSet s0 = ImmutableBitSet.of(0);
+    ImmutableBitSet s1 = ImmutableBitSet.of(1);
+    ImmutableBitSet s01 = ImmutableBitSet.of(0, 1);
+
+    assertTrue(se.isMinimal(Collections.singletonList(se)));
+    assertTrue(se.isMinimal(Collections.singletonList(s0)));
+    assertTrue(se.isMinimal(Collections.singletonList(s1)));
+    assertTrue(se.isMinimal(Collections.singletonList(s01)));
+
+    assertFalse(s0.isMinimal(Collections.singletonList(se)));
+    assertTrue(s0.isMinimal(Collections.singletonList(s0)));
+    assertTrue(s0.isMinimal(Collections.singletonList(s1)));
+    assertTrue(s0.isMinimal(Collections.singletonList(s01)));
+
+    assertFalse(s1.isMinimal(Collections.singletonList(se)));
+    assertTrue(s1.isMinimal(Collections.singletonList(s0)));
+    assertTrue(s1.isMinimal(Collections.singletonList(s1)));
+    assertTrue(s1.isMinimal(Collections.singletonList(s01)));
+
+    assertFalse(s01.isMinimal(Collections.singletonList(se)));
+    assertFalse(s01.isMinimal(Collections.singletonList(s0)));
+    assertFalse(s01.isMinimal(Collections.singletonList(s1)));
+    assertTrue(s01.isMinimal(Collections.singletonList(s01)));
+
+    assertTrue(se.isMinimal(Arrays.asList(s0, s1)));
+    assertTrue(s0.isMinimal(Arrays.asList(s0, s1)));
+    assertTrue(s1.isMinimal(Arrays.asList(s0, s1)));
+    assertFalse(s01.isMinimal(Arrays.asList(s0, s1)));
+  }
+
   /**
    * Test case for {@link ImmutableBitSet#anyMatch(IntPredicate)}
    * and {@link ImmutableBitSet#allMatch(IntPredicate)}.

testkit/src/main/java/org/apache/calcite/test/RelMetadataFixture.java

@@ -356,6 +356,7 @@ public RelMetadataFixture assertThatUniqueKeysAre(
         ImmutableSortedSet.copyOf(result),
         is(ImmutableSortedSet.copyOf(expectedUniqueKeys)));
     checkUniqueConsistent(rel);
+    assertThat(result.stream().allMatch(s -> s.isMinimal(result)), is(true));
     return this;
   }