[CALCITE-7257] Subqueries cannot be decorrelated if join condition contains RexFieldAccess

Author: iwanttobepowerful

core/src/main/java/org/apache/calcite/sql2rel/RelDecorrelator.java

@@ -1772,8 +1772,11 @@ private static boolean isWidening(RelDataType type, RelDataType type1) {
       return null;
     }
 
+
+    Frame newLeftFrame = getNewLeftFrame(rel, oldLeft, leftFrame, isCorVarDefined);
+
     RelNode newJoin = relBuilder
-        .push(leftFrame.r)
+        .push(newLeftFrame.r)
         .push(rightFrame.r)
         .join(rel.getJoinType(),
             decorrelateExpr(castNonNull(currentRel), map, cm, rel.getCondition()),
@@ -1785,25 +1788,23 @@ private static boolean isWidening(RelDataType type, RelDataType type1) {
     Map<Integer, Integer> mapOldToNewOutputs = new HashMap<>();
 
     int oldLeftFieldCount = oldLeft.getRowType().getFieldCount();
-    int newLeftFieldCount = leftFrame.r.getRowType().getFieldCount();
+    int newLeftFieldCount = newLeftFrame.r.getRowType().getFieldCount();
 
     int oldRightFieldCount = oldRight.getRowType().getFieldCount();
     //noinspection AssertWithSideEffects
     assert rel.getRowType().getFieldCount()
         == oldLeftFieldCount + oldRightFieldCount;
 
     // Left input positions are not changed.
-    mapOldToNewOutputs.putAll(leftFrame.oldToNewOutputs);
-
+    mapOldToNewOutputs.putAll(newLeftFrame.oldToNewOutputs);
     // Right input positions are shifted by newLeftFieldCount.
     for (int i = 0; i < oldRightFieldCount; i++) {
       mapOldToNewOutputs.put(i + oldLeftFieldCount,
           requireNonNull(rightFrame.oldToNewOutputs.get(i)) + newLeftFieldCount);
     }
 
     final NavigableMap<CorDef, Integer> corDefOutputs =
-        new TreeMap<>(leftFrame.corDefOutputs);
-
+        new TreeMap<>(newLeftFrame.corDefOutputs);
     // Right input positions are shifted by newLeftFieldCount.
     for (Map.Entry<CorDef, Integer> entry
         : rightFrame.corDefOutputs.entrySet()) {
@@ -1813,6 +1814,46 @@ private static boolean isWidening(RelDataType type, RelDataType type1) {
     return register(rel, newJoin, mapOldToNewOutputs, corDefOutputs);
   }
 
+  private Frame getNewLeftFrame(Join rel, RelNode oldLeft, Frame leftFrame,
+      boolean isCorVarDefined) {
+    Frame newLeftFrame = leftFrame;
+    boolean joinConditionContainsFieldAccess = RexUtil.containsFieldAccess(rel.getCondition());
+    if (joinConditionContainsFieldAccess && isCorVarDefined) {
+      final CorelMap localCorelMap = new CorelMapBuilder().build(rel);
+      final List<CorRef> corVarList = new ArrayList<>(localCorelMap.mapRefRelToCorRef.values());
+      Collections.sort(corVarList);
+
+      // 1. Build valueGen and a new corDefOutputs based on corVarList.
+      final NavigableMap<CorDef, Integer> newLeftCorDefOutputs = new TreeMap<>();
+      final RelNode valueGen =
+          requireNonNull(createValueGenerator(corVarList, 0, newLeftCorDefOutputs));
+      int valueGenFieldCount = valueGen.getRowType().getFieldCount();
+      RelNode newLeft = relBuilder
+          .push(valueGen)
+          .push(leftFrame.r)
+          .join(JoinRelType.INNER, relBuilder.literal(true))
+          .build();
+
+      // 3. Merge leftFrame.corDefOutputs into newLeftCorDefOutputs, with valueGenFieldCount offset
+      for (Map.Entry<CorDef,Integer> e : leftFrame.corDefOutputs.entrySet()) {
+        newLeftCorDefOutputs.put(e.getKey(), valueGenFieldCount + e.getValue());
+      }
+
+      // 4. Compute the new positions of oldLeft's oldToNewOutputs in newLeft.
+      final Map<Integer,Integer> newLeftOldToNew = new HashMap<>();
+      for (Map.Entry<Integer,Integer> e : leftFrame.oldToNewOutputs.entrySet()) {
+        newLeftOldToNew.put(e.getKey(), valueGenFieldCount + e.getValue());
+      }
+
+      // 5. Build a new Frame with oldLeft, newLeft, and the mappings above.
+      newLeftFrame = new Frame(oldLeft, newLeft, newLeftCorDefOutputs, newLeftOldToNew);
+
+      // 6. Update map so that oldLeft is associated with the new Frame.
+      map.put(oldLeft, newLeftFrame);
+    }
+    return newLeftFrame;
+  }
+
   private static RexInputRef getNewForOldInputRef(RelNode currentRel,
       Map<RelNode, Frame> map, RexInputRef oldInputRef) {
     requireNonNull(currentRel, "currentRel");

core/src/test/java/org/apache/calcite/sql2rel/RelDecorrelatorTest.java

@@ -584,4 +584,78 @@ public static Frameworks.ConfigBuilder config() {
         + "        LogicalTableScan(table=[[scott, DEPT]])\n";
     assertThat(decorrelatedNoRules, hasTree(planDecorrelatedNoRules));
   }
+
+  /** Test case for <a href="https://issues.apache.org/jira/browse/CALCITE-7257">[CALCITE-7257]
+   * Subqueries cannot be decorrelated if join condition contains RexFieldAccess</a>. */
+  @Test void testJoinConditionContainsRexFieldAccess() {
+    final FrameworkConfig frameworkConfig = config().build();
+    final RelBuilder builder = RelBuilder.create(frameworkConfig);
+    final RelOptCluster cluster = builder.getCluster();
+    final Planner planner = Frameworks.getPlanner(frameworkConfig);
+    final String sql = ""
+        + "SELECT E1.* \n"
+        + "FROM\n"
+        + "  EMP E1\n"
+        + "WHERE\n"
+        + "  E1.EMPNO = (\n"
+        + "    SELECT D1.DEPTNO FROM DEPT D1\n"
+        + "    WHERE E1.ENAME IN (SELECT B1.ENAME FROM BONUS B1))";
+    final RelNode originalRel;
+    try {
+      final SqlNode parse = planner.parse(sql);
+      final SqlNode validate = planner.validate(parse);
+      originalRel = planner.rel(validate).rel;
+    } catch (Exception e) {
+      throw TestUtil.rethrow(e);
+    }
+
+    final HepProgram hepProgram = HepProgram.builder()
+        .addRuleCollection(
+            ImmutableList.of(
+                // SubQuery program rules
+                CoreRules.FILTER_SUB_QUERY_TO_CORRELATE,
+                CoreRules.PROJECT_SUB_QUERY_TO_CORRELATE,
+                CoreRules.JOIN_SUB_QUERY_TO_CORRELATE))
+        .build();
+    final Program program =
+        Programs.of(hepProgram, true,
+            requireNonNull(cluster.getMetadataProvider()));
+    final RelNode before =
+        program.run(cluster.getPlanner(), originalRel, cluster.traitSet(),
+            Collections.emptyList(), Collections.emptyList());
+    final String planBefore = ""
+        + "LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7])\n"
+        + "  LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7])\n"
+        + "    LogicalFilter(condition=[=($0, CAST($8):SMALLINT)])\n"
+        + "      LogicalCorrelate(correlation=[$cor0], joinType=[left], requiredColumns=[{1}])\n"
+        + "        LogicalTableScan(table=[[scott, EMP]])\n"
+        + "        LogicalAggregate(group=[{}], agg#0=[SINGLE_VALUE($0)])\n"
+        + "          LogicalProject(DEPTNO=[$0])\n"
+        + "            LogicalProject(DEPTNO=[$0], DNAME=[$1], LOC=[$2])\n"
+        + "              LogicalJoin(condition=[=($cor0.ENAME, $3)], joinType=[inner])\n"
+        + "                LogicalTableScan(table=[[scott, DEPT]])\n"
+        + "                LogicalProject(ENAME=[$0])\n"
+        + "                  LogicalTableScan(table=[[scott, BONUS]])\n";
+    assertThat(before, hasTree(planBefore));
+
+    // Decorrelate without any rules, just "purely" decorrelation algorithm on RelDecorrelator
+    final RelNode after =
+        RelDecorrelator.decorrelateQuery(before, builder, RuleSets.ofList(Collections.emptyList()),
+            RuleSets.ofList(Collections.emptyList()));
+    final String planAfter = ""
+        + "LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7])\n"
+        + "  LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], ENAME0=[$8], $f1=[CAST($9):TINYINT])\n"
+        + "    LogicalJoin(condition=[AND(=($1, $8), =($0, CAST($9):SMALLINT))], joinType=[inner])\n"
+        + "      LogicalTableScan(table=[[scott, EMP]])\n"
+        + "      LogicalAggregate(group=[{0}], agg#0=[SINGLE_VALUE($1)])\n"
+        + "        LogicalProject(ENAME=[$0], DEPTNO=[$1])\n"
+        + "          LogicalJoin(condition=[=($0, $4)], joinType=[inner])\n"
+        + "            LogicalJoin(condition=[true], joinType=[inner])\n"
+        + "              LogicalProject(ENAME=[$1])\n"
+        + "                LogicalTableScan(table=[[scott, EMP]])\n"
+        + "              LogicalTableScan(table=[[scott, DEPT]])\n"
+        + "            LogicalProject(ENAME=[$0])\n"
+        + "              LogicalTableScan(table=[[scott, BONUS]])\n";
+    assertThat(after, hasTree(planAfter));
+  }
 }