Skip to content

Track column lineage #7354

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Praveen2112 wants to merge 3 commits into from
Closed

Track column lineage #7354

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b74988c
Remove unused variable
Praveen2112 Mar 19, 2021
9c6e746
Introduce OriginColumnDetail
Praveen2112 Mar 19, 2021
ba5b4ef
Track lineage for each field.
Praveen2112 Mar 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/Analysis.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
package io.trino.sql.analyzer;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.HashMultimap;
import com.google.common.collect.HashMultiset;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
Expand All @@ -38,6 +37,7 @@
import io.trino.spi.security.Identity;
import io.trino.spi.security.ViewExpression;
import io.trino.spi.type.Type;
import io.trino.sql.analyzer.Field.OriginColumnDetail;
import io.trino.sql.tree.AllColumns;
import io.trino.sql.tree.ExistsPredicate;
import io.trino.sql.tree.Expression;
Expand Down Expand Up @@ -121,9 +121,6 @@ public class Analysis
// a map of users to the columns per table that they access
private final Map<AccessControlInfo, Map<QualifiedObjectName, Set<String>>> tableColumnReferences = new LinkedHashMap<>();

// Track referenced fields from source relation node
private final Multimap<NodeRef<? extends Node>, Field> referencedFields = HashMultimap.create();

private final Map<NodeRef<QuerySpecification>, List<FunctionCall>> aggregates = new LinkedHashMap<>();
private final Map<NodeRef<OrderBy>, List<Expression>> orderByAggregates = new LinkedHashMap<>();
private final Map<NodeRef<QuerySpecification>, GroupingSetAnalysis> groupingSets = new LinkedHashMap<>();
Expand Down Expand Up @@ -194,6 +191,7 @@ public class Analysis

// row id field for update/delete queries
private final Map<NodeRef<Table>, FieldReference> rowIdField = new LinkedHashMap<>();
private final Map<NodeRef<Expression>, List<OriginColumnDetail>> originColumnDetails = new LinkedHashMap<>();

public Analysis(@Nullable Statement root, Map<NodeRef<Parameter>, Expression> parameters, boolean isDescribe)
{
Expand Down Expand Up @@ -839,9 +837,9 @@ public void addEmptyColumnReferencesForTable(AccessControl accessControl, Identi
tableColumnReferences.computeIfAbsent(accessControlInfo, k -> new LinkedHashMap<>()).computeIfAbsent(table, k -> new HashSet<>());
}

public void addReferencedFields(Multimap<NodeRef<Node>, Field> references)
public void addReferencedFields(Expression expression, Multimap<NodeRef<Node>, Field> references)

@kasiafi kasiafi Mar 19, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

addColumnOriginDetails ?

{
referencedFields.putAll(references);
originColumnDetails.put(NodeRef.of(expression), references.values().stream().flatMap(x -> x.getOriginColumnDetails().stream()).collect(toImmutableList()));
}

public Map<AccessControlInfo, Map<QualifiedObjectName, Set<String>>> getTableColumnReferences()
Expand Down Expand Up @@ -885,6 +883,11 @@ public List<Expression> getRowFilters(Table node)
return rowFilters.getOrDefault(NodeRef.of(node), ImmutableList.of());
}

public List<OriginColumnDetail> getColumnOriginDetails(Expression expression)
{
return originColumnDetails.get(NodeRef.of(expression));
}

public boolean hasColumnMask(QualifiedObjectName table, String column, String identity)
{
return columnMaskScopes.contains(new ColumnMaskScopeEntry(table, column, identity));
Expand Down
6 changes: 2 additions & 4 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/ExpressionAnalyzer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -503,9 +503,7 @@ private Type handleResolvedField(Expression node, ResolvedField resolvedField, S
}
}

if (field.getOriginTable().isPresent() && field.getOriginColumnName().isPresent()) {
tableColumnReferences.put(field.getOriginTable().get(), field.getOriginColumnName().get());
}
field.getOriginColumnDetails().forEach(columnDetail -> tableColumnReferences.put(columnDetail.getTableName(), columnDetail.getColumnName()));

fieldId.getRelationId()
.getSourceNode()
Expand Down Expand Up @@ -1973,6 +1971,7 @@ public static ExpressionAnalysis analyzeExpression(
analyzer.analyze(expression, scope);

updateAnalysis(analysis, analyzer, session, accessControl);
analysis.addReferencedFields(expression, analyzer.getReferencedFields());

@kasiafi kasiafi Mar 19, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably referencedFields could be now simplified to Set<Field>. I don't think the mapping by source node is used.


return new ExpressionAnalysis(
analyzer.getExpressionTypes(),
Expand Down Expand Up @@ -2030,7 +2029,6 @@ private static void updateAnalysis(Analysis analysis, ExpressionAnalyzer analyze
analysis.addColumnReferences(analyzer.getColumnReferences());
analysis.addLambdaArgumentReferences(analyzer.getLambdaArgumentReferences());
analysis.addTableColumnReferences(accessControl, session.getIdentity(), analyzer.getTableColumnReferences());
analysis.addReferencedFields(analyzer.getReferencedFields());
}

public static ExpressionAnalyzer create(
Expand Down
60 changes: 38 additions & 22 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/Field.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,18 +13,19 @@
*/
package io.trino.sql.analyzer;

import com.google.common.collect.ImmutableList;
import io.trino.metadata.QualifiedObjectName;
import io.trino.spi.type.Type;
import io.trino.sql.tree.QualifiedName;

import java.util.List;
import java.util.Optional;

import static java.util.Objects.requireNonNull;

public class Field
{
private final Optional<QualifiedObjectName> originTable;
private final Optional<String> originColumnName;
private final List<OriginColumnDetail> originColumnDetails;

@kasiafi kasiafi Mar 19, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you explain how it is possible to have more than 1 element in this list?

@Praveen2112 Praveen2112 Mar 19, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Like if we have field for a expression like func(func2(col1, col2), col3) then we might need fetch the OriginColumnDetail for col1, col2 and col3

private final Optional<QualifiedName> relationAlias;
private final Optional<String> name;
private final Type type;
Expand All @@ -36,61 +37,54 @@ public static Field newUnqualified(String name, Type type)
requireNonNull(name, "name is null");
requireNonNull(type, "type is null");

return new Field(Optional.empty(), Optional.of(name), type, false, Optional.empty(), Optional.empty(), false);
return new Field(Optional.empty(), Optional.of(name), type, false, ImmutableList.of(), false);
}

public static Field newUnqualified(Optional<String> name, Type type)
{
requireNonNull(name, "name is null");
requireNonNull(type, "type is null");

return new Field(Optional.empty(), name, type, false, Optional.empty(), Optional.empty(), false);
return new Field(Optional.empty(), name, type, false, ImmutableList.of(), false);
}

public static Field newUnqualified(Optional<String> name, Type type, Optional<QualifiedObjectName> originTable, Optional<String> originColumn, boolean aliased)
public static Field newUnqualified(Optional<String> name, Type type, List<OriginColumnDetail> originColumnDetails, boolean aliased)
{
requireNonNull(name, "name is null");
requireNonNull(type, "type is null");
requireNonNull(originTable, "originTable is null");
requireNonNull(originColumnDetails, "originColumnDetails is null");

return new Field(Optional.empty(), name, type, false, originTable, originColumn, aliased);
return new Field(Optional.empty(), name, type, false, originColumnDetails, aliased);
}

public static Field newQualified(QualifiedName relationAlias, Optional<String> name, Type type, boolean hidden, Optional<QualifiedObjectName> originTable, Optional<String> originColumn, boolean aliased)
public static Field newQualified(QualifiedName relationAlias, Optional<String> name, Type type, boolean hidden, List<OriginColumnDetail> originColumnDetails, boolean aliased)
{
requireNonNull(relationAlias, "relationAlias is null");
requireNonNull(name, "name is null");
requireNonNull(type, "type is null");
requireNonNull(originTable, "originTable is null");
requireNonNull(originColumnDetails, "originColumnDetails is null");

return new Field(Optional.of(relationAlias), name, type, hidden, originTable, originColumn, aliased);
return new Field(Optional.of(relationAlias), name, type, hidden, originColumnDetails, aliased);
}

public Field(Optional<QualifiedName> relationAlias, Optional<String> name, Type type, boolean hidden, Optional<QualifiedObjectName> originTable, Optional<String> originColumnName, boolean aliased)
public Field(Optional<QualifiedName> relationAlias, Optional<String> name, Type type, boolean hidden, List<OriginColumnDetail> originColumnDetails, boolean aliased)
{
requireNonNull(relationAlias, "relationAlias is null");
requireNonNull(name, "name is null");
requireNonNull(type, "type is null");
requireNonNull(originTable, "originTable is null");
requireNonNull(originColumnName, "originColumnName is null");
requireNonNull(originColumnDetails, "originColumnDetails is null");

this.relationAlias = relationAlias;
this.name = name;
this.type = type;
this.hidden = hidden;
this.originTable = originTable;
this.originColumnName = originColumnName;
this.originColumnDetails = ImmutableList.copyOf(originColumnDetails);
this.aliased = aliased;
}

public Optional<QualifiedObjectName> getOriginTable()
public List<OriginColumnDetail> getOriginColumnDetails()
{
return originTable;
}

public Optional<String> getOriginColumnName()
{
return originColumnName;
return originColumnDetails;
}

public Optional<QualifiedName> getRelationAlias()
Expand Down Expand Up @@ -169,4 +163,26 @@ public String toString()

return result.toString();
}

public static class OriginColumnDetail

@kasiafi kasiafi Mar 19, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The details are distinct per expression, but if you want to reason about them further, e.g. collect all references from a query, this class will need equals().

{
private final QualifiedObjectName tableName;
private final String columnName;

public OriginColumnDetail(QualifiedObjectName tableName, String columnName)
{
this.tableName = requireNonNull(tableName, "tableName is null");
this.columnName = requireNonNull(columnName, "columnName is null");
}

public QualifiedObjectName getTableName()
{
return tableName;
}

public String getColumnName()
{
return columnName;
}
}
}
6 changes: 2 additions & 4 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/RelationType.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,8 +179,7 @@ public RelationType withAlias(String relationAlias, List<String> columnAliases)
columnAlias,
field.getType(),
field.isHidden(),
field.getOriginTable(),
field.getOriginColumnName(),
field.getOriginColumnDetails(),
field.isAliased()));
}
else if (!field.isHidden()) {
Expand All @@ -192,8 +191,7 @@ else if (!field.isHidden()) {
columnAlias,
field.getType(),
false,
field.getOriginTable(),
field.getOriginColumnName(),
field.getOriginColumnDetails(),
field.isAliased()));
}
}
Expand Down
53 changes: 20 additions & 33 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/StatementAnalyzer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
import io.trino.sql.analyzer.Analysis.ResolvedWindow;
import io.trino.sql.analyzer.Analysis.SelectExpression;
import io.trino.sql.analyzer.Analysis.UnnestAnalysis;
import io.trino.sql.analyzer.Field.OriginColumnDetail;
import io.trino.sql.analyzer.Scope.AsteriskedIdentifierChainBasis;
import io.trino.sql.parser.ParsingException;
import io.trino.sql.parser.SqlParser;
Expand Down Expand Up @@ -191,6 +192,7 @@
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Verify.verify;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableListMultimap.toImmutableListMultimap;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.collect.Iterables.getLast;
Expand Down Expand Up @@ -1253,8 +1255,7 @@ protected Scope visitTable(Table table, Optional<Scope> scope)
Optional.of(column.getName()),
column.getType(),
column.isHidden(),
Optional.of(name),
Optional.of(column.getName()),
ImmutableList.of(new OriginColumnDetail(name, column.getName())),
false);
fields.add(field);
ColumnHandle columnHandle = columnHandles.get(column.getName());
Expand Down Expand Up @@ -1363,8 +1364,7 @@ private Scope createScopeForCommonTableExpression(Table table, Optional<Scope> s
Optional.of(aliases.next().getValue()),
inputField.getType(),
false,
inputField.getOriginTable(),
inputField.getOriginColumnName(),
inputField.getOriginColumnDetails(),
inputField.isAliased()));
}
}
Expand All @@ -1380,8 +1380,7 @@ private Scope createScopeForCommonTableExpression(Table table, Optional<Scope> s
inputField.getName(),
inputField.getType(),
false,
inputField.getOriginTable(),
inputField.getOriginColumnName(),
inputField.getOriginColumnDetails(),
inputField.isAliased()));
}
}
Expand Down Expand Up @@ -1423,8 +1422,7 @@ private Scope createScopeForView(Table table, QualifiedObjectName name, Optional
Optional.of(column.getName()),
getViewColumnType(column, name, table),
false,
Optional.of(name),
Optional.of(column.getName()),
ImmutableList.of(new OriginColumnDetail(name, column.getName())),
false))
.collect(toImmutableList());

Expand Down Expand Up @@ -1477,8 +1475,7 @@ private Scope createScopeForMaterializedView(Table table, QualifiedObjectName na
Optional.of(column.getName()),
getViewColumnType(column, name, table),
false,
Optional.of(name),
Optional.of(column.getName()),
ImmutableList.of(new OriginColumnDetail(name, column.getName())),
false))
.collect(toImmutableList());

Expand Down Expand Up @@ -1597,7 +1594,7 @@ protected Scope visitQuerySpecification(QuerySpecification node, Optional<Scope>
GroupingSetAnalysis groupByAnalysis = analyzeGroupBy(node, sourceScope, outputExpressions);
analyzeHaving(node, sourceScope);

Scope outputScope = computeAndAssignOutputScope(node, scope, sourceScope);
Scope outputScope = computeAndAssignOutputScope(node, scope);

List<Expression> orderByExpressions = emptyList();
Optional<Scope> orderByScope = Optional.empty();
Expand Down Expand Up @@ -1737,8 +1734,7 @@ protected Scope visitSetOperation(SetOperation node, Optional<Scope> scope)
oldField.getName(),
outputFieldTypes[i],
oldField.isHidden(),
oldField.getOriginTable(),
oldField.getOriginColumnName(),
oldField.getOriginColumnDetails(),
oldField.isAliased());
}

Expand Down Expand Up @@ -2006,12 +2002,11 @@ private Scope analyzeJoinUsing(Join node, List<Identifier> columns, Optional<Sco

private void recordColumnAccess(Field field)
{
if (field.getOriginTable().isPresent() && field.getOriginColumnName().isPresent()) {
analysis.addTableColumnReferences(
accessControl,
session.getIdentity(),
ImmutableMultimap.of(field.getOriginTable().get(), field.getOriginColumnName().get()));
}
analysis.addTableColumnReferences(
accessControl,
session.getIdentity(),
field.getOriginColumnDetails().stream()
.collect(toImmutableListMultimap(OriginColumnDetail::getTableName, OriginColumnDetail::getColumnName)));
}

private boolean isLateralRelation(Relation node)
Expand Down Expand Up @@ -2474,7 +2469,7 @@ private boolean hasAggregates(QuerySpecification node)
return !aggregates.isEmpty();
}

private Scope computeAndAssignOutputScope(QuerySpecification node, Optional<Scope> scope, Scope sourceScope)
private Scope computeAndAssignOutputScope(QuerySpecification node, Optional<Scope> scope)
{
ImmutableList.Builder<Field> outputFields = ImmutableList.builder();

Expand All @@ -2495,7 +2490,7 @@ private Scope computeAndAssignOutputScope(QuerySpecification node, Optional<Scop
name = field.getName();
}

outputFields.add(Field.newUnqualified(name, field.getType(), field.getOriginTable(), field.getOriginColumnName(), false));
outputFields.add(Field.newUnqualified(name, field.getType(), field.getOriginColumnDetails(), false));
}
}
else if (item instanceof SingleColumn) {
Expand All @@ -2504,8 +2499,7 @@ else if (item instanceof SingleColumn) {
Expression expression = column.getExpression();
Optional<Identifier> field = column.getAlias();

Optional<QualifiedObjectName> originTable = Optional.empty();
Optional<String> originColumn = Optional.empty();
ImmutableList.Builder<OriginColumnDetail> builder = ImmutableList.builder();
QualifiedName name = null;

if (expression instanceof Identifier) {
Expand All @@ -2515,21 +2509,15 @@ else if (expression instanceof DereferenceExpression) {
name = DereferenceExpression.getQualifiedName((DereferenceExpression) expression);
}

if (name != null) {
List<Field> matchingFields = sourceScope.getRelationType().resolveFields(name);
if (!matchingFields.isEmpty()) {
originTable = matchingFields.get(0).getOriginTable();
originColumn = matchingFields.get(0).getOriginColumnName();
}
}
builder.addAll(analysis.getColumnOriginDetails(expression));
Comment thread
Praveen2112 marked this conversation as resolved.
Outdated

if (field.isEmpty()) {
if (name != null) {
field = Optional.of(getLast(name.getOriginalParts()));
}
}

outputFields.add(Field.newUnqualified(field.map(Identifier::getValue), analysis.getType(expression), originTable, originColumn, column.getAlias().isPresent())); // TODO don't use analysis as a side-channel. Use outputExpressions to look up the type
outputFields.add(Field.newUnqualified(field.map(Identifier::getValue), analysis.getType(expression), builder.build(), column.getAlias().isPresent())); // TODO don't use analysis as a side-channel. Use outputExpressions to look up the type
}
else {
throw new IllegalArgumentException("Unsupported SelectItem type: " + item.getClass().getName());
Expand Down Expand Up @@ -2673,8 +2661,7 @@ private void analyzeAllColumnsFromTable(
alias,
field.getType(),
false,
field.getOriginTable(),
field.getOriginColumnName(),
field.getOriginColumnDetails(),
!allColumns.getAliases().isEmpty() || field.isAliased()));

Type type = field.getType();
Expand Down
Loading