From 2a0f9013d8c06e0efdc8a803897bc57fd887e3f0 Mon Sep 17 00:00:00 2001
From: "kyle.cao" <kyle.cao@vesoft.com>
Date: Thu, 21 Oct 2021 15:01:38 +0800
Subject: [PATCH] fix PushLimitDownProjRule (#3167)

---
 .../rule/PushLimitDownProjectRule.cpp         |  8 ++--
 .../PushLimitDownProjectRule.feature          | 40 +++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 tests/tck/features/optimizer/PushLimitDownProjectRule.feature

diff --git a/src/graph/optimizer/rule/PushLimitDownProjectRule.cpp b/src/graph/optimizer/rule/PushLimitDownProjectRule.cpp
index 4bc6463c5bb..21f3d65420b 100644
--- a/src/graph/optimizer/rule/PushLimitDownProjectRule.cpp
+++ b/src/graph/optimizer/rule/PushLimitDownProjectRule.cpp
@@ -49,15 +49,17 @@ StatusOr<OptRule::TransformResult> PushLimitDownProjectRule::transform(
   auto newLimit = static_cast<Limit *>(limit->clone());
   auto newLimitGroup = OptGroup::create(octx);
   auto newLimitGroupNode = newLimitGroup->makeGroupNode(newLimit);
+  auto projInputVar = proj->inputVar();
   newLimit->setOutputVar(proj->outputVar());
-  newLimit->setInputVar(proj->inputVar());
-  //  newLimit->setColNames(proj->colNames());
+  newLimit->setInputVar(projInputVar);
+  auto *varPtr = octx->qctx()->symTable()->getVar(projInputVar);
+  DCHECK(!!varPtr);
+  newLimit->setColNames(varPtr->colNames);
 
   auto newProj = static_cast<Project *>(proj->clone());
   auto newProjGroupNode = OptGroupNode::create(octx, newProj, limitGroupNode->group());
   newProj->setOutputVar(limit->outputVar());
   newProj->setInputVar(newLimit->outputVar());
-  //  newProj->setColNames(limit->colNames());
 
   newProjGroupNode->dependsOn(const_cast<OptGroup *>(newLimitGroupNode->group()));
   for (auto dep : projGroupNode->dependencies()) {
diff --git a/tests/tck/features/optimizer/PushLimitDownProjectRule.feature b/tests/tck/features/optimizer/PushLimitDownProjectRule.feature
new file mode 100644
index 00000000000..ea36f95c118
--- /dev/null
+++ b/tests/tck/features/optimizer/PushLimitDownProjectRule.feature
@@ -0,0 +1,40 @@
+# Copyright (c) 2021 vesoft inc. All rights reserved.
+#
+# This source code is licensed under Apache 2.0 License,
+# attached with Common Clause Condition 1.0, found in the LICENSES directory.
+Feature: Push Limit down project rule
+
+  Background:
+    Given a graph with space named "nba"
+
+  Scenario: push limit down to Project
+    When profiling query:
+      """
+      MATCH p=(v:player)-[]->(n)
+      WHERE id(v)=="Tim Duncan" and n.age>30
+      RETURN p LIMIT 100
+      """
+    Then the result should be, in any order:
+      | p                                                                                                                                                                                                                                 |
+      | <("Tim Duncan" :bachelor{name: "Tim Duncan", speciality: "psychology"} :player{age: 42, name: "Tim Duncan"})-[:teammate@0 {end_year: 2016, start_year: 2010}]->("Danny Green" :player{age: 31, name: "Danny Green"})>             |
+      | <("Tim Duncan" :bachelor{name: "Tim Duncan", speciality: "psychology"} :player{age: 42, name: "Tim Duncan"})-[:teammate@0 {end_year: 2016, start_year: 2015}]->("LaMarcus Aldridge" :player{age: 33, name: "LaMarcus Aldridge"})> |
+      | <("Tim Duncan" :bachelor{name: "Tim Duncan", speciality: "psychology"} :player{age: 42, name: "Tim Duncan"})-[:teammate@0 {end_year: 2016, start_year: 2002}]->("Manu Ginobili" :player{age: 41, name: "Manu Ginobili"})>         |
+      | <("Tim Duncan" :bachelor{name: "Tim Duncan", speciality: "psychology"} :player{age: 42, name: "Tim Duncan"})-[:teammate@0 {end_year: 2016, start_year: 2001}]->("Tony Parker" :player{age: 36, name: "Tony Parker"})>             |
+      | <("Tim Duncan" :bachelor{name: "Tim Duncan", speciality: "psychology"} :player{age: 42, name: "Tim Duncan"})-[:like@0 {likeness: 95}]->("Manu Ginobili" :player{age: 41, name: "Manu Ginobili"})>                                 |
+      | <("Tim Duncan" :bachelor{name: "Tim Duncan", speciality: "psychology"} :player{age: 42, name: "Tim Duncan"})-[:like@0 {likeness: 95}]->("Tony Parker" :player{age: 36, name: "Tony Parker"})>                                     |
+    And the execution plan should be:
+      | id | name         | dependencies | operator info |
+      | 18 | DataCollect  | 26           |               |
+      | 26 | Project      | 25           |               |
+      | 25 | Limit        | 20           |               |
+      | 20 | Filter       | 13           |               |
+      | 13 | Project      | 12           |               |
+      | 12 | InnerJoin    | 11           |               |
+      | 11 | Project      | 22           |               |
+      | 22 | GetVertices  | 7            |               |
+      | 7  | Filter       | 6            |               |
+      | 6  | Project      | 5            |               |
+      | 5  | Filter       | 24           |               |
+      | 24 | GetNeighbors | 1            |               |
+      | 1  | PassThrough  | 0            |               |
+      | 0  | Start        |              |               |