From 3444d97014bd61a41a43eb1b8e1026e3ce754a38 Mon Sep 17 00:00:00 2001 From: praveenkrishna Date: Thu, 16 Dec 2021 14:58:18 +0530 Subject: [PATCH] Capture column level dependencies for CTEs and sub-queries Previously they weren't captured if some transformation is applied on columns referenced from CTEs or Aliased relation. --- .../trino/sql/analyzer/StatementAnalyzer.java | 30 +++++++++++++--- .../execution/TestEventListenerBasic.java | 35 ++++++++++++++++++- 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java b/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java index 8a15b5c8f83e..8902360781f8 100644 --- a/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java +++ b/core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java @@ -1616,14 +1616,16 @@ private Scope createScopeForCommonTableExpression(Table table, Optional s for (int i = 0; i < queryDescriptor.getAllFieldCount(); i++) { Field inputField = queryDescriptor.getFieldByIndex(i); if (!inputField.isHidden()) { - fieldBuilder.add(Field.newQualified( + Field field = Field.newQualified( QualifiedName.of(table.getName().getSuffix()), Optional.of(aliases.next().getValue()), inputField.getType(), false, inputField.getOriginTable(), inputField.getOriginColumnName(), - inputField.isAliased())); + inputField.isAliased()); + fieldBuilder.add(field); + analysis.addSourceColumns(field, analysis.getSourceColumns(inputField)); } } fields = fieldBuilder.build(); @@ -1633,14 +1635,16 @@ private Scope createScopeForCommonTableExpression(Table table, Optional s for (int i = 0; i < queryDescriptor.getAllFieldCount(); i++) { Field inputField = queryDescriptor.getFieldByIndex(i); if (!inputField.isHidden()) { - fieldBuilder.add(Field.newQualified( + Field field = Field.newQualified( QualifiedName.of(table.getName().getSuffix()), inputField.getName(), inputField.getType(), false, inputField.getOriginTable(), inputField.getOriginColumnName(), - inputField.isAliased())); + inputField.isAliased()); + fieldBuilder.add(field); + analysis.addSourceColumns(field, analysis.getSourceColumns(inputField)); } } fields = fieldBuilder.build(); @@ -2027,14 +2031,27 @@ protected Scope visitAliasedRelation(AliasedRelation relation, Optional s } List aliases = null; + Collection inputFields = relationType.getAllFields(); if (relation.getColumnNames() != null) { aliases = relation.getColumnNames().stream() .map(Identifier::getValue) .collect(Collectors.toList()); + // hidden fields are not exposed when there are column aliases + inputFields = relationType.getVisibleFields(); } RelationType descriptor = relationType.withAlias(relation.getAlias().getValue(), aliases); + checkArgument(inputFields.size() == descriptor.getAllFieldCount(), + "Expected %s fields, got %s", + descriptor.getAllFieldCount(), + inputFields.size()); + + Streams.forEachPair( + descriptor.getAllFields().stream(), + inputFields.stream(), + (newField, field) -> analysis.addSourceColumns(newField, analysis.getSourceColumns(field))); + return createAndAssignScope(relation, scope, descriptor); } @@ -3919,6 +3936,11 @@ private Scope setAliases(Scope scope, Identifier tableName, List col RelationType oldDescriptor = scope.getRelationType(); validateColumnAliases(columnNames, oldDescriptor.getVisibleFieldCount()); RelationType newDescriptor = oldDescriptor.withAlias(tableName.getValue(), columnNames.stream().map(Identifier::getValue).collect(toImmutableList())); + + Streams.forEachPair( + oldDescriptor.getAllFields().stream(), + newDescriptor.getAllFields().stream(), + (newField, field) -> analysis.addSourceColumns(newField, analysis.getSourceColumns(field))); return scope.withRelationType(newDescriptor); } diff --git a/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java b/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java index 8e23e0b275e7..ca66c413e307 100644 --- a/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java +++ b/testing/trino-tests/src/test/java/io/trino/execution/TestEventListenerBasic.java @@ -779,7 +779,40 @@ public void testOutputColumnsWithClause() throws Exception { assertLineage( - "WITH w AS (SELECT * FROM orders) SELECT clerk AS test_varchar, orderkey AS test_bigint FROM w", + "WITH w AS (SELECT * FROM orders) SELECT lower(clerk) AS test_varchar, orderkey AS test_bigint FROM w", + ImmutableSet.of("tpch.tiny.orders"), + new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "clerk"))), + new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "orderkey")))); + } + + @Test + public void testOutputColumnsColumnAliasInWithClause() + throws Exception + { + assertLineage( + "WITH w(aliased_clerk, aliased_orderkey) AS (SELECT clerk, orderkey FROM orders) SELECT lower(aliased_clerk) AS test_varchar, aliased_orderkey AS test_bigint FROM w", + ImmutableSet.of("tpch.tiny.orders"), + new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "clerk"))), + new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "orderkey")))); + } + + @Test + public void testOutputColumnsWithAliasedRelation() + throws Exception + { + assertLineage( + "SELECT lower(clerk) AS test_varchar, orderkey AS test_bigint FROM (SELECT * FROM orders) w", + ImmutableSet.of("tpch.tiny.orders"), + new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "clerk"))), + new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "orderkey")))); + } + + @Test + public void testOutputColumnsWithColumnAliasInAliasedRelation() + throws Exception + { + assertLineage( + "SELECT lower(aliased_clerk) AS test_varchar, aliased_orderkey AS test_bigint FROM (SELECT clerk, orderkey FROM orders) w(aliased_clerk, aliased_orderkey)", ImmutableSet.of("tpch.tiny.orders"), new OutputColumnMetadata("test_varchar", VARCHAR_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "clerk"))), new OutputColumnMetadata("test_bigint", BIGINT_TYPE, ImmutableSet.of(new ColumnDetail("tpch", "tiny", "orders", "orderkey"))));