Final analyzer/planner/optimizer changes for tvf except LocalExecutionPlanner and ExcludeColumns optimizer rule.

Co-authored-by: kasiafi <[email protected]>
Co-authored-by: Xin Zhang <[email protected]>

Author: 3 people

presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Field.java

@@ -86,14 +86,6 @@ public Field(Optional<NodeLocation> nodeLocation, Optional<QualifiedName> relati
         this.aliased = aliased;
     }
 
-    public static Field newUnqualified(Optional<String> name, Type type)
-    {
-        requireNonNull(name, "name is null");
-        requireNonNull(type, "type is null");
-
-        return new Field(Optional.empty(), Optional.empty(), name, type, false, Optional.empty(), Optional.empty(), false);
-    }
-
     public Optional<NodeLocation> getNodeLocation()
     {
         return nodeLocation;

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java

@@ -1523,7 +1523,7 @@ private void verifyRequiredColumns(TableFunctionInvocation node, Map<String, Lis
                 // the scope is recorded, because table arguments are already analyzed
                 Scope inputScope = analysis.getScope(tableArgumentsByName.get(name).getRelation());
                 columns.stream()
-                        .filter(column -> column < 0 || column >= inputScope.getRelationType().getAllFieldCount()) // hidden columns can be required as well as visible columns
+                        .filter(column -> column < 0 || column >= inputScope.getRelationType().getVisibleFieldCount())
                         .findFirst()
                         .ifPresent(column -> {
                             throw new SemanticException(TABLE_FUNCTION_IMPLEMENTATION_ERROR, "Invalid index: %s of required column from table argument %s", column, name);

presto-main-base/src/main/java/com/facebook/presto/sql/planner/BasePlanFragmenter.java

@@ -47,6 +47,8 @@
 import com.facebook.presto.sql.planner.plan.SequenceNode;
 import com.facebook.presto.sql.planner.plan.SimplePlanRewriter;
 import com.facebook.presto.sql.planner.plan.StatisticsWriterNode;
+import com.facebook.presto.sql.planner.plan.TableFunctionNode;
+import com.facebook.presto.sql.planner.plan.TableFunctionProcessorNode;
 import com.facebook.presto.sql.planner.sanity.PlanChecker;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
@@ -270,6 +272,22 @@ public PlanNode visitValues(ValuesNode node, RewriteContext<FragmentProperties>
         return context.defaultRewrite(node, context.get());
     }
 
+    @Override
+    public PlanNode visitTableFunction(TableFunctionNode node, RewriteContext<FragmentProperties> context)
+    {
+        throw new IllegalStateException(format("Unexpected node: TableFunctionNode (%s)", node.getName()));
+    }
+
+    @Override
+    public PlanNode visitTableFunctionProcessor(TableFunctionProcessorNode node, RewriteContext<FragmentProperties> context)
+    {
+        if (!node.getSource().isPresent()) {
+            // context is mutable. The leaf node should set the PartitioningHandle.
+            context.get().addSourceDistribution(node.getId(), SOURCE_DISTRIBUTION, metadata, session);
+        }
+        return context.defaultRewrite(node, context.get());
+    }
+
     @Override
     public PlanNode visitExchange(ExchangeNode exchange, RewriteContext<FragmentProperties> context)
     {

presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java

@@ -50,6 +50,7 @@
 import com.facebook.presto.sql.planner.iterative.rule.ImplementBernoulliSampleAsFilter;
 import com.facebook.presto.sql.planner.iterative.rule.ImplementFilteredAggregations;
 import com.facebook.presto.sql.planner.iterative.rule.ImplementOffset;
+import com.facebook.presto.sql.planner.iterative.rule.ImplementTableFunctionSource;
 import com.facebook.presto.sql.planner.iterative.rule.InlineProjections;
 import com.facebook.presto.sql.planner.iterative.rule.InlineProjectionsOnValues;
 import com.facebook.presto.sql.planner.iterative.rule.InlineSqlFunctions;
@@ -375,6 +376,14 @@ public PlanOptimizers(
         PlanOptimizer predicatePushDown = new StatsRecordingPlanOptimizer(optimizerStats, new PredicatePushDown(metadata, sqlParser, expressionOptimizerManager, featuresConfig.isNativeExecutionEnabled()));
         PlanOptimizer prefilterForLimitingAggregation = new StatsRecordingPlanOptimizer(optimizerStats, new PrefilterForLimitingAggregation(metadata, statsCalculator));
 
+        builder.add(
+                new IterativeOptimizer(
+                        metadata,
+                        ruleStats,
+                        statsCalculator,
+                        costCalculator,
+                        ImmutableSet.of(new RewriteTableFunctionToTableScan(metadata))));
+
         builder.add(
                 new IterativeOptimizer(
                         metadata,
@@ -417,6 +426,7 @@ public PlanOptimizers(
                                 .addAll(columnPruningRules)
                                 .addAll(ImmutableSet.of(
                                         new MergeDuplicateAggregation(metadata.getFunctionAndTypeManager()),
+                                        new ImplementTableFunctionSource(metadata),
                                         new RemoveRedundantIdentityProjections(),
                                         new RemoveFullSample(),
                                         new EvaluateZeroSample(),

presto-main-base/src/main/java/com/facebook/presto/sql/planner/PlannerUtils.java

@@ -43,6 +43,7 @@
 import com.facebook.presto.spi.relation.SpecialFormExpression;
 import com.facebook.presto.spi.relation.VariableReferenceExpression;
 import com.facebook.presto.sql.analyzer.Field;
+import com.facebook.presto.sql.analyzer.RelationType;
 import com.facebook.presto.sql.planner.iterative.Lookup;
 import com.facebook.presto.sql.planner.plan.ExchangeNode;
 import com.facebook.presto.sql.planner.planPrinter.PlanPrinter;
@@ -205,6 +206,9 @@ public static PlanNode addOverrideProjection(PlanNode source, PlanNodeIdAllocato
                 || source.getOutputVariables().stream().distinct().count() != source.getOutputVariables().size()) {
             return source;
         }
+        if (source instanceof ProjectNode && ((ProjectNode) source).getAssignments().getMap().equals(variableMap)) {
+            return source;
+        }
         Assignments.Builder assignmentsBuilder = Assignments.builder();
         assignmentsBuilder.putAll(source.getOutputVariables().stream().collect(toImmutableMap(identity(), x -> variableMap.containsKey(x) ? variableMap.get(x) : x)));
         return new ProjectNode(source.getSourceLocation(), planNodeIdAllocator.getNextId(), source, assignmentsBuilder.build(), LOCAL);
@@ -574,4 +578,19 @@ public static RowExpression randomizeJoinKey(Session session, FunctionAndTypeMan
         }
         return new SpecialFormExpression(COALESCE, VARCHAR, ImmutableList.of(castToVarchar, concatExpression));
     }
+
+    public static int[] getFieldIndexesForVisibleColumns(RelationPlan sourcePlan)
+    {
+        // required columns are a subset of visible columns of the source. remap required column indexes to field indexes in source relation type.
+        RelationType sourceRelationType = sourcePlan.getScope().getRelationType();
+        int[] fieldIndexForVisibleColumn = new int[sourceRelationType.getVisibleFieldCount()];
+        int visibleColumn = 0;
+        for (int i = 0; i < sourceRelationType.getAllFieldCount(); i++) {
+            if (!sourceRelationType.getFieldByIndex(i).isHidden()) {
+                fieldIndexForVisibleColumn[visibleColumn] = i;
+                visibleColumn++;
+            }
+        }
+        return fieldIndexForVisibleColumn;
+    }
 }

presto-main-base/src/main/java/com/facebook/presto/sql/planner/QueryPlanner.java

@@ -143,7 +143,7 @@
 import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
-class QueryPlanner
+public class QueryPlanner
 {
     private final Analysis analysis;
     private final VariableAllocator variableAllocator;
@@ -524,7 +524,7 @@ private PlanBuilder project(PlanBuilder subPlan, Iterable<Expression> expression
      *
      * @return the new subplan and a mapping of each expression to the symbol representing the coercion or an existing symbol if a coercion wasn't needed
      */
-    private PlanAndMappings coerce(PlanBuilder subPlan, List<Expression> expressions, Analysis analysis, PlanNodeIdAllocator idAllocator, VariableAllocator variableAllocator, Metadata metadata)
+    public PlanAndMappings coerce(PlanBuilder subPlan, List<Expression> expressions, Analysis analysis, PlanNodeIdAllocator idAllocator, VariableAllocator variableAllocator, Metadata metadata)
     {
         Assignments.Builder assignments = Assignments.builder();
         assignments.putAll(subPlan.getRoot().getOutputVariables().stream().collect(toImmutableMap(Function.identity(), Function.identity())));
@@ -1346,7 +1346,7 @@ private RowExpression rowExpression(Expression expression, SqlPlannerContext con
                 context.getTranslatorContext());
     }
 
-    private static List<Expression> toSymbolReferences(List<VariableReferenceExpression> variables)
+    public static List<Expression> toSymbolReferences(List<VariableReferenceExpression> variables)
     {
         return variables.stream()
                 .map(variable -> new SymbolReference(
@@ -1355,6 +1355,11 @@ private static List<Expression> toSymbolReferences(List<VariableReferenceExpress
                 .collect(toImmutableList());
     }
 
+    public static SymbolReference toSymbolReference(VariableReferenceExpression variable)
+    {
+        return new SymbolReference(variable.getSourceLocation().map(location -> new NodeLocation(location.getLine(), location.getColumn())), variable.getName());
+    }
+
     public static class PlanAndMappings
     {
         private final PlanBuilder subPlan;