[CALCITE-6640] RelMdUniqueKeys grows exponentially when key columns a…

…re repeated in projections Avoid generating non-minimal/redundant key sets when computing the unique keys for columns that are repeated in the output.
apache · Oct 23, 2024 · b5e2fd2 · b5e2fd2
1 parent b0ecd06
commit b5e2fd2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 21 deletions.
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
@@ -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();
   }

diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java b/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() {