Add method for pruning invisible fields

trinodb · Jul 28, 2020 · b57236f · b57236f
1 parent 696c41a
commit b57236f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 8 deletions.
diff --git a/presto-main/src/main/java/io/prestosql/sql/planner/QueryPlanner.java b/presto-main/src/main/java/io/prestosql/sql/planner/QueryPlanner.java
@@ -834,6 +834,13 @@ private static List<Symbol> visibleFields(RelationPlan subPlan)
                 .collect(toImmutableList());
     }
 
+    public static NodeAndMappings pruneInvisibleFields(RelationPlan plan, PlanNodeIdAllocator idAllocator)
+    {
+        List<Symbol> visibleFields = visibleFields(plan);
+        ProjectNode pruned = new ProjectNode(idAllocator.getNextId(), plan.getRoot(), Assignments.identity(visibleFields));
+        return new NodeAndMappings(pruned, visibleFields);
+    }
+
     private PlanBuilder distinct(PlanBuilder subPlan, QuerySpecification node, List<Expression> expressions)
     {
         if (node.getSelect().isDistinct()) {

diff --git a/presto-main/src/main/java/io/prestosql/sql/planner/RelationPlanner.java b/presto-main/src/main/java/io/prestosql/sql/planner/RelationPlanner.java
@@ -94,6 +94,7 @@
 import static io.prestosql.sql.planner.PlanBuilder.newPlanBuilder;
 import static io.prestosql.sql.planner.QueryPlanner.coerce;
 import static io.prestosql.sql.planner.QueryPlanner.coerceIfNecessary;
+import static io.prestosql.sql.planner.QueryPlanner.pruneInvisibleFields;
 import static io.prestosql.sql.planner.plan.AggregationNode.singleGroupingSet;
 import static io.prestosql.sql.tree.BooleanLiteral.TRUE_LITERAL;
 import static io.prestosql.sql.tree.Join.Type.CROSS;
@@ -874,17 +875,16 @@ private SetOperationPlan process(SetOperation node)
         for (Relation child : node.getRelations()) {
             RelationPlan plan = process(child, null);
 
+            NodeAndMappings planAndMappings;
             List<Type> types = analysis.getRelationCoercion(child);
             if (types == null) {
-                // If there are no coercions, we still want to call the coerce method below to create a plan
-                // with just the visible fields
-                types = plan.getDescriptor()
-                        .getVisibleFields().stream()
-                        .map(Field::getType)
-                        .collect(toImmutableList());
+                // no coercion required, only prune invisible fields from child outputs
+                planAndMappings = pruneInvisibleFields(plan, idAllocator);
+            }
+            else {
+                // apply required coercion and prune invisible fields from child outputs
+                planAndMappings = coerce(plan, types, symbolAllocator, idAllocator);
             }
-
-            NodeAndMappings planAndMappings = coerce(plan, types, symbolAllocator, idAllocator);
             for (int i = 0; i < outputFields.getAllFields().size(); i++) {
                 symbolMapping.put(outputs.get(i), planAndMappings.getFields().get(i));
             }