apache · vlsi · Jan 2, 2020 · Jan 2, 2020
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
@@ -20,8 +20,10 @@
 import org.apache.calcite.plan.hep.HepRelVertex;
 import org.apache.calcite.plan.volcano.RelSubset;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.SingleRel;
 import org.apache.calcite.rel.convert.Converter;
 import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.Calc;
 import org.apache.calcite.rel.core.Correlate;
 import org.apache.calcite.rel.core.Exchange;
 import org.apache.calcite.rel.core.Filter;
@@ -42,12 +44,14 @@
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexProgram;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.util.BuiltInMethod;
 import org.apache.calcite.util.ImmutableBitSet;
 import org.apache.calcite.util.Pair;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Lists;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -206,7 +210,22 @@ public Boolean areColumnsUnique(Project rel, RelMetadataQuery mq,
     // Also need to map the input column set to the corresponding child
     // references
 
-    List<RexNode> projExprs = rel.getProjects();
+    return areProjectColumnsUnique(rel, mq, columns, ignoreNulls, rel.getProjects());
+  }
+
+  public Boolean areColumnsUnique(Calc rel, RelMetadataQuery mq,
+      ImmutableBitSet columns, boolean ignoreNulls) {
+    columns = decorateWithConstantColumnsFromPredicates(columns, rel, mq);
+    RexProgram program = rel.getProgram();
+
+    return areProjectColumnsUnique(rel, mq, columns, ignoreNulls,
+        Lists.transform(program.getProjectList(), program::expandLocalRef));
+  }
+
+  private Boolean areProjectColumnsUnique(
+      SingleRel rel, RelMetadataQuery mq,
+      ImmutableBitSet columns, boolean ignoreNulls, List<RexNode> projExprs) {
+    RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
     ImmutableBitSet.Builder childColumns = ImmutableBitSet.builder();
     for (int bit : columns) {
       RexNode projExpr = projExprs.get(bit);
@@ -226,8 +245,6 @@ public Boolean areColumnsUnique(Project rel, RelMetadataQuery mq,
         if (!(castOperand instanceof RexInputRef)) {
           continue;
         }
-        RelDataTypeFactory typeFactory =
-            rel.getCluster().getTypeFactory();
         RelDataType castType =
             typeFactory.createTypeWithNullability(
                 projExpr.getType(), true);

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
@@ -16,8 +16,11 @@
  */
 package org.apache.calcite.rel.metadata;
 
+import org.apache.calcite.linq4j.Linq4j;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.SingleRel;
 import org.apache.calcite.rel.core.Aggregate;
+import org.apache.calcite.rel.core.Calc;
 import org.apache.calcite.rel.core.Correlate;
 import org.apache.calcite.rel.core.Filter;
 import org.apache.calcite.rel.core.Join;
@@ -29,12 +32,17 @@
 import org.apache.calcite.rel.core.TableScan;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexProgram;
 import org.apache.calcite.util.BuiltInMethod;
 import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.Permutation;
 
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -82,57 +90,78 @@ public Set<ImmutableBitSet> getUniqueKeys(TableModify rel, RelMetadataQuery mq,
 
   public Set<ImmutableBitSet> getUniqueKeys(Project rel, RelMetadataQuery mq,
       boolean ignoreNulls) {
+    return getProjectUniqueKeys(rel, mq, ignoreNulls, rel.getProjects());
+  }
+
+  public Set<ImmutableBitSet> getUniqueKeys(Calc rel, RelMetadataQuery mq,
+      boolean ignoreNulls) {
+    RexProgram program = rel.getProgram();
+    Permutation permutation = program.getPermutation();
+    return getProjectUniqueKeys(rel, mq, ignoreNulls,
+        Lists.transform(program.getProjectList(), program::expandLocalRef));
+  }
+
+  private Set<ImmutableBitSet> getProjectUniqueKeys(SingleRel rel, RelMetadataQuery mq,
+      boolean ignoreNulls, List<RexNode> projExprs) {
     // LogicalProject maps a set of rows to a different set;
     // Without knowledge of the mapping function(whether it
     // preserves uniqueness), it is only safe to derive uniqueness
     // info from the child of a project when the mapping is f(a) => a.
     //
     // Further more, the unique bitset coming from the child needs
     // to be mapped to match the output of the project.
-    final Map<Integer, Integer> mapInToOutPos = new HashMap<>();
-    final List<RexNode> projExprs = rel.getProjects();
-    final Set<ImmutableBitSet> projUniqueKeySet = new HashSet<>();
+
+    // Single input can be mapped to multiple outputs
+    ImmutableMultimap.Builder<Integer, Integer> inToOutPosBuilder = ImmutableMultimap.builder();
+    ImmutableBitSet.Builder mappedInColumnsBuilder = ImmutableBitSet.builder();
 
     // Build an input to output position map.
     for (int i = 0; i < projExprs.size(); i++) {
       RexNode projExpr = projExprs.get(i);
       if (projExpr instanceof RexInputRef) {
-        mapInToOutPos.put(((RexInputRef) projExpr).getIndex(), i);
+        int inputIndex = ((RexInputRef) projExpr).getIndex();
+        inToOutPosBuilder.put(inputIndex, i);
+        mappedInColumnsBuilder.set(inputIndex);
       }
     }
+    ImmutableBitSet inColumnsUsed = mappedInColumnsBuilder.build();
 
-    if (mapInToOutPos.isEmpty()) {
+    if (inColumnsUsed.isEmpty()) {
       // if there's no RexInputRef in the projected expressions
       // return empty set.
-      return projUniqueKeySet;
+      return ImmutableSet.of();
     }
 
     Set<ImmutableBitSet> childUniqueKeySet =
         mq.getUniqueKeys(rel.getInput(), ignoreNulls);
 
-    if (childUniqueKeySet != null) {
-      // Now add to the projUniqueKeySet the child keys that are fully
-      // projected.
-      for (ImmutableBitSet colMask : childUniqueKeySet) {
-        ImmutableBitSet.Builder tmpMask = ImmutableBitSet.builder();
-        boolean completeKeyProjected = true;
-        for (int bit : colMask) {
-          if (mapInToOutPos.containsKey(bit)) {
-            tmpMask.set(mapInToOutPos.get(bit));
-          } else {
-            // Skip the child unique key if part of it is not
-            // projected.
-            completeKeyProjected = false;
-            break;
-          }
-        }
-        if (completeKeyProjected) {
-          projUniqueKeySet.add(tmpMask.build());
-        }
-      }
+    if (childUniqueKeySet == null) {
+      return ImmutableSet.of();
     }
 
-    return projUniqueKeySet;
+    Map<Integer, ImmutableBitSet> mapInToOutPos =
+        Maps.transformValues(inToOutPosBuilder.build().asMap(), ImmutableBitSet::of);
+
+    ImmutableSet.Builder<ImmutableBitSet> resultBuilder = ImmutableSet.builder();
+    // Now add to the projUniqueKeySet the child keys that are fully
+    // projected.
+    for (ImmutableBitSet colMask : childUniqueKeySet) {
+      ImmutableBitSet.Builder tmpMask = ImmutableBitSet.builder();
+      if (!inColumnsUsed.contains(colMask)) {
+        // colMask contains a column that is not projected as RexInput => the key is not unique
+        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(
+          Iterables.transform(colMask,
+              in -> Iterables.filter(mapInToOutPos.get(in).powerSet(), bs -> !bs.isEmpty())));
+
+      resultBuilder.addAll(Iterables.transform(product, ImmutableBitSet::union));
+    }
+    return resultBuilder.build();
   }
 
   public Set<ImmutableBitSet> getUniqueKeys(Join rel, RelMetadataQuery mq,

diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -23,6 +23,7 @@
 import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.calcite.plan.RelOptPredicateList;
 import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.plan.RelTraitSet;
 import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelCollationTraitDef;
@@ -49,6 +50,7 @@
 import org.apache.calcite.rel.core.Union;
 import org.apache.calcite.rel.core.Values;
 import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.logical.LogicalCalc;
 import org.apache.calcite.rel.logical.LogicalExchange;
 import org.apache.calcite.rel.logical.LogicalFilter;
 import org.apache.calcite.rel.logical.LogicalJoin;
@@ -80,8 +82,10 @@
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexProgram;
 import org.apache.calcite.rex.RexTableInputRef;
 import org.apache.calcite.rex.RexTableInputRef.RelTableRef;
+import org.apache.calcite.runtime.SqlFunctions;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlSpecialOperator;
@@ -101,6 +105,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSortedSet;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Sets;
@@ -126,6 +131,8 @@
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Function;
+import javax.annotation.Nonnull;
 
 import static org.apache.calcite.test.Matchers.within;
 
@@ -189,17 +196,15 @@ private RelNode convertSql(String sql) {
     return convertSql(tester, sql);
   }
 
-  private RelNode convertSql(Tester tester, String sql) {
+  private static RelNode convertSql(Tester tester, String sql) {
     final RelRoot root = tester.convertSqlToRel(sql);
     root.rel.getCluster().setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
     return root.rel;
   }
 
   private RelNode convertSql(String sql, boolean typeCoercion) {
     final Tester tester = typeCoercion ? this.tester : this.strictTester;
-    final RelRoot root = tester.convertSqlToRel(sql);
-    root.rel.getCluster().setMetadataProvider(DefaultRelMetadataProvider.INSTANCE);
-    return root.rel;
+    return convertSql(tester, sql);
   }
 
   private void checkPercentageOriginalRows(String sql, double expected) {
@@ -889,24 +894,31 @@ private void checkRelSelectivity(
   /**
    * Checks result of getting unique keys for sql, using specific tester.
    */
-  private void checkGetUniqueKeys(Tester tester,
-      String sql, Set<ImmutableBitSet> expectedUniqueKeySet) {
-    RelNode rel = convertSql(tester, sql);
+  private void checkGetUniqueKeys(String sql, Set<ImmutableBitSet> expectedUniqueKeySet,
+      Function<String, RelNode> converter) {
+    RelNode rel = converter.apply(sql);
     final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
     Set<ImmutableBitSet> result = mq.getUniqueKeys(rel);
-    assertThat(result, CoreMatchers.equalTo(expectedUniqueKeySet));
+    assertEquals(
+        ImmutableSortedSet.copyOf(expectedUniqueKeySet),
+        ImmutableSortedSet.copyOf(result),
+        () -> "unique keys, sql: " + sql + ", rel: " + RelOptUtil.toString(rel));
     assertUniqueConsistent(rel);
   }
 
+  /**
+   * Checks result of getting unique keys for sql, using specific tester.
+   */
+  private void checkGetUniqueKeys(Tester tester,
+      String sql, Set<ImmutableBitSet> expectedUniqueKeySet) {
+    checkGetUniqueKeys(sql, expectedUniqueKeySet, s -> convertSql(tester, s));
+  }
+
   /**
    * Checks result of getting unique keys for sql, using default tester.
    */
   private void checkGetUniqueKeys(String sql, Set<ImmutableBitSet> expectedUniqueKeySet) {
-    RelNode rel = convertSql(sql);
-    final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
-    Set<ImmutableBitSet> result = mq.getUniqueKeys(rel);
-    assertThat(result, CoreMatchers.equalTo(expectedUniqueKeySet));
-    assertUniqueConsistent(rel);
+    checkGetUniqueKeys(tester, sql, expectedUniqueKeySet);
   }
 
   /** Asserts that {@link RelMetadataQuery#getUniqueKeys(RelNode)}
@@ -919,7 +931,9 @@ private void assertUniqueConsistent(RelNode rel) {
         ImmutableBitSet.range(0, rel.getRowType().getFieldCount());
     for (ImmutableBitSet key : allCols.powerSet()) {
       Boolean result2 = mq.areColumnsUnique(rel, key);
-      assertTrue(result2 == null || result2 == isUnique(uniqueKeys, key));
+      assertEquals(isUnique(uniqueKeys, key), SqlFunctions.isTrue(result2),
+          () -> "areColumnsUnique. key: " + key + ", uniqueKeys: " + uniqueKeys
+              + ", rel: " + RelOptUtil.toString(rel));
     }
   }
 
@@ -1180,6 +1194,66 @@ private void checkColumnUniquenessForFilterWithConstantColumns(String sql) {
         ImmutableSet.of());
   }
 
+  private static ImmutableBitSet bitSetOf(int... bits) {
+    return ImmutableBitSet.of(bits);
+  }
+
+  @Test public void calcColumnsAreUniqueSimpleCalc() {
+    checkGetUniqueKeys("select empno, empno*0 from emp",
+        ImmutableSet.of(bitSetOf(0)),
+        convertProjectAsCalc());
+  }
+
+  @Test public void calcColumnsAreUniqueCalcWithFirstConstant() {
+    checkGetUniqueKeys("select 1, empno, empno*0 from emp",
+        ImmutableSet.of(bitSetOf(1)),
+        convertProjectAsCalc());
+
+  }
+  @Test public void calcMultipleColumnsAreUniqueCalc() {
+    checkGetUniqueKeys("select empno, empno from emp",
+        ImmutableSet.of(bitSetOf(0), bitSetOf(1), bitSetOf(0, 1)),
+        convertProjectAsCalc());
+  }
+
+  @Test public void calcMultipleColumnsAreUniqueCalc2() {
+    checkGetUniqueKeys(
+        "select a1.empno, a2.empno from emp a1 join emp a2 on (a1.empno=a2.empno)",
+        ImmutableSet.of(bitSetOf(0), bitSetOf(1), bitSetOf(0, 1)),
+        convertProjectAsCalc());
+  }
+
+  @Test public void calcMultipleColumnsAreUniqueCalc3() {
+    checkGetUniqueKeys(
+        "select a1.empno, a2.empno, a2.empno\n"
+        + " from emp a1 join emp a2\n"
+        + " on (a1.empno=a2.empno)",
+        ImmutableSet.of(
+            bitSetOf(0), bitSetOf(0, 1), bitSetOf(0, 1, 2), bitSetOf(0, 2),
+            bitSetOf(1), bitSetOf(1, 2), bitSetOf(2)),
+        convertProjectAsCalc());
+  }
+
+  @Test public void calcColumnsAreNonUniqueCalc() {
+    checkGetUniqueKeys("select empno*0 from emp",
+        ImmutableSet.of(),
+        convertProjectAsCalc());
+  }
+
+  @Nonnull
+  private Function<String, RelNode> convertProjectAsCalc() {
+    return s -> {
+      Project project = (Project) convertSql(s);
+      RexProgram program = RexProgram.create(
+          project.getInput().getRowType(),
+          project.getProjects(),
+          null,
+          project.getRowType(),
+          project.getCluster().getRexBuilder());
+      return LogicalCalc.create(project.getInput(), program);
+    };
+  }
+
   @Test public void testBrokenCustomProviderWithMetadataFactory() {
     final List<String> buf = new ArrayList<>();
     ColTypeImpl.THREAD_LIST.set(buf);