Skip to content

Fix CTE reference in analyzer #22515

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

Merged
feilong-liu merged 1 commit into from
Apr 16, 2024
Merged

Fix CTE reference in analyzer #22515

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
14 changes: 14 additions & 0 deletions presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1153,6 +1153,20 @@ public void testWrittenIntemediateByteLimit()
assertQueryFails(session, testQuery, "Query has exceeded WrittenIntermediate Limit of 0MB.*");
}

@Test
public void testNestedCteWithSameName()
{
String testQuery = "with t1 as ( select orderkey k from orders where orderkey > 5), t2 as ( select orderkey k from orders where orderkey < 10 ), t3 as " +
"( select t1.k, t2.k from t1 left join t2 on t1.k=t2.k ), t4 as ( with t2 as ( select orderkey k from orders where orderkey > 5 ), " +
"t1 as ( select orderkey k from orders where orderkey < 10 ), t3 as ( select t1.k, t2.k from t1 left join t2 on t1.k=t2.k ) select * from t3 ) " +
"select * from t3 except select * from t4";
QueryRunner queryRunner = getQueryRunner();
compareResults(queryRunner.execute(getMaterializedSession(),
testQuery),
queryRunner.execute(getSession(),
testQuery));
}

private void compareResults(MaterializedResult actual, MaterializedResult expected)
{
compareResults(actual, expected, false);
Expand Down
2 changes: 1 addition & 1 deletion presto-main/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,7 @@ protected RelationPlan visitTable(Table node, SqlPlannerContext context)
}
else {
// cte considered for materialization
String normalizedCteId = context.getCteInfo().normalize(analysis, namedQuery.getQuery(), cteName);
String normalizedCteId = context.getCteInfo().normalize(NodeRef.of(namedQuery.getQuery()), cteName);
Copy link
Copy Markdown
Contributor Author

@feilong-liu feilong-liu Apr 13, 2024

Choose a reason for hiding this comment

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

Use NodeRef, because we want to use == rather than equals comparison.

session.getCteInformationCollector().addCTEReference(cteName, normalizedCteId, namedQuery.isFromView(), true);
subPlan = new RelationPlan(
new CteReferenceNode(getSourceLocation(node.getLocation()),
Expand Down
69 changes: 10 additions & 59 deletions presto-main/src/main/java/com/facebook/presto/sql/planner/SqlPlannerContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,13 @@

import com.facebook.presto.Session;
import com.facebook.presto.spi.PrestoException;
import com.facebook.presto.sql.analyzer.Analysis;
import com.facebook.presto.sql.relational.SqlToRowExpressionTranslator;
import com.facebook.presto.sql.tree.DefaultTraversalVisitor;
import com.facebook.presto.sql.tree.NodeRef;
import com.facebook.presto.sql.tree.Query;
import com.facebook.presto.sql.tree.Table;
import com.google.common.annotations.VisibleForTesting;

import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeSet;

import static com.facebook.presto.SystemSessionProperties.getMaxLeafNodesInPlan;
import static com.facebook.presto.SystemSessionProperties.isLeafNodeLimitEnabled;
Expand Down Expand Up @@ -73,64 +69,19 @@ public class CteInfo
@VisibleForTesting
public static final String delimiter = "_*%$_";
// never decreases
private int currentQueryScopeId;
private int prefix;

// Maps a set of Query objects, including the parent query statement and all its referenced statements,
// to a unique scope identifier. Each set of related queries shares the same scope.
Map<TreeSet<Query>, String> queryNodeScopeIdMap = new HashMap<>();
// Map a cte Query to a unique ID, which will be used in CTE reference node to identify the same CTE
private final Map<NodeRef<Query>, String> cteQueryUniqueIdMap = new HashMap<>();

public String normalize(Analysis analysis, Query query, String cteName)
public String normalize(NodeRef<Query> queryNodeRef, String cteName)
{
QueryReferenceCollectorContext context = new QueryReferenceCollectorContext();
context.getReferencedQuerySet().add(query);
query.accept(new QueryReferenceCollector(analysis), context);
TreeSet<Query> normalizedKey = context.getReferencedQuerySet();
if (!queryNodeScopeIdMap.containsKey(normalizedKey)) {
queryNodeScopeIdMap.put(normalizedKey, String.valueOf(currentQueryScopeId++));
}
return queryNodeScopeIdMap.get(normalizedKey) + delimiter + cteName;
}

private class QueryReferenceCollector
extends DefaultTraversalVisitor<Void, QueryReferenceCollectorContext>
{
private final Analysis analysis;

public QueryReferenceCollector(Analysis analysis)
{
this.analysis = analysis;
}

@Override
protected Void visitTable(Table node, QueryReferenceCollectorContext context)
{
Analysis.NamedQuery namedQuery = analysis.getNamedQuery(node);
if (namedQuery != null) {
context.addQuery(namedQuery.getQuery());
process(namedQuery.getQuery(), context);
}
return null;
}
}

private class QueryReferenceCollectorContext
{
private final TreeSet<Query> referencedQuerySet;

public QueryReferenceCollectorContext()
{
this.referencedQuerySet = new TreeSet<>(Comparator.comparingInt(Query::hashCode));
}

public void addQuery(Query ref)
{
this.referencedQuerySet.add(ref);
}

public TreeSet<Query> getReferencedQuerySet()
{
return referencedQuerySet;
if (cteQueryUniqueIdMap.containsKey(queryNodeRef)) {
return cteQueryUniqueIdMap.get(queryNodeRef) + delimiter + cteName;
}
String identityString = String.valueOf(prefix++);
cteQueryUniqueIdMap.put(queryNodeRef, identityString);
return identityString + delimiter + cteName;
}
}
}
8 changes: 4 additions & 4 deletions .../src/test/java/com/facebook/presto/sql/planner/optimizations/TestLogicalCteOptimizer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -624,13 +624,13 @@ public void testHeuristicMaterializationWithDeepNestedCteUsage3()
"SELECT * FROM a\n",
anyTree(
sequence(
cteProducer(addQueryScopeDelimiter("a", 1), anyTree(tableScan("orders"))),
cteProducer(addQueryScopeDelimiter("a", 2), anyTree(tableScan("orders"))),
anyTree(PlanMatchPattern.union(
PlanMatchPattern.union(
PlanMatchPattern.union(
anyTree(tableScan("orders")), anyTree(cteConsumer(addQueryScopeDelimiter("a", 1)))),
anyTree(cteConsumer(addQueryScopeDelimiter("a", 1)))),
anyTree(cteConsumer(addQueryScopeDelimiter("a", 1))))))));
anyTree(tableScan("orders")), anyTree(cteConsumer(addQueryScopeDelimiter("a", 2)))),
anyTree(cteConsumer(addQueryScopeDelimiter("a", 2)))),
anyTree(cteConsumer(addQueryScopeDelimiter("a", 2))))))));
}

@Test
Expand Down