Repository: alibaba/GraphScope (language: python)

Issue:
[BUG] Unexpected results when dealing with HEAD in the new compilation stack
**Describe the bug**

Run the following gremlin query on the new compilation stack, and it outputs errors:

```
gremlin> g.V().match(__.as("a").out().as("b")).count()
{alias=_} not found; expected aliases are: [a, b]
```




Discussion / hints:

Fix patch:
diff --git a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/rex/RexGraphVariableList.java b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/rex/RexGraphVariableList.java
new file mode 100644
index 000000000000..aa41e537a6e7
--- /dev/null
+++ b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/rex/RexGraphVariableList.java
@@ -0,0 +1,198 @@
+/*
+ *
+ *  * Copyright 2020 Alibaba Group Holding Limited.
+ *  *
+ *  * Licensed under the Apache License, Version 2.0 (the "License");
+ *  * you may not use this file except in compliance with the License.
+ *  * You may obtain a copy of the License at
+ *  *
+ *  * http://www.apache.org/licenses/LICENSE-2.0
+ *  *
+ *  * Unless required by applicable law or agreed to in writing, software
+ *  * distributed under the License is distributed on an "AS IS" BASIS,
+ *  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  * See the License for the specific language governing permissions and
+ *  * limitations under the License.
+ *
+ */
+
+package com.alibaba.graphscope.common.ir.rex;
+
+import com.alibaba.graphscope.common.ir.tools.AliasInference;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFieldImpl;
+import org.apache.calcite.rel.type.RelRecordType;
+import org.apache.calcite.rel.type.StructKind;
+import org.apache.commons.lang3.ObjectUtils;
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.List;
+import java.util.ListIterator;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+// denote a list of RexGraphVariable, to handle the special case of gremlin `head` if the input
+// operator outputs multiple columns
+public class RexGraphVariableList extends RexGraphVariable implements List<RexGraphVariable> {
+    private final List<RexGraphVariable> innerList;
+
+    protected RexGraphVariableList(
+            int aliasId,
+            int columnId,
+            @Nullable String name,
+            RelDataType type,
+            List<RexGraphVariable> innerList) {
+        super(aliasId, columnId, name, type);
+        this.innerList = ObjectUtils.requireNonEmpty(innerList);
+    }
+
+    public static RexGraphVariable of(List<RexGraphVariable> innerList) {
+        if (innerList.size() == 1) return innerList.get(0);
+        List<String> digest =
+                innerList.stream().map(var -> var.getName()).collect(Collectors.toList());
+        RelRecordType fullyType =
+                new RelRecordType(
+                        StructKind.FULLY_QUALIFIED,
+                        innerList.stream()
+                                .map(
+                                        var -> {
+                                            String[] splits =
+                                                    var.getName()
+                                                            .split(
+                                                                    Pattern.quote(
+                                                                            AliasInference
+                                                                                    .DELIMITER));
+                                            String aliasName =
+                                                    splits.length > 0
+                                                            ? splits[0]
+                                                            : AliasInference.DEFAULT_NAME;
+                                            return new RelDataTypeFieldImpl(
+                                                    aliasName, var.getAliasId(), var.getType());
+                                        })
+                                .collect(Collectors.toList()));
+        return new RexGraphVariableList(
+                AliasInference.DEFAULT_ID,
+                AliasInference.DEFAULT_COLUMN_ID,
+                digest.toString(),
+                fullyType,
+                innerList);
+    }
+
+    @Override
+    public int size() {
+        return innerList.size();
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return innerList.isEmpty();
+    }
+
+    @Override
+    public boolean contains(Object o) {
+        return innerList.contains(o);
+    }
+
+    @Override
+    public Iterator<RexGraphVariable> iterator() {
+        return innerList.iterator();
+    }
+
+    @Override
+    public Object[] toArray() {
+        return innerList.toArray();
+    }
+
+    @Override
+    public <T> T[] toArray(T[] a) {
+        return innerList.toArray(a);
+    }
+
+    @Override
+    public boolean add(RexGraphVariable variable) {
+        return innerList.add(variable);
+    }
+
+    @Override
+    public boolean remove(Object o) {
+        return innerList.remove(o);
+    }
+
+    @Override
+    public boolean containsAll(Collection<?> c) {
+        return innerList.containsAll(c);
+    }
+
+    @Override
+    public boolean addAll(Collection<? extends RexGraphVariable> c) {
+        return innerList.addAll(c);
+    }
+
+    @Override
+    public boolean addAll(int index, Collection<? extends RexGraphVariable> c) {
+        return innerList.addAll(index, c);
+    }
+
+    @Override
+    public boolean removeAll(Collection<?> c) {
+        return innerList.removeAll(c);
+    }
+
+    @Override
+    public boolean retainAll(Collection<?> c) {
+        return innerList.retainAll(c);
+    }
+
+    @Override
+    public void clear() {
+        innerList.clear();
+    }
+
+    @Override
+    public RexGraphVariable get(int index) {
+        return innerList.get(index);
+    }
+
+    @Override
+    public RexGraphVariable set(int index, RexGraphVariable element) {
+        return innerList.set(index, element);
+    }
+
+    @Override
+    public void add(int index, RexGraphVariable element) {
+        innerList.add(index, element);
+    }
+
+    @Override
+    public RexGraphVariable remove(int index) {
+        return innerList.remove(index);
+    }
+
+    @Override
+    public int indexOf(Object o) {
+        return innerList.indexOf(o);
+    }
+
+    @Override
+    public int lastIndexOf(Object o) {
+        return innerList.lastIndexOf(o);
+    }
+
+    @Override
+    public ListIterator<RexGraphVariable> listIterator() {
+        return innerList.listIterator();
+    }
+
+    @Override
+    public ListIterator<RexGraphVariable> listIterator(int index) {
+        return innerList.listIterator(index);
+    }
+
+    @Override
+    public List<RexGraphVariable> subList(int fromIndex, int toIndex) {
+        return innerList.subList(fromIndex, toIndex);
+    }
+}
diff --git a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/AliasInference.java b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/AliasInference.java
index bcd5fd200529..36b6330ee2b8 100644
--- a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/AliasInference.java
+++ b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/AliasInference.java
@@ -47,6 +47,9 @@ public abstract class AliasInference {
     public static final String DEFAULT_NAME = "_";
     public static final int DEFAULT_ID = -1;
 
+    // denote all columns, to support count star
+    public static final String STAR = "*";
+
     public static final int DEFAULT_COLUMN_ID = 100;
 
     public static final String DELIMITER = ".";
diff --git a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/GraphBuilder.java b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/GraphBuilder.java
index 1f971d03f2c7..6d43089b45a8 100644
--- a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/GraphBuilder.java
+++ b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/tools/GraphBuilder.java
@@ -454,13 +454,28 @@ private JoinRelType getJoinRelType(GraphOpt.Match opt) {
      */
     public RexGraphVariable variable(@Nullable String alias) {
         alias = AliasInference.isDefaultAlias(alias) ? AliasInference.DEFAULT_NAME : alias;
-        ColumnField columnField = getAliasField(alias);
-        RelDataTypeField aliasField = columnField.right;
-        return RexGraphVariable.of(
-                aliasField.getIndex(),
-                columnField.left,
-                AliasInference.SIMPLE_NAME(alias),
-                aliasField.getType());
+        List<ColumnField> columnFields = getAliasField(alias);
+        if (columnFields.size() == 1) {
+            ColumnField columnField = columnFields.get(0);
+            RelDataTypeField aliasField = columnField.right;
+            return RexGraphVariable.of(
+                    aliasField.getIndex(),
+                    columnField.left,
+                    aliasField.getName(),
+                    aliasField.getType());
+        } else {
+            return RexGraphVariableList.of(
+                    columnFields.stream()
+                            .map(
+                                    field ->
+                                            RexGraphVariable.of(
+                                                    field.right.getIndex(),
+                                                    field.left,
+                                                    AliasInference.SIMPLE_NAME(
+                                                            field.right.getName()),
+                                                    field.right.getType()))
+                            .collect(Collectors.toList()));
+        }
     }
 
     /**
@@ -474,7 +489,17 @@ public RexGraphVariable variable(@Nullable String alias, String property) {
         alias = AliasInference.isDefaultAlias(alias) ? AliasInference.DEFAULT_NAME : alias;
         Objects.requireNonNull(property);
         String varName = AliasInference.SIMPLE_NAME(alias) + AliasInference.DELIMITER + property;
-        ColumnField columnField = getAliasField(alias);
+        List<ColumnField> columnFields = getAliasField(alias);
+        if (columnFields.size() != 1) {
+            throw new IllegalArgumentException(
+                    "cannot get property="
+                            + property
+                            + " from alias="
+                            + alias
+                            + ", expected one column, but found "
+                            + columnFields.size());
+        }
+        ColumnField columnField = columnFields.get(0);
         RelDataTypeField aliasField = columnField.right;
         if (property.equals(GraphProperty.LEN_KEY)) {
             if (!(aliasField.getType() instanceof GraphPathType)) {
@@ -557,8 +582,23 @@ public RexGraphVariable variable(@Nullable String alias, String property) {
      * @param alias
      * @return
      */
-    private ColumnField getAliasField(String alias) {
+    private List<ColumnField> getAliasField(String alias) {
         Objects.requireNonNull(alias);
+        if (alias.equals(AliasInference.STAR)) {
+            RelNode peek = requireNonNull(peek(), "frame stack is empty");
+            RelDataType outputType =
+                    com.alibaba.graphscope.common.ir.tools.Utils.getOutputType(peek);
+            return outputType.getFieldList().stream()
+                    .map(
+                            field -> {
+                                int columnIdx =
+                                        AliasInference.isDefaultAlias(field.getName())
+                                                ? AliasInference.DEFAULT_COLUMN_ID
+                                                : getColumnIndex(peek, field);
+                                return new ColumnField(columnIdx, field);
+                            })
+                    .collect(Collectors.toList());
+        }
         Set<String> aliases = new HashSet<>();
         int nodeIdx = 0;
         for (int inputOrdinal = 0; inputOrdinal < size(); ++inputOrdinal) {
@@ -569,27 +609,34 @@ private ColumnField getAliasField(String alias) {
                 // to support `head` in gremlin
                 if (nodeIdx++ == 0 && AliasInference.isDefaultAlias(alias)) {
                     if (fields.size() == 1) {
-                        return new ColumnField(
-                                AliasInference.DEFAULT_COLUMN_ID,
-                                new RelDataTypeFieldImpl(
-                                        AliasInference.DEFAULT_NAME,
-                                        AliasInference.DEFAULT_ID,
-                                        fields.get(0).getType()));
+                        return ImmutableList.of(
+                                new ColumnField(
+                                        AliasInference.DEFAULT_COLUMN_ID,
+                                        new RelDataTypeFieldImpl(
+                                                AliasInference.DEFAULT_NAME,
+                                                AliasInference.DEFAULT_ID,
+                                                fields.get(0).getType())));
                     } else if (cur
                             instanceof
                             CommonTableScan) { // specific implementation for gremlin, to get `head`
                         // in nested traversal
-                        return new ColumnField(
-                                AliasInference.DEFAULT_COLUMN_ID,
-                                new RelDataTypeFieldImpl(
-                                        AliasInference.DEFAULT_NAME,
-                                        AliasInference.DEFAULT_ID,
-                                        fields.get(fields.size() - 1).getType()));
+                        return ImmutableList.of(
+                                new ColumnField(
+                                        AliasInference.DEFAULT_COLUMN_ID,
+                                        new RelDataTypeFieldImpl(
+                                                AliasInference.DEFAULT_NAME,
+                                                AliasInference.DEFAULT_ID,
+                                                fields.get(fields.size() - 1).getType())));
+                    } else { // If the output of the current operator consists of multiple columns,
+                        // we transform `head` into all columns of the current operator.
+                        return fields.stream()
+                                .map(field -> new ColumnField(getColumnIndex(cur, field), field))
+                                .collect(Collectors.toList());
                     }
                 }
                 for (RelDataTypeField field : fields) {
                     if (!AliasInference.isDefaultAlias(alias) && field.getName().equals(alias)) {
-                        return new ColumnField(getColumnIndex(cur, field), field);
+                        return ImmutableList.of(new ColumnField(getColumnIndex(cur, field), field));
                     }
                     aliases.add(AliasInference.SIMPLE_NAME(field.getName()));
                 }
@@ -779,6 +826,8 @@ public GraphBuilder filter(RexNode... conditions) {
 
     @Override
     public GraphBuilder filter(Iterable<? extends RexNode> conditions) {
+        conditions = flatExprs((Iterable<RexNode>) conditions);
+
         RexVisitor propertyChecker = new RexPropertyChecker(true, this);
         for (RexNode condition : conditions) {
             RelDataType type = condition.getType();
@@ -1186,6 +1235,8 @@ public GraphBuilder project(
         Config config = Utils.getFieldValue(RelBuilder.class, this, "config");
         RexSimplify simplifier = Utils.getFieldValue(RelBuilder.class, this, "simplifier");
 
+        nodes = flatExprs((Iterable<RexNode>) nodes);
+
         List<RexNode> nodeList = Lists.newArrayList(nodes);
         List<@Nullable String> fieldNameList = Lists.newArrayList(aliases);
 
@@ -1368,7 +1419,7 @@ public GroupKey groupKey(List<RexNode> variables, List<@Nullable String> aliases
      * @return
      */
     private GroupKey groupKey_(List<RexNode> variables, List<@Nullable String> aliases) {
-        return new GraphGroupKeys(variables, aliases);
+        return new GraphGroupKeys((List<RexNode>) flatExprs(variables), aliases);
     }
 
     // build aggregate functions
@@ -1446,15 +1497,11 @@ protected AggCall aggregateCall(
             ImmutableList<RexNode> orderKeys,
             @Nullable String alias,
             ImmutableList<RexNode> operands) {
-        // if operands of the aggregate call is empty, set a variable with (alias = null) by default
-        return new GraphAggCall(
-                        getCluster(),
-                        aggFunction,
-                        ObjectUtils.isNotEmpty(operands)
-                                ? operands
-                                : ImmutableList.of(this.variable((String) null)))
-                .as(alias)
-                .distinct(distinct);
+        if (operands.isEmpty()) { // to support count star
+            operands = ImmutableList.of(variable(AliasInference.STAR));
+        }
+        operands = ImmutableList.copyOf(flatExprs(operands));
+        return new GraphAggCall(getCluster(), aggFunction, operands).as(alias).distinct(distinct);
     }
 
     @Override
@@ -1536,6 +1583,8 @@ public GraphBuilder sortLimit(
             throw new IllegalArgumentException("FETCH node must be RexLiteral");
         }
 
+        nodes = flatExprs((Iterable<RexNode>) nodes);
+
         RelNode input = requireNonNull(peek(), "frame stack is empty");
 
         List<RelDataTypeField> originalFields = input.getRowType().getFieldList();
@@ -1641,6 +1690,8 @@ public GraphBuilder sortLimit(
     public GraphBuilder dedupBy(Iterable<? extends RexNode> nodes) {
         RelNode input = requireNonNull(peek(), "frame stack is empty");
 
+        nodes = flatExprs((Iterable<RexNode>) nodes);
+
         List<RelDataTypeField> originalFields = input.getRowType().getFieldList();
 
         Registrar registrar = new Registrar(this, input, true);
@@ -1946,4 +1997,39 @@ private LabelConfig getLabelConfig(TableConfig tableConfig) {
         labels.forEach(k -> labelConfig.addLabel(k));
         return labelConfig;
     }
+
+    // flat each expression in the given expression list, i.e. flat `RexGraphVariableList([a, b,
+    // c])` to `List of [RexGraphVariable(a), RexGraphVariable(b), RexGraphVariable(c)]`
+    private Iterable<RexNode> flatExprs(Iterable<RexNode> exprs) {
+        List<RexNode> flatExprs = Lists.newArrayList();
+        exprs.forEach(
+                expr -> {
+                    if (expr instanceof RexCall) {
+                        List<RexNode> operands = ((RexCall) expr).getOperands();
+                        if (operands.stream()
+                                .anyMatch(operand -> operand instanceof RexGraphVariableList)) {
+                            switch (expr.getKind()) {
+                                case DESCENDING:
+                                    operands =
+                                            (List<RexNode>)
+                                                    flatExprs(ImmutableList.of(operands.get(0)));
+                                    operands.forEach(operand -> flatExprs.add(desc(operand)));
+                                    break;
+                                default:
+                                    throw new IllegalArgumentException(
+                                            "cannot flat operands of "
+                                                    + expr.getKind()
+                                                    + " operator");
+                            }
+                        } else {
+                            flatExprs.add(expr);
+                        }
+                    } else if (expr instanceof RexGraphVariableList) {
+                        flatExprs.addAll((RexGraphVariableList) expr);
+                    } else {
+                        flatExprs.add(expr);
+                    }
+                });
+        return flatExprs;
+    }
 }
diff --git a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/gremlin/antlr4x/visitor/GraphBuilderVisitor.java b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/gremlin/antlr4x/visitor/GraphBuilderVisitor.java
index b58230e9802a..10aa08c94f72 100644
--- a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/gremlin/antlr4x/visitor/GraphBuilderVisitor.java
+++ b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/gremlin/antlr4x/visitor/GraphBuilderVisitor.java
@@ -548,7 +548,8 @@ public GraphBuilder visitOC_AggregateFunctionInvocation(
         switch (functionName) {
             case "count":
                 return (GraphBuilder)
-                        builder.aggregate(builder.groupKey(), builder.countStar(null));
+                        builder.aggregate(
+                                builder.groupKey(), builder.count(builder.variable((String) null)));
             case "sum":
                 return (GraphBuilder)
                         builder.aggregate(


Test patch:
diff --git a/interactive_engine/compiler/src/test/java/com/alibaba/graphscope/gremlin/antlr4x/GraphBuilderTest.java b/interactive_engine/compiler/src/test/java/com/alibaba/graphscope/gremlin/antlr4x/GraphBuilderTest.java
index ee50b47e33f9..d2a2219b49ba 100644
--- a/interactive_engine/compiler/src/test/java/com/alibaba/graphscope/gremlin/antlr4x/GraphBuilderTest.java
+++ b/interactive_engine/compiler/src/test/java/com/alibaba/graphscope/gremlin/antlr4x/GraphBuilderTest.java
@@ -1522,4 +1522,21 @@ public void g_V_as_a_select_expr_id_a_test() {
                     + " alias=[a], opt=[VERTEX], uniqueKeyFilters=[=(_.~id, 2)])",
                 node.explain().trim());
     }
+
+    @Test
+    public void g_V_match_as_a_out_as_b_count_test() {
+        RelNode node = eval("g.V().match(as(\"a\").out().as(\"b\")).count()");
+        Assert.assertEquals(
+                "GraphLogicalAggregate(keys=[{variables=[], aliases=[]}], values=[[{operands=[a,"
+                    + " b], aggFunction=COUNT, alias='$f0', distinct=false}]])\n"
+                    + "  GraphLogicalSingleMatch(input=[null],"
+                    + " sentence=[GraphLogicalGetV(tableConfig=[{isAll=true, tables=[software,"
+                    + " person]}], alias=[b], opt=[END])\n"
+                    + "  GraphLogicalExpand(tableConfig=[{isAll=true, tables=[created, knows]}],"
+                    + " alias=[_], opt=[OUT])\n"
+                    + "    GraphLogicalSource(tableConfig=[{isAll=false, tables=[person]}],"
+                    + " alias=[a], opt=[VERTEX])\n"
+                    + "], matchOpt=[INNER])",
+                node.explain().trim());
+    }
 }