Fix CTE reference in analyzer

prestodb · Apr 16, 2024 · 29fd544 · 29fd544
1 parent a38afd6
commit 29fd544
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 64 deletions.
diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java b/presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java
@@ -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);

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java
@@ -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);
                 session.getCteInformationCollector().addCTEReference(cteName, normalizedCteId, namedQuery.isFromView(), true);
                 subPlan = new RelationPlan(
                         new CteReferenceNode(getSourceLocation(node.getLocation()),

diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/SqlPlannerContext.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/SqlPlannerContext.java
@@ -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;
@@ -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;
         }
     }
 }
diff --git a/.../src/test/java/com/facebook/presto/sql/planner/optimizations/TestLogicalCteOptimizer.java b/.../src/test/java/com/facebook/presto/sql/planner/optimizations/TestLogicalCteOptimizer.java
@@ -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