[CALCITE-6570] Add SCALAR_QUERY to sourceExpressionList

apache · Sep 13, 2024 · 54216ca · 54216ca
1 parent 91fe118
commit 54216ca
Show file tree

Hide file tree

Showing 3 changed files with 31 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,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
@@ -3119,6 +3119,13 @@ void checkCorrelatedMapSubQuery(boolean expand) {
     sql(sql).ok();
   }
 
+  @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
@@ -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>