[CALCITE-6570] Add SCALAR_QUERY to sourceExpressionList

apache · Sep 16, 2024 · 5385399 · 5385399
1 parent 91fe118
commit 5385399
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 0 deletions.
diff --git a/core/src/main/java/org/apache/calcite/sql/validate/SqlValidatorImpl.java b/core/src/main/java/org/apache/calcite/sql/validate/SqlValidatorImpl.java
@@ -2941,6 +2941,19 @@ private void registerQuery(
           SqlSelect.HAVING_OPERAND);
       registerSubQueries(selectScope2,
           SqlNonNullableAccessors.getSelectList(select));
+
+      if (enclosingNode.getKind() == SqlKind.UPDATE) {
+        registerSubQueries(selectScope2,
+            ((SqlUpdate) enclosingNode).getSourceExpressionList());
+      } else if (enclosingNode.getKind() == SqlKind.MERGE) {
+        SqlUpdate updateCall = ((SqlMerge) enclosingNode).getUpdateCall();
+
+        if (updateCall != null) {
+          registerSubQueries(selectScope2,
+              updateCall.getSourceExpressionList());
+        }
+      }
+
       final SqlNodeList orderList = select.getOrderList();
       if (orderList != null) {
         // If the query is 'SELECT DISTINCT', restrict the columns

diff --git a/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java b/core/src/test/java/org/apache/calcite/test/SqlToRelConverterTest.java
@@ -3119,6 +3119,18 @@ void checkCorrelatedMapSubQuery(boolean expand) {
     sql(sql).ok();
   }
 
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6570">[CALCITE-6570]
+   * UPDATE with sub-query that requires type cast gives AssertionError</a>.
+   */
+  @Test void testUpdateSubQueryWithCast() {
+    final String sql = "update emp\n"
+        + "set empno = (\n"
+        + "  select cast(min(empno) as BIGINT) from emp as e where e.deptno = emp.deptno)";
+    sql(sql).ok();
+  }
+
   /**
    * Test case for
    * <a href="https://issues.apache.org/jira/browse/CALCITE-3229">[CALCITE-3229]
@@ -3213,6 +3225,20 @@ void checkCorrelatedMapSubQuery(boolean expand) {
     sql(sql).ok();
   }
 
+  /**
+   * Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6570">[CALCITE-6570]
+   * UPDATE with sub-query that requires type cast gives AssertionError</a>.
+   */
+  @Test void testMergeSubQueryWithCast() {
+    final String sql = "merge into emp t0\n" +
+        "using emp t1 ON t0.empno = t1.empno\n" +
+        "when matched then\n" +
+        "update set deptno = (select cast(deptno as BIGINT) from emp where deptno > 1)";
+
+    sql(sql).ok();
+  }
+
   @Test void testSelectView() {
     // translated condition: deptno = 20 and sal > 1000 and empno > 100
     final String sql = "select * from emp_20 where empno > 100";

diff --git a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
@@ -4708,6 +4708,28 @@ LogicalTableModify(table=[[CATALOG, SALES, EMPNULLABLES]], operation=[MERGE], up
         LogicalFilter(condition=[IS NULL($7)])
           LogicalTableScan(table=[[CATALOG, SALES, EMP]])
       LogicalTableScan(table=[[CATALOG, SALES, EMPNULLABLES]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testMergeSubQueryWithCast">
+    <Resource name="sql">
+      <![CDATA[merge into emp t0
+using emp t1 ON t0.empno = t1.empno
+when matched then
+update set deptno = (select cast(deptno as BIGINT) from emp where deptno > 1)]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalTableModify(table=[[CATALOG, SALES, EMP]], operation=[MERGE], updateColumnList=[[DEPTNO]], flattened=[true])
+  LogicalProject($f0=[$18])
+    LogicalJoin(condition=[true], joinType=[left])
+      LogicalJoin(condition=[=($9, $0)], joinType=[inner])
+        LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+        LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+      LogicalAggregate(group=[{}], agg#0=[SINGLE_VALUE($0)])
+        LogicalProject(EXPR$0=[CAST($7):BIGINT NOT NULL])
+          LogicalFilter(condition=[>($7, 1)])
+            LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>
@@ -8933,6 +8955,24 @@ LogicalTableModify(table=[[CATALOG, SALES, EMP]], operation=[UPDATE], updateColu
       LogicalAggregate(group=[{0}], EXPR$0=[MIN($1)])
         LogicalProject(DEPTNO=[$7], EMPNO=[$0])
           LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+]]>
+    </Resource>
+  </TestCase>
+  <TestCase name="testUpdateSubQueryWithCast">
+    <Resource name="sql">
+      <![CDATA[update emp
+set empno = (
+  select cast(min(empno) as BIGINT) from emp as e where e.deptno = emp.deptno)]]>
+    </Resource>
+    <Resource name="plan">
+      <![CDATA[
+LogicalTableModify(table=[[CATALOG, SALES, EMP]], operation=[UPDATE], updateColumnList=[[EMPNO]], sourceExpressionList=[[CAST($0):INTEGER]], flattened=[true])
+  LogicalProject(EMPNO=[$0], ENAME=[$1], JOB=[$2], MGR=[$3], HIREDATE=[$4], SAL=[$5], COMM=[$6], DEPTNO=[$7], SLACKER=[$8], EXPR$0=[CAST($10):BIGINT])
+    LogicalJoin(condition=[=($7, $9)], joinType=[left])
+      LogicalTableScan(table=[[CATALOG, SALES, EMP]])
+      LogicalAggregate(group=[{0}], agg#0=[MIN($1)])
+        LogicalProject(DEPTNO=[$7], EMPNO=[$0])
+          LogicalTableScan(table=[[CATALOG, SALES, EMP]])
 ]]>
     </Resource>
   </TestCase>