Use Proper query scope in CTE Materialization

presto-hive/src/test/java/com/facebook/presto/hive/TestCteExecution.java

@@ -167,6 +167,28 @@ public void testRefinedCtesOutsideScope()
                 queryRunner.execute(getSession(), testQuery));
     }
 
+    @Test
+    public void testRedefinedCteWithSameDefinitionDifferentBase()
+    {
+        String testQuery = "SELECT (with test_base AS (SELECT colB FROM (VALUES (1)) AS TempTable(colB)), \n" +
+                "test_cte as (  SELECT colB FROM test_base)\n" +
+                "SELECT * FROM test_cte\n" +
+                "),\n" +
+                "(WITH test_base AS (\n" +
+                "    SELECT text_column\n" +
+                "    FROM (VALUES ('Some Text', 9)) AS t (text_column, number_column)\n" +
+                "), \n" +
+                "test_cte AS (\n" +
+                "    SELECT * FROM test_base\n" +
+                ")\n" +
+                "SELECT  CONCAT(text_column , 'XYZ') FROM test_cte\n" +
+                ")\n";
+        QueryRunner queryRunner = getQueryRunner();
+        compareResults(
+                queryRunner.execute(getMaterializedSession(), testQuery),
+                queryRunner.execute(getSession(), testQuery));
+    }
+
     @Test
     public void testComplexRefinedCtesOutsideScope()
     {

presto-main/src/main/java/com/facebook/presto/sql/planner/RelationPlanner.java

@@ -184,15 +184,13 @@ protected RelationPlan visitTable(Table node, SqlPlannerContext context)
             if (namedQuery.isFromView()) {
                 cteName = createQualifiedObjectName(session, node, node.getName()).toString();
             }
-            context.getNestedCteStack().push(cteName, namedQuery.getQuery());
             RelationPlan subPlan = process(namedQuery.getQuery(), context);
-            context.getNestedCteStack().pop(namedQuery.getQuery());
             boolean shouldBeMaterialized = getCteMaterializationStrategy(session).equals(ALL) && isCteMaterializable(subPlan.getRoot().getOutputVariables());
             session.getCteInformationCollector().addCTEReference(cteName, namedQuery.isFromView(), shouldBeMaterialized);
             if (shouldBeMaterialized) {
                 subPlan = new RelationPlan(
                         new CteReferenceNode(getSourceLocation(node.getLocation()),
-                                idAllocator.getNextId(), subPlan.getRoot(), context.getNestedCteStack().getRawPath(cteName)),
+                                idAllocator.getNextId(), subPlan.getRoot(), context.getCteInfo().normalize(analysis, namedQuery.getQuery(), cteName)),
                         subPlan.getScope(),
                         subPlan.getFieldMappings());
             }

presto-main/src/main/java/com/facebook/presto/sql/planner/SqlPlannerContext.java

@@ -15,13 +15,17 @@
 
 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.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.Stack;
+import java.util.TreeSet;
 
 import static com.facebook.presto.SystemSessionProperties.getMaxLeafNodesInPlan;
 import static com.facebook.presto.SystemSessionProperties.isLeafNodeLimitEnabled;
@@ -34,18 +38,18 @@ public class SqlPlannerContext
     private int leafNodesInLogicalPlan;
     private final SqlToRowExpressionTranslator.Context translatorContext;
 
-    private final NestedCteStack nestedCteStack;
+    private final CteInfo cteInfo;
 
     public SqlPlannerContext(int leafNodesInLogicalPlan)
     {
         this.leafNodesInLogicalPlan = leafNodesInLogicalPlan;
         this.translatorContext = new SqlToRowExpressionTranslator.Context();
-        this.nestedCteStack = new NestedCteStack();
+        this.cteInfo = new CteInfo();
     }
 
-    public NestedCteStack getNestedCteStack()
+    public CteInfo getCteInfo()
     {
-        return nestedCteStack;
+        return cteInfo;
     }
 
     public SqlToRowExpressionTranslator.Context getTranslatorContext()
@@ -64,57 +68,69 @@ public void incrementLeafNodes(Session session)
         }
     }
 
-    public class NestedCteStack
+    public class CteInfo
     {
         @VisibleForTesting
         public static final String delimiter = "_*%$_";
-        private final Stack<String> cteStack;
-        private final Map<String, String> rawCtePathMap;
+        // never decreases
+        private int currentQueryScopeId;
 
-        public NestedCteStack()
-        {
-            this.cteStack = new Stack<>();
-            this.rawCtePathMap = new HashMap<>();
-        }
+        // 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<>();
 
-        public void push(String cteName, Query query)
+        public String normalize(Analysis analysis, Query query, String cteName)
         {
-            this.cteStack.push(cteName);
-            if (query.getWith().isPresent()) {
-                // All ctes defined in this context should have their paths updated
-                query.getWith().get().getQueries().forEach(with -> this.addNestedCte(with.getName().toString()));
+            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;
         }
 
-        public void pop(Query query)
+        private class QueryReferenceCollector
+                extends DefaultTraversalVisitor<Void, QueryReferenceCollectorContext>
         {
-            this.cteStack.pop();
-            if (query.getWith().isPresent()) {
-                query.getWith().get().getQueries().forEach(with -> this.removeNestedCte(with.getName().toString()));
+            private final Analysis analysis;
+
+            public QueryReferenceCollector(Analysis analysis)
+            {
+                this.analysis = analysis;
             }
-        }
 
-        public String getRawPath(String cteName)
-        {
-            if (!this.rawCtePathMap.containsKey(cteName)) {
-                return cteName;
+            @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;
             }
-            return this.rawCtePathMap.get(cteName);
         }
 
-        private void addNestedCte(String cteName)
+        private class QueryReferenceCollectorContext
         {
-            this.rawCtePathMap.put(cteName, getCurrentRelativeCtePath() + delimiter + cteName);
-        }
+            private final TreeSet<Query> referencedQuerySet;
 
-        private void removeNestedCte(String cteName)
-        {
-            this.rawCtePathMap.remove(cteName);
-        }
+            public QueryReferenceCollectorContext()
+            {
+                this.referencedQuerySet = new TreeSet<>(Comparator.comparingInt(Query::hashCode));
+            }
 
-        public String getCurrentRelativeCtePath()
-        {
-            return String.join(delimiter, cteStack);
+            public void addQuery(Query ref)
+            {
+                this.referencedQuerySet.add(ref);
+            }
+
+            public TreeSet<Query> getReferencedQuerySet()
+            {
+                return referencedQuerySet;
+            }
         }
     }
 }