Allow nested rowtype reference

presto-main/src/main/java/com/facebook/presto/metadata/SessionPropertyManager.java

@@ -30,6 +30,7 @@
 import com.facebook.presto.type.MapType;
 import com.google.common.collect.ComparisonChain;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Ordering;
 import io.airlift.json.JsonCodec;
 import io.airlift.json.JsonCodecFactory;
@@ -189,7 +190,7 @@ public <T> T decodeProperty(String name, @Nullable String value, Class<T> type)
     @NotNull
     public static Object evaluatePropertyValue(Expression expression, Type expectedType, Session session, Metadata metadata)
     {
-        Object value = evaluateConstantExpression(expression, expectedType, metadata, session);
+        Object value = evaluateConstantExpression(expression, expectedType, metadata, session, ImmutableSet.of());
 
         // convert to object value type of SQL type
         BlockBuilder blockBuilder = expectedType.createBlockBuilder(new BlockBuilderStatus(), 1);

presto-main/src/main/java/com/facebook/presto/operator/scalar/TestingRowConstructor.java

@@ -151,6 +151,20 @@ DOUBLE, new ArrayType(BIGINT),
         return toStackRepresentation(parameterTypes, arg1, arg2, arg3);
     }
 
+    @ScalarFunction("test_row")
+    @SqlType("row<double,array<row<bigint,double>('col0','col1')>,row<bigint,double>('col0','col1')>('col0','col1','col2')")
+    public static Block testNestedRowWithArray(
+            @Nullable @SqlType(StandardTypes.DOUBLE) Double arg1,
+            @Nullable @SqlType("array<row<bigint,double>('col0','col1')>") Block arg2,
+            @Nullable @SqlType("row<bigint,double>('col0','col1')") Block arg3)
+    {
+        List<Type> parameterTypes = ImmutableList.of(
+                DOUBLE,
+                new ArrayType(new RowType(ImmutableList.of(BIGINT, DOUBLE), Optional.of(ImmutableList.of("col0", "col1")))),
+                new RowType(ImmutableList.of(BIGINT, DOUBLE), Optional.of(ImmutableList.of("col0", "col1"))));
+        return toStackRepresentation(parameterTypes, arg1, arg2, arg3);
+    }
+
     @ScalarFunction("test_row")
     @SqlType("row<timestamp>('col0')")
     public static Block testRowBigintBigint(@Nullable @SqlType(StandardTypes.TIMESTAMP) Long arg1)

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

@@ -23,6 +23,7 @@
 import com.facebook.presto.sql.tree.CoalesceExpression;
 import com.facebook.presto.sql.tree.ComparisonExpression;
 import com.facebook.presto.sql.tree.CurrentTime;
+import com.facebook.presto.sql.tree.DereferenceExpression;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.Extract;
 import com.facebook.presto.sql.tree.FunctionCall;
@@ -48,22 +49,24 @@
 import com.facebook.presto.sql.tree.WhenClause;
 import com.facebook.presto.sql.tree.Window;
 import com.facebook.presto.sql.tree.WindowFrame;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 
 import javax.annotation.Nullable;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MUST_BE_AGGREGATE_OR_GROUP_BY;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NESTED_AGGREGATION;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NESTED_WINDOW;
 import static com.facebook.presto.sql.analyzer.SemanticErrorCode.NOT_SUPPORTED;
 import static com.facebook.presto.util.ImmutableCollectors.toImmutableList;
+import static com.facebook.presto.util.Types.checkType;
+import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Predicates.equalTo;
-import static com.google.common.base.Predicates.instanceOf;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -76,18 +79,20 @@ public class AggregationAnalyzer
     private final List<Expression> expressions;
 
     private final Metadata metadata;
+    private final Set<Expression> columnReferences;
 
     private final TupleDescriptor tupleDescriptor;
 
-    public AggregationAnalyzer(List<FieldOrExpression> groupByExpressions, Metadata metadata, TupleDescriptor tupleDescriptor)
+    public AggregationAnalyzer(List<FieldOrExpression> groupByExpressions, Metadata metadata, TupleDescriptor tupleDescriptor, Set<Expression> columnReferences)
     {
         requireNonNull(groupByExpressions, "groupByExpressions is null");
         requireNonNull(metadata, "metadata is null");
         requireNonNull(tupleDescriptor, "tupleDescriptor is null");
+        requireNonNull(columnReferences, "columnReferences is null");
 
         this.tupleDescriptor = tupleDescriptor;
         this.metadata = metadata;
-
+        this.columnReferences = ImmutableSet.copyOf(columnReferences);
         this.expressions = groupByExpressions.stream()
                 .filter(FieldOrExpression::isExpression)
                 .map(FieldOrExpression::getExpression)
@@ -103,18 +108,23 @@ public AggregationAnalyzer(List<FieldOrExpression> groupByExpressions, Metadata
         // For a query like "SELECT * FROM T GROUP BY a", groupByExpressions will contain "a",
         // and the '*' will be expanded to Field references. Therefore we translate all simple name expressions
         // in the group by clause to fields they reference so that the expansion from '*' can be matched against them
-        for (Expression expression : Iterables.filter(expressions, instanceOf(QualifiedNameReference.class))) {
-            QualifiedName name = ((QualifiedNameReference) expression).getName();
+        for (Expression expression : Iterables.filter(expressions, columnReferences::contains)) {
+            QualifiedName name;
+            if (expression instanceof QualifiedNameReference) {
+                name = ((QualifiedNameReference) expression).getName();
+            }
+            else {
+                name = DereferenceExpression.getQualifiedName(checkType(expression, DereferenceExpression.class, "expression"));
+            }
 
             List<Field> fields = tupleDescriptor.resolveFields(name);
-            Preconditions.checkState(fields.size() <= 1, "Found more than one field for name '%s': %s", name, fields);
+            checkState(fields.size() <= 1, "Found more than one field for name '%s': %s", name, fields);
 
             if (fields.size() == 1) {
                 Field field = Iterables.getOnlyElement(fields);
                 fieldIndexes.add(tupleDescriptor.indexOf(field));
             }
         }
-
         this.fieldIndexes = fieldIndexes.build();
     }
 
@@ -334,11 +344,25 @@ public Boolean visitWindowFrame(WindowFrame node, Void context)
         @Override
         protected Boolean visitQualifiedNameReference(QualifiedNameReference node, Void context)
         {
-            QualifiedName name = node.getName();
+            return isField(node.getName());
+        }
 
-            List<Field> fields = tupleDescriptor.resolveFields(name);
-            Preconditions.checkState(!fields.isEmpty(), "No fields for name '%s'", name);
-            Preconditions.checkState(fields.size() <= 1, "Found more than one field for name '%s': %s", name, fields);
+        @Override
+        protected Boolean visitDereferenceExpression(DereferenceExpression node, Void context)
+        {
+            if (columnReferences.contains(node)) {
+                return isField(DereferenceExpression.getQualifiedName(node));
+            }
+
+            // Allow SELECT col1.f1 FROM table1 GROUP BY col1
+            return process(node.getBase(), context);
+        }
+
+        private Boolean isField(QualifiedName qualifiedName)
+        {
+            List<Field> fields = tupleDescriptor.resolveFields(qualifiedName);
+            checkState(!fields.isEmpty(), "No fields for name '%s'", qualifiedName);
+            checkState(fields.size() <= 1, "Found more than one field for name '%s': %s", qualifiedName, fields);
 
             Field field = Iterables.getOnlyElement(fields);
             return fieldIndexes.contains(tupleDescriptor.indexOf(field));

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

@@ -24,8 +24,6 @@
 import com.facebook.presto.sql.tree.InPredicate;
 import com.facebook.presto.sql.tree.Join;
 import com.facebook.presto.sql.tree.Node;
-import com.facebook.presto.sql.tree.QualifiedName;
-import com.facebook.presto.sql.tree.QualifiedNameReference;
 import com.facebook.presto.sql.tree.Query;
 import com.facebook.presto.sql.tree.QuerySpecification;
 import com.facebook.presto.sql.tree.Relation;
@@ -44,6 +42,7 @@
 import java.util.Optional;
 import java.util.Set;
 
+import static com.google.common.collect.Sets.newIdentityHashSet;
 import static java.util.Objects.requireNonNull;
 
 public class Analysis
@@ -55,7 +54,7 @@ public class Analysis
 
     private TupleDescriptor outputDescriptor;
     private final IdentityHashMap<Node, TupleDescriptor> outputDescriptors = new IdentityHashMap<>();
-    private final IdentityHashMap<Expression, Map<QualifiedName, Integer>> resolvedNames = new IdentityHashMap<>();
+    private final IdentityHashMap<Expression, Map<Expression, Integer>> resolvedNames = new IdentityHashMap<>();
 
     private final IdentityHashMap<QuerySpecification, List<FunctionCall>> aggregates = new IdentityHashMap<>();
     private final IdentityHashMap<QuerySpecification, List<FieldOrExpression>> groupByExpressions = new IdentityHashMap<>();
@@ -71,11 +70,11 @@ public class Analysis
 
     private final IdentityHashMap<Table, TableHandle> tables = new IdentityHashMap<>();
 
-    private final IdentityHashMap<Expression, Boolean> rowFieldReferences = new IdentityHashMap<>();
     private final IdentityHashMap<Expression, Type> types = new IdentityHashMap<>();
     private final IdentityHashMap<Expression, Type> coercions = new IdentityHashMap<>();
     private final IdentityHashMap<Relation, Type[]> relationCoercions = new IdentityHashMap<>();
     private final IdentityHashMap<FunctionCall, Signature> functionSignature = new IdentityHashMap<>();
+    private final Set<Expression> columnReferences = newIdentityHashSet();
 
     private final IdentityHashMap<Field, ColumnHandle> columns = new IdentityHashMap<>();
 
@@ -122,12 +121,12 @@ public void setCreateTableAsSelectWithData(boolean createTableAsSelectWithData)
         this.createTableAsSelectWithData = createTableAsSelectWithData;
     }
 
-    public void addResolvedNames(Expression expression, Map<QualifiedName, Integer> mappings)
+    public void addResolvedNames(Expression expression, Map<Expression, Integer> mappings)
     {
         resolvedNames.put(expression, mappings);
     }
 
-    public Map<QualifiedName, Integer> getResolvedNames(Expression expression)
+    public Map<Expression, Integer> getResolvedNames(Expression expression)
     {
         return resolvedNames.get(expression);
     }
@@ -147,11 +146,6 @@ public IdentityHashMap<Expression, Type> getTypes()
         return new IdentityHashMap<>(types);
     }
 
-    public boolean isRowFieldReference(QualifiedNameReference qualifiedNameReference)
-    {
-        return rowFieldReferences.containsKey(qualifiedNameReference);
-    }
-
     public Type getType(Expression expression)
     {
         Preconditions.checkArgument(types.containsKey(expression), "Expression not analyzed: %s", expression);
@@ -309,14 +303,19 @@ public void addFunctionSignatures(IdentityHashMap<FunctionCall, Signature> infos
         functionSignature.putAll(infos);
     }
 
-    public void addTypes(IdentityHashMap<Expression, Type> types)
+    public Set<Expression> getColumnReferences()
     {
-        this.types.putAll(types);
+        return ImmutableSet.copyOf(columnReferences);
     }
 
-    public void addRowFieldReferences(IdentityHashMap<Expression, Boolean> rowFieldReferences)
+    public void addColumnReferences(Set<Expression> references)
     {
-        this.rowFieldReferences.putAll(rowFieldReferences);
+        columnReferences.addAll(references);
+    }
+
+    public void addTypes(IdentityHashMap<Expression, Type> types)
+    {
+        this.types.putAll(types);
     }
 
     public void addCoercion(Expression expression, Type type)

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

@@ -16,6 +16,7 @@
 import com.facebook.presto.spi.type.Type;
 import com.facebook.presto.sql.tree.Expression;
 import com.facebook.presto.sql.tree.InPredicate;
+import com.google.common.collect.ImmutableSet;
 
 import java.util.IdentityHashMap;
 import java.util.Set;
@@ -27,15 +28,18 @@ public class ExpressionAnalysis
     private final IdentityHashMap<Expression, Type> expressionTypes;
     private final IdentityHashMap<Expression, Type> expressionCoercions;
     private final Set<InPredicate> subqueryInPredicates;
+    private final Set<Expression> columnReferences;
 
     public ExpressionAnalysis(
             IdentityHashMap<Expression, Type> expressionTypes,
             IdentityHashMap<Expression, Type> expressionCoercions,
-            Set<InPredicate> subqueryInPredicates)
+            Set<InPredicate> subqueryInPredicates,
+            Set<Expression> columnReferences)
     {
         this.expressionTypes = requireNonNull(expressionTypes, "expressionTypes is null");
         this.expressionCoercions = requireNonNull(expressionCoercions, "expressionCoercions is null");
         this.subqueryInPredicates = requireNonNull(subqueryInPredicates, "subqueryInPredicates is null");
+        this.columnReferences = ImmutableSet.copyOf(requireNonNull(columnReferences, "columnReferences is null"));
     }
 
     public Type getType(Expression expression)
@@ -57,4 +61,9 @@ public Set<InPredicate> getSubqueryInPredicates()
     {
         return subqueryInPredicates;
     }
+
+    public Set<Expression> getColumnReferences()
+    {
+        return columnReferences;
+    }
 }