Skip to content

Commit

Permalink
Fix column lineage for UPDATE statements with subqueries
Browse files Browse the repository at this point in the history 
This commit fixes the issue where the query output columns were not setting
source columns for UPDATE statements setting a target column to a subquery
result.
  • Loading branch information
@lucasdlemos @kokosing
lucasdlemos authored and kokosing committed Aug 29, 2023
1 parent ceb3e8b commit 482f755
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 4 deletions.
20 changes: 16 additions & 4 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3315,14 +3315,24 @@ protected Scope visitUpdate(Update update, Optional<Scope> scope)

ImmutableList.Builder<ExpressionAnalysis> analysesBuilder = ImmutableList.builder();
ImmutableList.Builder<Type> expressionTypesBuilder = ImmutableList.builder();
ImmutableMap.Builder<String, Set<SourceColumn>> sourceColumnsByColumnNameBuilder = ImmutableMap.builder();
for (UpdateAssignment assignment : update.getAssignments()) {
String targetColumnName = assignment.getName().getValue();
Expression expression = assignment.getValue();
ExpressionAnalysis analysis = analyzeExpression(expression, tableScope);
analysesBuilder.add(analysis);
expressionTypesBuilder.add(analysis.getType(expression));
ExpressionAnalysis expressionAnalysis = analyzeExpression(expression, tableScope);
analysesBuilder.add(expressionAnalysis);
expressionTypesBuilder.add(expressionAnalysis.getType(expression));

Set<SourceColumn> sourceColumns = expressionAnalysis.getSubqueries().stream()
.map(query -> analyze(query.getNode(), tableScope))
.flatMap(subqueryScope -> subqueryScope.getRelationType().getVisibleFields().stream())
.flatMap(field -> analysis.getSourceColumns(field).stream())
.collect(toImmutableSet());
sourceColumnsByColumnNameBuilder.put(targetColumnName, sourceColumns);
}
List<ExpressionAnalysis> analyses = analysesBuilder.build();
List<Type> expressionTypes = expressionTypesBuilder.build();
Map<String, Set<SourceColumn>> sourceColumnsByColumnName = sourceColumnsByColumnNameBuilder.buildOrThrow();

List<Type> tableTypes = update.getAssignments().stream()
.map(assignment -> requireNonNull(columns.get(assignment.getName().getValue())))
Expand Down Expand Up @@ -3353,7 +3363,9 @@ protected Scope visitUpdate(Update update, Optional<Scope> scope)
tableName,
Optional.of(table),
Optional.of(updatedColumnSchemas.stream()
.map(column -> new OutputColumn(new Column(column.getName(), column.getType().toString()), ImmutableSet.of()))
.map(column -> new OutputColumn(
new Column(column.getName(), column.getType().toString()),
sourceColumnsByColumnName.getOrDefault(column.getName(), ImmutableSet.of())))
.collect(toImmutableList())));

createMergeAnalysis(table, handle, tableSchema, tableScope, tableScope, ImmutableList.of(updatedColumnHandles));
Expand Down
59 changes: 59 additions & 0 deletions testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1077,6 +1077,65 @@ public void testOutputColumnsForUpdatingSingleColumn()
.containsExactly(new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of()));
}

@Test
public void testOutputColumnsForUpdatingColumnWithSelectQuery()
throws Exception
{
QueryEvents queryEvents = runQueryAndWaitForEvents("UPDATE mock.default.table_for_output SET test_varchar = (SELECT name from nation LIMIT 1)").getQueryEvents();
QueryCompletedEvent event = queryEvents.getQueryCompletedEvent();
assertThat(event.getIoMetadata().getOutput().get().getColumns().get())
.containsExactly(new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "nation", "name"))));
}

@Test
public void testOutputColumnsForUpdatingColumnWithSelectQueryWithAliasedField()
throws Exception
{
QueryEvents queryEvents = runQueryAndWaitForEvents("UPDATE mock.default.table_for_output SET test_varchar = (SELECT name AS aliased_name from nation LIMIT 1)").getQueryEvents();
QueryCompletedEvent event = queryEvents.getQueryCompletedEvent();
assertThat(event.getIoMetadata().getOutput().get().getColumns().get())
.containsExactly(new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "nation", "name"))));
}

@Test
public void testOutputColumnsForUpdatingColumnsWithSelectQueries()
throws Exception
{
QueryEvents queryEvents = runQueryAndWaitForEvents("""
UPDATE mock.default.table_for_output SET test_varchar = (SELECT name AS aliased_name from nation LIMIT 1), test_bigint = (SELECT nationkey FROM nation LIMIT 1)
""").getQueryEvents();
QueryCompletedEvent event = queryEvents.getQueryCompletedEvent();
assertThat(event.getIoMetadata().getOutput().get().getColumns().get())
.containsExactlyInAnyOrder(
new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "nation", "name"))),
new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "nation", "nationkey"))));
}

@Test
public void testOutputColumnsForUpdatingColumnsWithSelectQueryAndRawValue()
throws Exception
{
QueryEvents queryEvents = runQueryAndWaitForEvents("""
UPDATE mock.default.table_for_output SET test_varchar = (SELECT name AS aliased_name from nation LIMIT 1), test_bigint = 1
""").getQueryEvents();
QueryCompletedEvent event = queryEvents.getQueryCompletedEvent();
assertThat(event.getIoMetadata().getOutput().get().getColumns().get())
.containsExactlyInAnyOrder(
new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "nation", "name"))),
new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of()));
}

@Test
public void testOutputColumnsForUpdatingColumnWithSelectQueryAndWhereClauseWithOuterColumn()
throws Exception
{
QueryEvents queryEvents = runQueryAndWaitForEvents("""
UPDATE mock.default.table_for_output SET test_varchar = (SELECT name from nation WHERE test_bigint = nationkey)""").getQueryEvents();
QueryCompletedEvent event = queryEvents.getQueryCompletedEvent();
assertThat(event.getIoMetadata().getOutput().get().getColumns().get())
.containsExactly(new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "nation", "name"))));
}

@Test
public void testCreateTable()
throws Exception
Expand Down

0 comments on commit 482f755

Please sign in to comment.