From 1e3573f85922a333fe5271ca33396af7449c3de3 Mon Sep 17 00:00:00 2001
From: Xiaoli Zhou <yihe.zxl@alibaba-inc.com>
Date: Tue, 10 Sep 2024 17:49:33 +0800
Subject: [PATCH] fix(interactive): Fix Bugs of Type Inference in `both()`
 (#4199)

---
 .../common/ir/planner/GraphIOProcessor.java   |  8 +++-
 .../gremlin/antlr4x/GraphBuilderTest.java     | 41 +++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/planner/GraphIOProcessor.java b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/planner/GraphIOProcessor.java
index 0cddd1c22932..6d2af7f93997 100644
--- a/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/planner/GraphIOProcessor.java
+++ b/interactive_engine/compiler/src/main/java/com/alibaba/graphscope/common/ir/planner/GraphIOProcessor.java
@@ -475,12 +475,16 @@ private void checkPattern(Pattern pattern) {
                                 });
                 if (edge.getElementDetails().getRange() == null) {
                     Preconditions.checkArgument(
-                            Sets.newHashSet(src.getVertexTypeIds()).equals(expectedSrcIds),
+                            !edge.isBoth()
+                                    ? Sets.newHashSet(src.getVertexTypeIds()).equals(expectedSrcIds)
+                                    : expectedSrcIds.containsAll(src.getVertexTypeIds()),
                             "src vertex types %s not consistent with edge types %s",
                             src.getVertexTypeIds(),
                             edge.getEdgeTypeIds());
                     Preconditions.checkArgument(
-                            Sets.newHashSet(dst.getVertexTypeIds()).equals(expectedDstIds),
+                            !edge.isBoth()
+                                    ? Sets.newHashSet(dst.getVertexTypeIds()).equals(expectedDstIds)
+                                    : expectedDstIds.containsAll(dst.getVertexTypeIds()),
                             "dst vertex types %s not consistent with edge types %s",
                             dst.getVertexTypeIds(),
                             edge.getEdgeTypeIds());
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 bb53bff8a9b1..9141dff65570 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
@@ -1778,4 +1778,45 @@ public void g_V_path_as_a_select_a_valueMap() {
                     + " person]}], alias=[_], opt=[VERTEX])",
                 rel.explain().trim());
     }
+
+    @Test
+    public void g_V_match_as_a_person_both_as_b_test() {
+        GraphBuilder builder = Utils.mockGraphBuilder(optimizer, irMeta);
+        RelNode node1 =
+                eval("g.V().match(as('a').hasLabel('person').both().as('b')).count()", builder);
+        RelNode after1 = optimizer.optimize(node1, new GraphIOProcessor(builder, irMeta));
+        Assert.assertEquals(
+                "GraphLogicalAggregate(keys=[{variables=[], aliases=[]}], values=[[{operands=[a,"
+                        + " b], aggFunction=COUNT, alias='$f0', distinct=false}]])\n"
+                        + "  GraphPhysicalExpand(tableConfig=[[EdgeLabel(knows, person, person),"
+                        + " EdgeLabel(created, person, software)]], alias=[b], startAlias=[a],"
+                        + " opt=[BOTH], physicalOpt=[VERTEX])\n"
+                        + "    GraphLogicalSource(tableConfig=[{isAll=false, tables=[person]}],"
+                        + " alias=[a], opt=[VERTEX])",
+                after1.explain().trim());
+        RelNode node2 =
+                eval("g.V().match(as('a').hasLabel('software').both().as('b')).count()", builder);
+        RelNode after2 = optimizer.optimize(node2, new GraphIOProcessor(builder, irMeta));
+        Assert.assertEquals(
+                "GraphLogicalAggregate(keys=[{variables=[], aliases=[]}], values=[[{operands=[a,"
+                    + " b], aggFunction=COUNT, alias='$f0', distinct=false}]])\n"
+                    + "  GraphPhysicalGetV(tableConfig=[{isAll=false, tables=[person]}], alias=[b],"
+                    + " opt=[OTHER], physicalOpt=[ITSELF])\n"
+                    + "    GraphPhysicalExpand(tableConfig=[{isAll=false, tables=[created]}],"
+                    + " alias=[_], startAlias=[a], opt=[BOTH], physicalOpt=[VERTEX])\n"
+                    + "      GraphLogicalSource(tableConfig=[{isAll=false, tables=[software]}],"
+                    + " alias=[a], opt=[VERTEX])",
+                after2.explain().trim());
+        RelNode node3 = eval("g.V().match(as('a').both().as('b')).count()", builder);
+        RelNode after3 = optimizer.optimize(node3, new GraphIOProcessor(builder, irMeta));
+        Assert.assertEquals(
+                "GraphLogicalAggregate(keys=[{variables=[], aliases=[]}], values=[[{operands=[a,"
+                        + " b], aggFunction=COUNT, alias='$f0', distinct=false}]])\n"
+                        + "  GraphPhysicalExpand(tableConfig=[[EdgeLabel(knows, person, person),"
+                        + " EdgeLabel(created, person, software)]], alias=[b], startAlias=[a],"
+                        + " opt=[BOTH], physicalOpt=[VERTEX])\n"
+                        + "    GraphLogicalSource(tableConfig=[{isAll=false, tables=[software,"
+                        + " person]}], alias=[a], opt=[VERTEX])",
+                after3.explain().trim());
+    }
 }