[CALCITE-953] Improve RelMdPredicates to deal with RexLiteral (Pengcheng Xiong)

Author: Pengcheng Xiong

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

@@ -35,11 +35,13 @@
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexPermuteInputsShuttle;
 import org.apache.calcite.rex.RexUtil;
 import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.util.BitSets;
 import org.apache.calcite.util.BuiltInMethod;
 import org.apache.calcite.util.ImmutableBitSet;
@@ -147,6 +149,7 @@ public RelOptPredicateList getPredicates(TableScan table) {
    */
   public RelOptPredicateList getPredicates(Project project) {
     RelNode child = project.getInput();
+    final RexBuilder rexBuilder = project.getCluster().getRexBuilder();
     RelOptPredicateList childInfo =
         RelMetadataQuery.getPulledUpPredicates(child);
 
@@ -175,6 +178,20 @@ public RelOptPredicateList getPredicates(Project project) {
         projectPullUpPredicates.add(r);
       }
     }
+
+    // Project can also generate constants. We need to include them.
+    for (Ord<RexNode> expr : Ord.zip(project.getProjects())) {
+      if (RexLiteral.isNullLiteral(expr.e)) {
+        projectPullUpPredicates.add(
+            rexBuilder.makeCall(SqlStdOperatorTable.IS_NULL,
+                rexBuilder.makeInputRef(project, expr.i)));
+      } else if (expr.e instanceof RexLiteral) {
+        final RexLiteral literal = (RexLiteral) expr.e;
+        projectPullUpPredicates.add(
+            rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
+                rexBuilder.makeInputRef(project, expr.i), literal));
+      }
+    }
     return RelOptPredicateList.of(projectPullUpPredicates);
   }
 

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

@@ -1079,6 +1079,31 @@ private void checkPredicates(RelOptCluster cluster, RelOptTable empTable,
     assertThat(predicates.rightInferredPredicates.isEmpty(), is(true));
   }
 
+  /**
+   * Unit test for
+   * {@link org.apache.calcite.rel.metadata.RelMdPredicates#getPredicates(Aggregate)}.
+   */
+  @Test public void testPullUpPredicatesFromAggregation() {
+    final String sql = "select a, max(b) from (\n"
+        + "  select 1 as a, 2 as b from emp)subq\n"
+        + "group by a";
+    final Aggregate rel = (Aggregate) convertSql(sql);
+    RelOptPredicateList inputSet = RelMetadataQuery.getPulledUpPredicates(rel);
+    ImmutableList<RexNode> pulledUpPredicates = inputSet.pulledUpPredicates;
+    assertThat(pulledUpPredicates.toString(), is("[=($0, 1)]"));
+  }
+
+  @Test public void testPullUpPredicatesFromProject() {
+    final String sql = "select deptno, mgr, x, 'y' as y from (\n"
+        + "  select deptno, mgr, cast(null as integer) as x\n"
+        + "  from emp\n"
+        + "  where mgr is null and deptno < 10)";
+    final RelNode rel = convertSql(sql);
+    RelOptPredicateList list = RelMetadataQuery.getPulledUpPredicates(rel);
+    assertThat(list.pulledUpPredicates.toString(),
+        is("[IS NULL($1), <($0, 10), IS NULL($2), =($3, 'y')]"));
+  }
+
   /** Custom metadata interface. */
   public interface ColType extends Metadata {
     String getColType(int column);

core/src/test/resources/org/apache/calcite/test/RelOptRulesTest.xml

@@ -913,7 +913,7 @@ LogicalProject(EXPR$0=[CAST(CASE(IS NULL($1), IS NULL($0), IS NULL($0), IS NULL(
         </Resource>
         <Resource name="planAfter">
             <![CDATA[
-LogicalProject(EXPR$0=[CASE(IS NULL($1), IS NULL($0), CAST(=($1, $0)):BOOLEAN NOT NULL)])
+LogicalProject(EXPR$0=[CASE(IS NULL($1), false, CAST(=($1, 2)):BOOLEAN NOT NULL)])
   LogicalProject(EXPR$0=[2], EXPR$1=[null])
     LogicalValues(tuples=[[{ 0 }]])
 ]]>
@@ -2908,10 +2908,10 @@ LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$
         </Resource>
         <Resource name="planAfter">
             <![CDATA[
-LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], DEPTNO0=[$9], NAME=[$10])
+LogicalProject(EMPNO=[10], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], DEPTNO0=[10], NAME=[$10])
   LogicalProject(EMPNO=[10], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], DEPTNO0=[10], NAME=[$11])
-    LogicalJoin(condition=[AND(=(10, $10), =($9, $12))], joinType=[inner])
-      LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], $f9=[+($7, $0)])
+    LogicalJoin(condition=[=($9, 15)], joinType=[inner])
+      LogicalProject(EMPNO=[10], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], $f9=[+($7, 10)])
         LogicalProject(EMPNO=[10], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8])
           LogicalFilter(condition=[=($0, 10)])
             LogicalTableScan(table=[[CATALOG, SALES, EMP]])