From f29d170541c3851cf00f0b7ecdfe9b30ec6eca86 Mon Sep 17 00:00:00 2001
From: Pavel Pereslegin <xxtern@gmail.com>
Date: Fri, 13 Sep 2024 16:50:05 +0400
Subject: [PATCH] [CALCITE-6570] Add SCALAR_QUERY to sourceExpressionList

---
 .../calcite/sql/validate/SqlValidatorImpl.java |  6 ++++++
 .../calcite/test/SqlToRelConverterTest.java    | 12 ++++++++++++
 .../calcite/test/SqlToRelConverterTest.xml     | 18 ++++++++++++++++++
 3 files changed, 36 insertions(+)

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
index 706c8e6be0e5..461a1ecec588 100644
--- 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,12 @@ private void registerQuery(
           SqlSelect.HAVING_OPERAND);
       registerSubQueries(selectScope2,
           SqlNonNullableAccessors.getSelectList(select));
+
+      if (enclosingNode.getKind() == SqlKind.UPDATE) {
+        registerSubQueries(selectScope2,
+            ((SqlUpdate) enclosingNode).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
index 47efce21c62d..df14fc016b81 100644
--- 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]
diff --git a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
index 5d6f6f2f0ff6..1939c2593a2b 100644
--- a/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
+++ b/core/src/test/resources/org/apache/calcite/test/SqlToRelConverterTest.xml
@@ -8933,6 +8933,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>