Support uncorrelated subquery in UPDATE assignments (SET)

trinodb · Jun 14, 2021 · cab7620 · cab7620
1 parent 8270d00
commit cab7620
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java b/core/trino-main/src/main/java/io/trino/sql/planner/QueryPlanner.java
@@ -529,6 +529,7 @@ public UpdateNode plan(Update node)
             builder = filter(builder, node.getWhere().get(), node);
         }
 
+        builder = subqueryPlanner.handleSubqueries(builder, orderedColumnValues, analysis.getSubqueries(node));
         builder = builder.appendProjections(orderedColumnValues, symbolAllocator, idAllocator);
 
         PlanAndMappings planAndMappings = coerce(builder, orderedColumnValues, analysis, idAllocator, symbolAllocator, typeCoercion);

diff --git a/...o-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveTransactionalTable.java b/...o-product-tests/src/main/java/io/trino/tests/product/hive/TestHiveTransactionalTable.java
@@ -1395,6 +1395,35 @@ public void testAcidUpdateWithSubqueryPredicate()
         });
     }
 
+    @Test(groups = HIVE_TRANSACTIONAL, timeOut = TEST_TIMEOUT)
+    public void testAcidUpdateWithSubqueryAssignment()
+    {
+        // TODO support UPDATE with correlated subquery in assignment
+        withTemporaryTable("test_update_subquery", true, false, NONE, tableName -> {
+            onTrino().executeQuery(format("CREATE TABLE %s (column1 INT, column2 varchar) WITH (transactional = true)", tableName));
+            onTrino().executeQuery(format("INSERT INTO %s VALUES (1, 'x')", tableName));
+            onTrino().executeQuery(format("INSERT INTO %s VALUES (2, 'y')", tableName));
+
+            // SET with uncorrelated subquery
+            onTrino().executeQuery(format("UPDATE %s SET column2 = (SELECT max(name) FROM tpch.tiny.region)", tableName));
+            verifySelectForPrestoAndHive("SELECT * FROM " + tableName, "true", row(1, "MIDDLE EAST"), row(2, "MIDDLE EAST"));
+
+            withTemporaryTable("second_table", true, false, NONE, secondTable -> {
+                onTrino().executeQuery(format("CREATE TABLE %s WITH (transactional = true) AS TABLE tpch.tiny.region", secondTable));
+
+                // UPDATE while reading from another transactional table. Multiple transactional could interfere with ConnectorMetadata.beginQuery
+                onTrino().executeQuery(format("UPDATE %s SET column2 = (SELECT min(name) FROM %s)", tableName, secondTable));
+                // TODO (https://github.com/trinodb/trino/issues/8268) verifySelectForPrestoAndHive("SELECT * FROM " + tableName, "true", row(1, "AFRICA"), row(2, "AFRICA"));
+                verifySelect("onPresto", onTrino(), "SELECT * FROM " + tableName, "true", row(1, "AFRICA"), row(2, "AFRICA"));
+            });
+
+            // SET with correlated subquery
+            assertThat(() -> onTrino().executeQuery(format("UPDATE %s SET column2 = (SELECT name FROM tpch.tiny.region WHERE column1 = regionkey)", tableName)))
+                    // TODO (https://github.com/trinodb/trino/issues/3325) support correlated UPDATE
+                    .failsWithMessageMatching("\\Qjava.sql.SQLException: Query failed (#\\E\\S+\\Q): Invalid descendant for DeleteNode or UpdateNode: io.trino.sql.planner.plan.MarkDistinctNode");
+        });
+    }
+
     @Flaky(issue = ERROR_COMMITTING_WRITE_TO_HIVE_ISSUE, match = ERROR_COMMITTING_WRITE_TO_HIVE_MATCH)
     @Test(groups = HIVE_TRANSACTIONAL, timeOut = TEST_TIMEOUT)
     public void testInsertDeletUpdateWithPrestoAndHive()