[CALCITE-7859] SqlToRelConverter.convertMerge produces a wrong plan or fails when a sub-query in MERGE ... UPDATE SET is expanded

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

@@ -52,6 +52,7 @@
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.core.RelFactories;
 import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.core.TableModify;
 import org.apache.calcite.rel.hint.HintStrategyTable;
 import org.apache.calcite.rel.hint.Hintable;
 import org.apache.calcite.rel.hint.RelHint;
@@ -4884,61 +4885,103 @@
             () -> "sourceSelect for " + call);
     RelNode mergeSourceRel = convertSelect(sourceSelect, false);
 
-    // then, convert the insert statement so we can get the insert
-    // values expressions
-    SqlInsert insertCall = call.getInsertCall();
-    int nLevel1Exprs = 0;
-    List<RexNode> level1InsertExprs = null;
-    List<RexNode> level2InsertExprs = null;
-    if (insertCall != null) {
-      RelNode insertRel = convertInsert(insertCall);
-
-      // if there are 2 level of projections in the insert source, combine
-      // them into a single project; level1 refers to the topmost project;
-      // the level1 projection contains references to the level2
-      // expressions, except in the case where no target expression was
-      // provided, in which case, the expression is the default value for
-      // the column; or if the expressions directly map to the source
-      // table
-      level1InsertExprs =
-          ((LogicalProject) insertRel.getInput(0)).getProjects();
-      if (insertRel.getInput(0).getInput(0) instanceof LogicalProject) {
-        level2InsertExprs =
-            ((LogicalProject) insertRel.getInput(0).getInput(0))
-                .getProjects();
-      }
-      nLevel1Exprs = level1InsertExprs.size();
-    }
-
-    LogicalJoin join = (LogicalJoin) mergeSourceRel.getInput(0);
-    int nSourceFields = join.getLeft().getRowType().getFieldCount();
+    // The input of the source projection is the join of the source and the
+    // target, possibly wrapped by the joins or correlates that expand the
+    // sub-queries of the SET expressions. Source fields come first in all of
+    // them.
+    RelNode mergeInput = mergeSourceRel.getInput(0);
+    final SqlNode sourceFrom =
+        requireNonNull(sourceSelect.getFrom(), () -> "from for " + sourceSelect);
+    final SqlNode sourceRef = ((SqlJoin) sourceFrom).getLeft();
+    final SqlValidatorNamespace sourceNamespace =
+        requireNonNull(validator().getNamespace(sourceRef),
+            () -> "namespace for " + sourceRef);
+    final int nSourceFields = sourceNamespace.getRowType().getFieldCount();
+    final Set<CorrelationId> variablesSet =
+        new HashSet<>(((LogicalProject) mergeSourceRel).getVariablesSet());
+
     final List<RexNode> projects = new ArrayList<>();
-    for (int level1Idx = 0; level1Idx < nLevel1Exprs; level1Idx++) {
-      requireNonNull(level1InsertExprs, "level1InsertExprs");
-      if ((level2InsertExprs != null)
-          && (level1InsertExprs.get(level1Idx) instanceof RexInputRef)) {
-        int level2Idx =
-            ((RexInputRef) level1InsertExprs.get(level1Idx)).getIndex();
-        projects.add(level2InsertExprs.get(level2Idx));
-      } else {
-        projects.add(level1InsertExprs.get(level1Idx));
-      }
+    final SqlInsert insertCall = call.getInsertCall();
+    if (insertCall != null) {
+      final Project insertProject = convertMergeInsert(insertCall, mergeInput);
+      projects.addAll(insertProject.getProjects());
+      mergeInput = insertProject.getInput();
+      variablesSet.addAll(insertProject.getVariablesSet());
     }
+
     if (updateCall != null) {
       final LogicalProject project = (LogicalProject) mergeSourceRel;
       projects.addAll(
           Util.skip(project.getProjects(), nSourceFields));
     }
 
-    // Keep the correlation variables set by the source projection: the SET
-    // expressions may contain correlated sub-queries.
-    relBuilder.push(join)
-        .project(projects, ImmutableList.of(), false,
-            ((LogicalProject) mergeSourceRel).getVariablesSet());
+    // Keep the correlation variables: the SET expressions and the insert
+    // values may contain correlated sub-queries.
+    relBuilder.push(mergeInput)
+        .project(projects, ImmutableList.of(), false, variablesSet);
+    return createMergeModify(targetTable, relBuilder.build(),
+        targetColumnNameList);
+  }
+
+  /** Converts the values of the INSERT of a MERGE over {@code mergeInput}, the
+   * input of the MERGE source projection.
+   *
+   * @return A projection whose expressions are the inserted row, in the order
+   * of the target columns; its input starts with the fields of
+   * {@code mergeInput} */
+  private Project convertMergeInsert(SqlInsert insertCall, RelNode mergeInput) {
+    final SqlSelect insertSource = (SqlSelect) insertCall.getSource();
+    final Blackboard bb =
+        createBlackboard(validator().getSelectScope(insertSource), null, false);
+    bb.setRoot(mergeInput, true);
+    callStack.push(insertCall);
+    convertSelectList(bb, bb, insertSource, ImmutableList.of(),
+        ImmutableList.of(), null);
+    callStack.pop();
+    final Project valuesProject = (Project) bb.root();
+
+    // Order the values by the target columns and add the defaults. The
+    // projection is built by RelBuilder on top of the values projection, so
+    // it may be the values projection itself (identity), a projection over
+    // it, or a projection merged with it (over its input).
+    final RelNode columns = convertColumnList(insertCall, valuesProject);
+    if (columns == valuesProject) {
+      return valuesProject;
+    }
+    final Project columnsProject = (Project) columns;
+    final List<RexNode> exprs;
+    if (columnsProject.getInput() == valuesProject) {
+      exprs =
+          RelOptUtil.pushPastProject(columnsProject.getProjects(),
+              valuesProject);
+    } else if (columnsProject.getInput() == valuesProject.getInput()) {
+      exprs = columnsProject.getProjects();
+    } else {
+      throw new AssertionError("unexpected input of insert columns: "
+          + columnsProject.getInput());
+    }
+    return LogicalProject.create(valuesProject.getInput(), ImmutableList.of(),
+        exprs, columnsProject.getRowType().getFieldNames(),
+        valuesProject.getVariablesSet());
+  }
 
-    return LogicalTableModify.create(targetTable, catalogReader,
-        relBuilder.build(), LogicalTableModify.Operation.MERGE,
-        targetColumnNameList, null, false);
+  /** Creates the {@link TableModify} of a MERGE.
+   *
+   * <p>Like {@link #createModify}, it lets a {@link ModifiableTable} create the
+   * node, so a MERGE is converted the same way as an INSERT, UPDATE or DELETE
+   * on the same table. */
+  private RelNode createMergeModify(RelOptTable targetTable, RelNode source,
+      List<String> updateColumnList) {
+    final ModifiableTable modifiableTable =
+        targetTable.unwrap(ModifiableTable.class);
+    if (modifiableTable != null
+        && modifiableTable == targetTable.unwrap(Table.class)) {
+      return modifiableTable.toModificationRel(cluster, targetTable,
+          catalogReader, source, TableModify.Operation.MERGE,
+          updateColumnList, null, false);
+    }
+    return LogicalTableModify.create(targetTable, catalogReader, source,
+        TableModify.Operation.MERGE, updateColumnList, null, false);
   }
 
   /**
@@ -6615,7 +6658,7 @@
       final RelDataType type = call.getType();
       List<RexNode> exprs = call.getOperands();
 
       SqlFunction histogramOp = !ENABLE_HISTOGRAM_AGG
           ? null
           : getHistogramOp(aggOp);
 

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

@@ -4171,6 +4171,103 @@ void checkCorrelatedMapSubQuery(boolean expand) {
     sql(sql).withExpand(false).ok();
   }
 
+  /** Tests a MERGE whose UPDATE contains a correlated sub-query, with
+   * sub-queries expanded into joins ({@code expand = true}). */
+  @Test void testMergeUpdateCorrelatedSubQueryExpanded() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update\n"
+        + "set sal = (select d.deptno from dept d where d.deptno = e.deptno)";
+    sql(sql).withExpand(true).ok();
+  }
+
+  /** As {@link #testMergeUpdateCorrelatedSubQueryExpanded()}, but without
+   * decorrelation, so that the plan keeps the {@code LogicalCorrelate}. */
+  @Test void testMergeUpdateCorrelatedSubQueryExpandedNoDecorrelate() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update\n"
+        + "set sal = (select d.deptno from dept d where d.deptno = e.deptno)";
+    sql(sql).withExpand(true).withDecorrelate(false).ok();
+  }
+
+  /** Tests a MERGE with an expanded sub-query in the UPDATE and an INSERT
+   * branch, which turns the join of the source and the target into a LEFT
+   * join. */
+  @Test void testMergeUpdateSubQueryExpandedWithInsert() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update\n"
+        + "set sal = (select max(d.deptno) from dept d)\n"
+        + "when not matched then insert (empno, ename, sal)\n"
+        + "values (t.empno, t.ename, 10)";
+    sql(sql).withExpand(true).ok();
+  }
+
+  /** Tests a MERGE whose UPDATE contains an uncorrelated sub-query, with
+   * sub-queries expanded into joins ({@code expand = true}). */
+  @Test void testMergeUpdateUncorrelatedSubQueryExpanded() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update\n"
+        + "set sal = (select max(d.deptno) from dept d)";
+    sql(sql).withExpand(true).ok();
+  }
+
+  /** Tests a MERGE whose INSERT values contain an uncorrelated sub-query,
+   * with sub-queries expanded into joins ({@code expand = true}). */
+  @Test void testMergeInsertSubQueryExpanded() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update set sal = 1\n"
+        + "when not matched then insert (empno, sal)\n"
+        + "values (t.empno, (select max(deptno) from dept))";
+    sql(sql).withExpand(true).ok();
+  }
+
+  /** As {@link #testMergeInsertSubQueryExpanded()}, but with sub-queries not
+   * expanded. */
+  @Test void testMergeInsertSubQuery() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update set sal = 1\n"
+        + "when not matched then insert (empno, sal)\n"
+        + "values (t.empno, (select max(deptno) from dept))";
+    sql(sql).withExpand(false).ok();
+  }
+
+  /** Tests a MERGE whose INSERT values contain a sub-query correlated to the
+   * source, with sub-queries expanded into joins. */
+  @Test void testMergeInsertCorrelatedSubQueryExpanded() {
+    final String sql = "merge into empnullables e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when not matched then insert (empno, sal)\n"
+        + "values (t.empno,\n"
+        + "  (select d.deptno from dept d where d.deptno = t.deptno))";
+    sql(sql).withExpand(true).ok();
+  }
+
+  /** Tests a MERGE into a modifiable view. The insert row is built for the
+   * view, as the update part, and the modification is on the view; the MERGE
+   * is not rewritten to the underlying table, so the values of the columns
+   * constrained by the view (DEPTNO = 20) are not set. */
+  @Test void testMergeModifiableView() {
+    final String sql = "merge into EMP_MODIFIABLEVIEW e\n"
+        + "using (select * from emp) t\n"
+        + "on e.empno = t.empno\n"
+        + "when matched then update set sal = t.sal\n"
+        + "when not matched then insert (empno, ename, sal)\n"
+        + "values (t.empno, t.ename, t.sal)";
+    sql(sql).withExtendedTester().ok();
+  }
+
   @Test void testSelectView() {
     // translated condition: deptno = 20 and sal > 1000 and empno > 100
     final String sql = "select * from emp_20 where empno > 100";