optimize go yield (#3974)

* convert schema function to EdgePropExpression

* add test case

Co-authored-by: Sophie <84560950+Sophie-Xie@users.noreply.github.com>

src/graph/validator/GoValidator.cpp

@@ -133,20 +133,22 @@ Status GoValidator::validateYield(YieldClause* yield) {
  auto& exprProps = goCtx_->exprProps;
 
  for (auto col : yield->columns()) {
+ const auto& colName = col->name();
  auto vertexExpr = ExpressionUtils::findAny(col->expr(), {Expression::Kind::kVertex});
  if (vertexExpr != nullptr &&
  static_cast<const VertexExpression*>(vertexExpr)->name() == "VERTEX") {
  return Status::SemanticError("`%s' is not support in go sentence.", col->toString().c_str());
  }
 
+ col->setExpr(rewriteVertexEdge2EdgeProp(col->expr()));
  col->setExpr(ExpressionUtils::rewriteLabelAttr2EdgeProp(col->expr()));
  NG_RETURN_IF_ERROR(ValidateUtil::invalidLabelIdentifiers(col->expr()));
 
  auto* colExpr = col->expr();
  auto typeStatus = deduceExprType(colExpr);
  NG_RETURN_IF_ERROR(typeStatus);
  auto type = typeStatus.value();
- outputs_.emplace_back(col->name(), type);
+ outputs_.emplace_back(colName, type);
  NG_RETURN_IF_ERROR(deduceProps(colExpr, exprProps, &tagIds_, &goCtx_->over.edgeTypes));
  }
 
@@ -183,6 +185,24 @@ void GoValidator::extractPropExprs(const Expression* expr,
  const_cast<Expression*>(expr)->accept(&visitor);
 }
 
+Expression* GoValidator::rewriteVertexEdge2EdgeProp(const Expression* expr) {
+ auto pool = qctx_->objPool();
+ const auto& name = expr->toString();
+ if (name == "id($^)" || name == "src(edge)") {
+ return EdgeSrcIdExpression::make(pool, "*");
+ }
+ if (name == "id($$)" || name == "dst(edge)") {
+ return EdgeDstIdExpression::make(pool, "*");
+ }
+ if (name == "rank(edge)") {
+ return EdgeRankExpression::make(pool, "*");
+ }
+ if (name == "type(edge)") {
+ return EdgeTypeExpression::make(pool, "*");
+ }
+ return const_cast<Expression*>(expr);
+}
+
 // Rewrites the property expression to corresponding Variable/Input expression
 // which get related property from previous plan node.
 Expression* GoValidator::rewrite2VarProp(const Expression* expr) {

src/graph/validator/GoValidator.h

@@ -40,6 +40,8 @@ class GoValidator final : public Validator {
 
  Expression* rewrite2VarProp(const Expression* expr);
 
+ Expression* rewriteVertexEdge2EdgeProp(const Expression* expr);
+
  private:
  std::unique_ptr<GoContext> goCtx_;
 

src/graph/validator/test/QueryValidatorTest.cpp

@@ -185,6 +185,23 @@ TEST_F(QueryValidatorTest, GoWithPipe) {
  PK::kStart};
  EXPECT_TRUE(checkResult(query, expected));
  }
+ {
+ std::string query = "YIELD \"1\" AS id | GO FROM $-.id OVER like YIELD id($$) as id";
+ std::vector<PlanNode::Kind> expected = {PK::kProject,
+ PK::kInnerJoin,
+ PK::kProject,
+ PK::kGetNeighbors,
+ PK::kDedup,
+ PK::kProject,
+ PK::kProject,
+ PK::kStart};
+ EXPECT_TRUE(checkResult(query, expected));
+ }
+ {
+ std::string query = "GO FROM 'Tim' OVER * YIELD id($$) as id";
+ std::vector<PlanNode::Kind> expected = {PK::kProject, PK::kGetNeighbors, PK::kStart};
+ EXPECT_TRUE(checkResult(query, expected));
+ }
  {
  std::string query =
  "GO 1 STEPS FROM \"1\" OVER like YIELD like._dst AS "

src/graph/visitor/DeducePropsVisitor.cpp

@@ -264,7 +264,14 @@ void DeducePropsVisitor::visit(EdgeExpression *expr) {
 }
 
 void DeducePropsVisitor::visitEdgePropExpr(PropertyExpression *expr) {
- auto status = qctx_->schemaMng()->toEdgeType(space_, expr->sym());
+ const auto &edgeName = expr->sym();
+ if (edgeName == "*") {
+ for (const auto &edgeType : *edgeTypes_) {
+ exprProps_->insertEdgeProp(edgeType, expr->prop());
+ }
+ return;
+ }
+ auto status = qctx_->schemaMng()->toEdgeType(space_, edgeName);
  if (!status.ok()) {
  status_ = std::move(status).status();
  return;

tests/tck/features/optimizer/PushLimitDownGetNeighborsRule.feature

@@ -22,6 +22,34 @@ Feature: Push Limit down rule
  | 6 | Limit | 7 | |
  | 7 | GetNeighbors | 0 | {"limit": "2"} |
  | 0 | Start | | |
+ When profiling query:
+ """
+ GO FROM "James Harden" OVER * YIELD id($$) as dst | Limit 2
+ """
+ Then the result should be, in any order:
+ | dst |
+ | "Russell Westbrook" |
+ | "Rockets" |
+ And the execution plan should be:
+ | id | name | dependencies | operator info |
+ | 5 | Project | 6 | |
+ | 6 | Limit | 7 | |
+ | 7 | GetNeighbors | 0 | {"limit": "2"} |
+ | 0 | Start | | |
+ When profiling query:
+ """
+ GO FROM "James Harden" OVER * YIELD id($^) as src | Limit 2
+ """
+ Then the result should be, in any order:
+ | src |
+ | "James Harden" |
+ | "James Harden" |
+ And the execution plan should be:
+ | id | name | dependencies | operator info |
+ | 5 | Project | 6 | |
+ | 6 | Limit | 7 | |
+ | 7 | GetNeighbors | 0 | {"limit": "2"} |
+ | 0 | Start | | |
 
  Scenario: push limit down to GetNeighbors with offset
  When profiling query: