Analysis for table and descriptor arguments.

Changes adapted from trino/PR#13602, PR#13653, PR#14115
Original commit: dbef4d9be37494967496230573ab400e54aab0d9
Author: kasiafi

Co-authored-by: kasiafi <30203062+kasiafi@users.noreply.github.com>
Co-authored-by: Xin Zhang <desertsxin@gmail.com>

Author: 3 people

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

@@ -49,6 +49,7 @@
 import com.facebook.presto.sql.tree.Offset;
 import com.facebook.presto.sql.tree.OrderBy;
 import com.facebook.presto.sql.tree.Parameter;
+import com.facebook.presto.sql.tree.QualifiedName;
 import com.facebook.presto.sql.tree.QuantifiedComparisonExpression;
 import com.facebook.presto.sql.tree.Query;
 import com.facebook.presto.sql.tree.QuerySpecification;
@@ -214,7 +215,11 @@ public class Analysis
 
     private Optional<List<OutputColumnMetadata>> updatedSourceColumns = Optional.empty();
 
+    // names of tables and aliased relations. All names are resolved case-insensitive.
+    private final Map<NodeRef<Relation>, QualifiedName> relationNames = new LinkedHashMap<>();
     private final Map<NodeRef<TableFunctionInvocation>, TableFunctionInvocationAnalysis> tableFunctionAnalyses = new LinkedHashMap<>();
+    private final Set<NodeRef<Relation>> aliasedRelations = new LinkedHashSet<>();
+    private final Set<NodeRef<TableFunctionInvocation>> polymorphicTableFunctions = new LinkedHashSet<>();
 
     public Analysis(@Nullable Statement root, Map<NodeRef<Parameter>, Expression> parameters, boolean isDescribe)
     {
@@ -1120,6 +1125,36 @@ public TableFunctionInvocationAnalysis getTableFunctionAnalysis(TableFunctionInv
         return tableFunctionAnalyses.get(NodeRef.of(node));
     }
 
+    public void setRelationName(Relation relation, QualifiedName name)
+    {
+        relationNames.put(NodeRef.of(relation), name);
+    }
+
+    public QualifiedName getRelationName(Relation relation)
+    {
+        return relationNames.get(NodeRef.of(relation));
+    }
+
+    public void addAliased(Relation relation)
+    {
+        aliasedRelations.add(NodeRef.of(relation));
+    }
+
+    public boolean isAliased(Relation relation)
+    {
+        return aliasedRelations.contains(NodeRef.of(relation));
+    }
+
+    public void addPolymorphicTableFunction(TableFunctionInvocation invocation)
+    {
+        polymorphicTableFunctions.add(NodeRef.of(invocation));
+    }
+
+    public boolean isPolymorphicTableFunction(TableFunctionInvocation invocation)
+    {
+        return polymorphicTableFunctions.contains(NodeRef.of(invocation));
+    }
+
     @Immutable
     public static final class Insert
     {
@@ -1371,23 +1406,229 @@ public int hashCode()
         }
     }
 
+    public static class TableArgumentAnalysis
+    {
+        private final String argumentName;
+        private final Optional<QualifiedName> name;
+        private final Relation relation;
+        private final Optional<List<Expression>> partitionBy; // it is allowed to partition by empty list
+        private final Optional<OrderBy> orderBy;
+        private final boolean pruneWhenEmpty;
+        private final boolean rowSemantics;
+        private final boolean passThroughColumns;
+
+        private TableArgumentAnalysis(
+                String argumentName,
+                Optional<QualifiedName> name,
+                Relation relation,
+                Optional<List<Expression>> partitionBy,
+                Optional<OrderBy> orderBy,
+                boolean pruneWhenEmpty,
+                boolean rowSemantics,
+                boolean passThroughColumns)
+        {
+            this.argumentName = requireNonNull(argumentName, "argumentName is null");
+            this.name = requireNonNull(name, "name is null");
+            this.relation = requireNonNull(relation, "relation is null");
+            this.partitionBy = requireNonNull(partitionBy, "partitionBy is null").map(ImmutableList::copyOf);
+            this.orderBy = requireNonNull(orderBy, "orderBy is null");
+            this.pruneWhenEmpty = pruneWhenEmpty;
+            this.rowSemantics = rowSemantics;
+            this.passThroughColumns = passThroughColumns;
+        }
+
+        public String getArgumentName()
+        {
+            return argumentName;
+        }
+
+        public Optional<QualifiedName> getName()
+        {
+            return name;
+        }
+
+        public Relation getRelation()
+        {
+            return relation;
+        }
+
+        public Optional<List<Expression>> getPartitionBy()
+        {
+            return partitionBy;
+        }
+
+        public Optional<OrderBy> getOrderBy()
+        {
+            return orderBy;
+        }
+
+        public boolean isPruneWhenEmpty()
+        {
+            return pruneWhenEmpty;
+        }
+
+        public boolean isRowSemantics()
+        {
+            return rowSemantics;
+        }
+
+        public boolean isPassThroughColumns()
+        {
+            return passThroughColumns;
+        }
+
+        public static Builder builder()
+        {
+            return new Builder();
+        }
+
+        public static final class Builder
+        {
+            private String argumentName;
+            private Optional<QualifiedName> name = Optional.empty();
+            private Relation relation;
+            private Optional<List<Expression>> partitionBy = Optional.empty();
+            private Optional<OrderBy> orderBy = Optional.empty();
+            private boolean pruneWhenEmpty;
+            private boolean rowSemantics;
+            private boolean passThroughColumns;
+
+            private Builder() {}
+
+            public Builder withArgumentName(String argumentName)
+            {
+                this.argumentName = argumentName;
+                return this;
+            }
+
+            public Builder withName(QualifiedName name)
+            {
+                this.name = Optional.of(name);
+                return this;
+            }
+
+            public Builder withRelation(Relation relation)
+            {
+                this.relation = relation;
+                return this;
+            }
+
+            public Builder withPartitionBy(List<Expression> partitionBy)
+            {
+                this.partitionBy = Optional.of(partitionBy);
+                return this;
+            }
+
+            public Builder withOrderBy(OrderBy orderBy)
+            {
+                this.orderBy = Optional.of(orderBy);
+                return this;
+            }
+
+            public Builder withPruneWhenEmpty(boolean pruneWhenEmpty)
+            {
+                this.pruneWhenEmpty = pruneWhenEmpty;
+                return this;
+            }
+
+            public Builder withRowSemantics(boolean rowSemantics)
+            {
+                this.rowSemantics = rowSemantics;
+                return this;
+            }
+
+            public Builder withPassThroughColumns(boolean passThroughColumns)
+            {
+                this.passThroughColumns = passThroughColumns;
+                return this;
+            }
+
+            public TableArgumentAnalysis build()
+            {
+                return new TableArgumentAnalysis(argumentName, name, relation, partitionBy, orderBy, pruneWhenEmpty, rowSemantics, passThroughColumns);
+            }
+        }
+    }
+
     /**
      * Encapsulates the result of analyzing a table function invocation.
      * Includes the connector ID, function name, argument bindings, and the
      * connector-specific table function handle needed for planning and execution.
+     *
+     * Example of a TableFunctionInvocationAnalysis for a table function
+     * with two table arguments, required columns, and co-partitioning
+     * implemented by {@link TestingTableFunctions.TwoTableArgumentsFunction}
+     *
+     * SQL:
+     * SELECT * FROM TABLE(system.two_table_arguments_function(
+     *     input1 => TABLE(t1) PARTITION BY (a, b),
+     *     input2 => TABLE(SELECT 1, 2) t1(x, y) PARTITION BY (x, y)
+     *     COPARTITION(t1, s1.t1)))
+     *
+     * Table Function:
+     *   super(
+     *     SCHEMA_NAME,
+     *     "two_table_arguments_function",
+     *     ImmutableList.of(
+     *       TableArgumentSpecification.builder()
+     *         .name("INPUT1")
+     *         .build(),
+     *       TableArgumentSpecification.builder()
+     *         .name("INPUT2")
+     *         .build()),
+     *     GENERIC_TABLE);
+     *
+     * analyze:
+     *   return TableFunctionAnalysis.builder()
+     *     .handle(HANDLE)
+     *     .returnedType(new Descriptor(ImmutableList.of(new Descriptor.Field(COLUMN_NAME, Optional.of(BOOLEAN)))))
+     *     .requiredColumns("INPUT1", ImmutableList.of(0))
+     *     .requiredColumns("INPUT2", ImmutableList.of(0))
+     *     .build();
+     *
+     * Example values:
+     * connectorId = "tpch"
+     * functionName = "two_table_arguments_function"
+     * arguments = {
+     *     "input1" -> row(a bigint,b bigint,c bigint,d bigint),
+     *     "input2" -> row(x integer,y integer)
+     * }
+     * tableArgumentAnalyses = [
+     *     { argumentName="INPUT1", relation=Table{t1}, partitionBy=[a, b], orderBy=[], pruneWhenEmpty=false },
+     *     { argumentName="INPUT2", relation=AliasedRelation{SELECT 1, 2 AS x, y}, partitionBy=[x, y], orderBy=[], pruneWhenEmpty=false }
+     * ]
+     * requiredColumns = {
+     *     "INPUT1" -> [0],
+     *     "INPUT2" -> [0]
+     * }
+     * copartitioningLists = [
+     *     ["INPUT2", "INPUT1"]
+     * ]
+     * properColumnsCount = 1
+     * connectorTableFunctionHandle = TestingTableFunctionPushdownHandle
+     * transactionHandle = AbstractAnalyzerTest$1$1
+     *
      */
     public static class TableFunctionInvocationAnalysis
     {
         private final ConnectorId connectorId;
         private final String functionName;
         private final Map<String, Argument> arguments;
+        private final List<TableArgumentAnalysis> tableArgumentAnalyses;
+        private final Map<String, List<Integer>> requiredColumns;
+        private final List<List<String>> copartitioningLists;
+        private final int properColumnsCount;
         private final ConnectorTableFunctionHandle connectorTableFunctionHandle;
         private final ConnectorTransactionHandle transactionHandle;
 
         public TableFunctionInvocationAnalysis(
                 ConnectorId connectorId,
                 String functionName,
                 Map<String, Argument> arguments,
+                List<TableArgumentAnalysis> tableArgumentAnalyses,
+                Map<String, List<Integer>> requiredColumns,
+                List<List<String>> copartitioningLists,
+                int properColumnsCount,
                 ConnectorTableFunctionHandle connectorTableFunctionHandle,
                 ConnectorTransactionHandle transactionHandle)
         {
@@ -1396,6 +1637,11 @@ public TableFunctionInvocationAnalysis(
             this.arguments = ImmutableMap.copyOf(arguments);
             this.connectorTableFunctionHandle = requireNonNull(connectorTableFunctionHandle, "connectorTableFunctionHandle is null");
             this.transactionHandle = requireNonNull(transactionHandle, "transactionHandle is null");
+            this.tableArgumentAnalyses = ImmutableList.copyOf(tableArgumentAnalyses);
+            this.requiredColumns = requiredColumns.entrySet().stream()
+                    .collect(toImmutableMap(Map.Entry::getKey, entry -> ImmutableList.copyOf(entry.getValue())));
+            this.copartitioningLists = ImmutableList.copyOf(copartitioningLists);
+            this.properColumnsCount = properColumnsCount;
         }
 
         public ConnectorId getConnectorId()
@@ -1413,6 +1659,31 @@ public Map<String, Argument> getArguments()
             return arguments;
         }
 
+        public List<TableArgumentAnalysis> getTableArgumentAnalyses()
+        {
+            return tableArgumentAnalyses;
+        }
+
+        public Map<String, List<Integer>> getRequiredColumns()
+        {
+            return requiredColumns;
+        }
+
+        public List<List<String>> getCopartitioningLists()
+        {
+            return copartitioningLists;
+        }
+
+        /**
+         * Proper columns are the columns produced by the table function, as opposed to pass-through columns from input tables.
+         * Proper columns should be considered the actual result of the table function.
+         * @return the number of table function's proper columns
+         */
+        public int getProperColumnsCount()
+        {
+            return properColumnsCount;
+        }
+
         public ConnectorTableFunctionHandle getConnectorTableFunctionHandle()
         {
             return connectorTableFunctionHandle;

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

@@ -57,14 +57,6 @@ public enum SemanticErrorCode
     INVALID_FUNCTION_NAME,
     DUPLICATE_PARAMETER_NAME,
     EXCEPTIONS_WHEN_RESOLVING_FUNCTIONS,
-
-    MISSING_RETURN_TYPE,
-    AMBIGUOUS_RETURN_TYPE,
-    MISSING_ARGUMENT,
-    INVALID_FUNCTION_ARGUMENT,
-    INVALID_ARGUMENTS,
-    NOT_IMPLEMENTED,
-
     ORDER_BY_MUST_BE_IN_SELECT,
     ORDER_BY_MUST_BE_IN_AGGREGATE,
     REFERENCE_TO_OUTPUT_ATTRIBUTE_WITHIN_ORDER_BY_GROUPING,
@@ -119,4 +111,16 @@ public enum SemanticErrorCode
     TOO_MANY_GROUPING_SETS,
 
     INVALID_OFFSET_ROW_COUNT,
+
+    TABLE_FUNCTION_MISSING_RETURN_TYPE,
+    TABLE_FUNCTION_AMBIGUOUS_RETURN_TYPE,
+    TABLE_FUNCTION_MISSING_ARGUMENT,
+    TABLE_FUNCTION_INVALID_FUNCTION_ARGUMENT,
+    TABLE_FUNCTION_INVALID_ARGUMENTS,
+    TABLE_FUNCTION_IMPLEMENTATION_ERROR,
+    TABLE_FUNCTION_INVALID_COPARTITIONING,
+    TABLE_FUNCTION_INVALID_TABLE_FUNCTION_INVOCATION,
+    TABLE_FUNCTION_DUPLICATE_RANGE_VARIABLE,
+    TABLE_FUNCTION_INVALID_COLUMN_REFERENCE,
+    TABLE_FUNCTION_COLUMN_NOT_FOUND
 }

presto-main-base/src/main/java/com/facebook/presto/metadata/TableFunctionRegistry.java

@@ -21,6 +21,7 @@
 import com.facebook.presto.spi.function.SchemaFunctionName;
 import com.facebook.presto.spi.function.table.ArgumentSpecification;
 import com.facebook.presto.spi.function.table.ConnectorTableFunction;
+import com.facebook.presto.spi.function.table.ReturnTypeSpecification.DescribedTable;
 import com.facebook.presto.spi.function.table.TableArgumentSpecification;
 import com.facebook.presto.sql.analyzer.SemanticException;
 import com.facebook.presto.sql.tree.QualifiedName;
@@ -152,5 +153,10 @@ private static void validateTableFunction(ConnectorTableFunction tableFunction)
         // The 'keep when empty' or 'prune when empty' property must not be explicitly specified for a table argument with row semantics.
         // Such a table argument is implicitly 'prune when empty'. The TableArgumentSpecification.Builder enforces the 'prune when empty' property
         // for a table argument with row semantics.
+
+        if (tableFunction.getReturnTypeSpecification() instanceof DescribedTable) {
+            DescribedTable describedTable = (DescribedTable) tableFunction.getReturnTypeSpecification();
+            checkArgument(describedTable.getDescriptor().isTyped(), "field types missing in returned type specification");
+        }
     }
 }