Skip to content

Commit

Permalink
Avoid re-resolving column names after analysis
Browse files Browse the repository at this point in the history 
Record the fields directly and extract the column name from them when needed.
  • Loading branch information
@martint
martint committed Mar 7, 2025
1 parent 789eb18 commit ef453bb
Show file tree
Hide file tree
Showing 4 changed files with 31 additions and 33 deletions.
33 changes: 7 additions & 26 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/Analysis.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,6 @@
import java.util.Optional;
import java.util.OptionalLong;
import java.util.Set;
import java.util.stream.Stream;

import static com.google.common.base.MoreObjects.toStringHelper;
import static com.google.common.base.Preconditions.checkArgument;
Expand All @@ -119,7 +118,6 @@
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.collect.MoreCollectors.toOptional;
import static io.trino.sql.analyzer.QueryType.DESCRIBE;
import static io.trino.sql.analyzer.QueryType.EXPLAIN;
import static java.lang.Boolean.FALSE;
Expand Down Expand Up @@ -156,7 +154,7 @@ public class Analysis
private final Map<NodeRef<Expression>, ResolvedField> columnReferences = new LinkedHashMap<>();

// a map of users to the columns per table that they access
private final Map<AccessControlInfo, Map<QualifiedObjectName, Set<String>>> tableColumnReferences = new LinkedHashMap<>();
private final Map<AccessControlInfo, Map<QualifiedObjectName, Set<Field>>> tableColumnReferences = new LinkedHashMap<>();

// Record fields prefixed with labels in row pattern recognition context
private final Map<NodeRef<Expression>, Optional<String>> labels = new LinkedHashMap<>();
Expand Down Expand Up @@ -961,10 +959,10 @@ public UnnestAnalysis getUnnest(Unnest node)
return unnestAnalysis.get(NodeRef.of(node));
}

public void addTableColumnReferences(AccessControl accessControl, Identity identity, Multimap<QualifiedObjectName, String> tableColumnMap)
public void addTableColumnReferences(AccessControl accessControl, Identity identity, Multimap<QualifiedObjectName, Field> tableColumnMap)
{
AccessControlInfo accessControlInfo = new AccessControlInfo(accessControl, identity);
Map<QualifiedObjectName, Set<String>> references = tableColumnReferences.computeIfAbsent(accessControlInfo, k -> new LinkedHashMap<>());
Map<QualifiedObjectName, Set<Field>> references = tableColumnReferences.computeIfAbsent(accessControlInfo, k -> new LinkedHashMap<>());
tableColumnMap.asMap()
.forEach((key, value) -> references.computeIfAbsent(key, k -> new HashSet<>()).addAll(value));
}
Expand Down Expand Up @@ -1092,7 +1090,7 @@ public JsonTableAnalysis getJsonTableAnalysis(JsonTable jsonTable)
return jsonTableAnalyses.get(NodeRef.of(jsonTable));
}

public Map<AccessControlInfo, Map<QualifiedObjectName, Set<String>>> getTableColumnReferences()
public Map<AccessControlInfo, Map<QualifiedObjectName, Set<Field>>> getTableColumnReferences()
{
return tableColumnReferences;
}
Expand Down Expand Up @@ -1189,9 +1187,9 @@ public List<TableInfo> getReferencedTables()
.filter(Objects::nonNull)
.flatMap(Collection::stream)
.distinct()
.map(fieldName -> new ColumnInfo(
fieldName,
resolveColumnMask(table.getNode().getName(), fieldName, columnMasks.getOrDefault(table, ImmutableMap.of()))
.map(field -> new ColumnInfo(
field.getOriginColumnName().get(),
Optional.ofNullable(columnMasks.getOrDefault(table, ImmutableMap.of()).get(field))
.map(Expression::toString)))
.collect(toImmutableList());

Expand All @@ -1212,23 +1210,6 @@ public List<TableInfo> getReferencedTables()
.collect(toImmutableList());
}

private static Optional<Expression> resolveColumnMask(QualifiedName tableName, String fieldName, Map<Field, Expression> expressions)
{
QualifiedName qualifiedFieldName = concatIdentifier(tableName, fieldName);
return expressions.entrySet().stream()
.filter(fieldExpression -> fieldExpression.getKey().canResolve(qualifiedFieldName))
.collect(toOptional())
.map(Map.Entry::getValue);
}

private static QualifiedName concatIdentifier(QualifiedName tableName, String fieldName)
{
return QualifiedName.of(Stream.concat(
tableName.getOriginalParts().stream(),
Stream.of(new Identifier(fieldName)))
.collect(toImmutableList()));
}

public List<RoutineInfo> getRoutines()
{
return resolvedFunctions.values().stream()
Expand Down
7 changes: 6 additions & 1 deletion core/trino-main/src/main/java/io/trino/sql/analyzer/Analyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import static io.trino.spi.StandardErrorCode.EXPRESSION_NOT_SCALAR;
import static io.trino.sql.analyzer.ExpressionTreeUtils.extractAggregateFunctions;
Expand Down Expand Up @@ -105,7 +107,10 @@ public Analysis analyze(Statement statement, QueryType queryType)
accessControlInfo.getAccessControl().checkCanSelectFromColumns(
accessControlInfo.getSecurityContext(session.getRequiredTransactionId(), session.getQueryId(), session.getStart()),
tableName,
columns)));
columns.stream()
.map(Field::getOriginColumnName)
.map(Optional::get)
.collect(Collectors.toSet()))));
}

return analysis;
Expand Down
6 changes: 3 additions & 3 deletions core/trino-main/src/main/java/io/trino/sql/analyzer/ExpressionAnalyzer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -328,7 +328,7 @@ public class ExpressionAnalyzer
// For lambda argument references, maps each QualifiedNameReference to the referenced LambdaArgumentDeclaration
private final Map<NodeRef<Identifier>, LambdaArgumentDeclaration> lambdaArgumentReferences = new LinkedHashMap<>();
private final Set<NodeRef<FunctionCall>> windowFunctions = new LinkedHashSet<>();
private final Multimap<QualifiedObjectName, String> tableColumnReferences = HashMultimap.create();
private final Multimap<QualifiedObjectName, Field> tableColumnReferences = HashMultimap.create();

// Track referenced fields from source relation node
private final Multimap<NodeRef<Node>, Field> referencedFields = HashMultimap.create();
Expand Down Expand Up @@ -578,7 +578,7 @@ public Set<NodeRef<FunctionCall>> getWindowFunctions()
return unmodifiableSet(windowFunctions);
}

public Multimap<QualifiedObjectName, String> getTableColumnReferences()
public Multimap<QualifiedObjectName, Field> getTableColumnReferences()
{
return tableColumnReferences;
}
Expand Down Expand Up @@ -766,7 +766,7 @@ private Type handleResolvedField(Expression node, ResolvedField resolvedField, C
}

if (field.getOriginTable().isPresent() && field.getOriginColumnName().isPresent()) {
tableColumnReferences.put(field.getOriginTable().get(), field.getOriginColumnName().get());
tableColumnReferences.put(field.getOriginTable().get(), field);
}

sourceFields.add(field);
Expand Down
18 changes: 15 additions & 3 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 @@ -886,12 +886,24 @@ protected Scope visitAnalyze(Analyze node, Optional<Scope> scope)

analysis.setAnalyzeMetadata(metadata.getStatisticsCollectionMetadata(session, tableHandle, analyzeProperties));

Set<Field> fields = metadata.getTableSchema(session, tableHandle)
.columns().stream()
.map(column -> Field.newQualified(
node.getTableName(),
Optional.of(column.getName()),
column.getType(),
column.isHidden(),
Optional.of(tableName),
Optional.of(column.getName()),
false))
.collect(Collectors.toSet());

// user must have read and insert permission in order to analyze stats of a table
analysis.addTableColumnReferences(
accessControl,
session.getIdentity(),
ImmutableMultimap.<QualifiedObjectName, String>builder()
.putAll(tableName, metadata.getColumnHandles(session, tableHandle).keySet())
ImmutableMultimap.<QualifiedObjectName, Field>builder()
.putAll(tableName, fields)
.build());
try {
accessControl.checkCanInsertIntoTable(session.toSecurityContext(), tableName);
Expand Down Expand Up @@ -3856,7 +3868,7 @@ private void recordColumnAccess(Field field)
analysis.addTableColumnReferences(
accessControl,
session.getIdentity(),
ImmutableMultimap.of(field.getOriginTable().get(), field.getOriginColumnName().get()));
ImmutableMultimap.of(field.getOriginTable().get(), field));
}
}

Expand Down

0 comments on commit ef453bb

Please sign in to comment.