From 1e22423ac74f832ea9c5e3381be7511c61c86bd1 Mon Sep 17 00:00:00 2001
From: Cheng Xuntao <7731943+xtcyclist@users.noreply.github.com>
Date: Tue, 28 Mar 2023 15:05:13 +0800
Subject: [PATCH] Allow the expressions like 'v.tag' in return (#5440)

---
 src/graph/visitor/PropertyTrackerVisitor.cpp | 15 ++++++++++++---
 src/graph/visitor/PrunePropertiesVisitor.cpp |  6 +++++-
 tests/tck/features/match/Base.IntVid.feature |  4 ++--
 tests/tck/features/match/Base.feature        | 13 +++++++++++--
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/graph/visitor/PropertyTrackerVisitor.cpp b/src/graph/visitor/PropertyTrackerVisitor.cpp
index ede1c764fea..3e20f09f8e7 100644
--- a/src/graph/visitor/PropertyTrackerVisitor.cpp
+++ b/src/graph/visitor/PropertyTrackerVisitor.cpp
@@ -208,10 +208,19 @@ void PropertyTrackerVisitor::visit(AttributeExpression *expr) {
   auto &propName = constVal.getStr();
   switch (lhs->kind()) {
     case Expression::Kind::kInputProperty:
-    case Expression::Kind::kVarProperty: {  // $e.name
+    case Expression::Kind::kVarProperty: {
+      // maybe: $e.prop
       auto *varPropExpr = static_cast<PropertyExpression *>(lhs);
-      auto &edgeAlias = varPropExpr->prop();
-      propsUsed_.insertEdgeProp(edgeAlias, unknownType_, propName);
+      auto &entityAlias = varPropExpr->prop();
+      propsUsed_.insertEdgeProp(entityAlias, unknownType_, propName);
+      // maybe: $v.tag
+      auto ret = qctx_->schemaMng()->toTagID(space_, propName);
+      if (!ret.ok()) {
+        status_ = std::move(ret).status();
+        break;
+      }
+      auto tagId = ret.value();
+      propsUsed_.insertVertexProp(entityAlias, tagId, "*");
       break;
     }
     case Expression::Kind::kCase: {  // (case xxx).name
diff --git a/src/graph/visitor/PrunePropertiesVisitor.cpp b/src/graph/visitor/PrunePropertiesVisitor.cpp
index f36c288b2dd..9daed3f1626 100644
--- a/src/graph/visitor/PrunePropertiesVisitor.cpp
+++ b/src/graph/visitor/PrunePropertiesVisitor.cpp
@@ -424,7 +424,11 @@ void PrunePropertiesVisitor::pruneCurrent(AppendVertices *node) {
       usedProps.insert(unknownIter->second.begin(), unknownIter->second.end());
     }
     if (tagIter != usedVertexProps.end()) {
-      usedProps.insert(tagIter->second.begin(), tagIter->second.end());
+      if (tagIter->second.find("*") != tagIter->second.end()) {
+        usedProps.insert(props.begin(), props.end());
+      } else {
+        usedProps.insert(tagIter->second.begin(), tagIter->second.end());
+      }
     }
     if (usedProps.empty()) {
       continue;
diff --git a/tests/tck/features/match/Base.IntVid.feature b/tests/tck/features/match/Base.IntVid.feature
index d45577d3ca5..ae1318655ff 100644
--- a/tests/tck/features/match/Base.IntVid.feature
+++ b/tests/tck/features/match/Base.IntVid.feature
@@ -1045,8 +1045,8 @@ Feature: Basic match
       MATCH (v:player{name:"Tim Duncan"}) RETURN v.player AS vtag
       """
     Then the result should be, in any order:
-      | vtag                |
-      | {name:"Tim Duncan"} |
+      | vtag                          |
+      | {age: 42, name: "Tim Duncan"} |
     When executing query:
       """
       MATCH (v:player)-[]->(b) WHERE v.age > 30 RETURN v.player.name AS vname
diff --git a/tests/tck/features/match/Base.feature b/tests/tck/features/match/Base.feature
index 8295915e923..5ebdcef3dd8 100644
--- a/tests/tck/features/match/Base.feature
+++ b/tests/tck/features/match/Base.feature
@@ -1449,8 +1449,17 @@ Feature: Basic match
       MATCH (v:player{name:"Tim Duncan"}) RETURN v.player AS vtag
       """
     Then the result should be, in any order:
-      | vtag                |
-      | {name:"Tim Duncan"} |
+      | vtag                          |
+      | {age: 42, name: "Tim Duncan"} |
+    When executing query:
+      """
+      match (v)
+      where id(v) == "Tim Duncan"
+      return v.player, v.bachelor.name
+      """
+    Then the result should be, in any order:
+      | v.player                      | v.bachelor.name |
+      | {age: 42, name: "Tim Duncan"} | "Tim Duncan"    |
     When executing query:
       """
       MATCH (v:player)-[]->(b) WHERE v.age > 30 RETURN v.player.name AS vname