Refactor coerce() method

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

Hide file tree

Showing 2 changed files with 18 additions and 17 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
@@ -28,6 +28,7 @@
 import io.prestosql.sql.analyzer.Analysis.GroupingSetAnalysis;
 import io.prestosql.sql.analyzer.Analysis.SelectExpression;
 import io.prestosql.sql.analyzer.FieldId;
+import io.prestosql.sql.analyzer.RelationType;
 import io.prestosql.sql.planner.plan.AggregationNode;
 import io.prestosql.sql.planner.plan.AggregationNode.Aggregation;
 import io.prestosql.sql.planner.plan.Assignments;
@@ -797,15 +798,15 @@ public static Expression coerceIfNecessary(Analysis analysis, Expression origina
                 analysis.isTypeOnlyCoercion(original));
     }
 
-    public static NodeAndMappings coerce(PlanNode node, List<Symbol> fields, List<Type> types, SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator)
+    public static NodeAndMappings coerce(RelationPlan plan, List<Type> types, SymbolAllocator symbolAllocator, PlanNodeIdAllocator idAllocator)
     {
-        checkArgument(fields.size() == types.size());
+        List<Symbol> visibleFields = visibleFields(plan);
+        checkArgument(visibleFields.size() == types.size());
 
         Assignments.Builder assignments = Assignments.builder();
-
         ImmutableList.Builder<Symbol> mappings = ImmutableList.builder();
         for (int i = 0; i < types.size(); i++) {
-            Symbol input = fields.get(i);
+            Symbol input = visibleFields.get(i);
             Type type = types.get(i);
 
             if (!symbolAllocator.getTypes().get(input).equals(type)) {
@@ -819,10 +820,20 @@ public static NodeAndMappings coerce(PlanNode node, List<Symbol> fields, List<Ty
             }
         }
 
-        ProjectNode coerced = new ProjectNode(idAllocator.getNextId(), node, assignments.build());
+        ProjectNode coerced = new ProjectNode(idAllocator.getNextId(), plan.getRoot(), assignments.build());
         return new NodeAndMappings(coerced, mappings.build());
     }
 
+    private static List<Symbol> visibleFields(RelationPlan subPlan)
+    {
+        RelationType descriptor = subPlan.getDescriptor();
+        return descriptor.getAllFields().stream()
+                .filter(field -> !field.isHidden())
+                .map(descriptor::indexOf)
+                .map(subPlan.getFieldMappings()::get)
+                .collect(toImmutableList());
+    }
+
     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
@@ -166,7 +166,7 @@ protected RelationPlan visitTable(Table node, Void context)
                     .map(Field::getType)
                     .collect(toImmutableList());
 
-            NodeAndMappings coerced = coerce(subPlan.getRoot(), visibleFields(subPlan), types, symbolAllocator, idAllocator);
+            NodeAndMappings coerced = coerce(subPlan, types, symbolAllocator, idAllocator);
 
             plan = new RelationPlan(coerced.getNode(), scope, coerced.getFields(), outerContext);
         }
@@ -884,7 +884,7 @@ private SetOperationPlan process(SetOperation node)
                         .collect(toImmutableList());
             }
 
-            NodeAndMappings planAndMappings = coerce(plan.getRoot(), visibleFields(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));
             }
@@ -906,16 +906,6 @@ private PlanNode distinct(PlanNode node)
                 Optional.empty());
     }
 
-    private static List<Symbol> visibleFields(RelationPlan subPlan)
-    {
-        RelationType descriptor = subPlan.getDescriptor();
-        return descriptor.getAllFields().stream()
-                .filter(field -> !field.isHidden())
-                .map(descriptor::indexOf)
-                .map(subPlan.getFieldMappings()::get)
-                .collect(toImmutableList());
-    }
-
     private static class SetOperationPlan
     {
         private final List<PlanNode> sources;