Feature/index scan limit push down

src/clients/storage/GraphStorageClient.cpp

@@ -491,6 +491,7 @@ folly::SemiFuture<StorageRpcResponse<cpp2::LookupIndexResp>> GraphStorageClient:
     bool isEdge,
     int32_t tagOrEdge,
     const std::vector<std::string>& returnCols,
+    int64_t limit,
     folly::EventBase* evb) {
   // TODO(sky) : instead of isEdge and tagOrEdge to nebula::cpp2::SchemaID for graph layer.
   auto status = getHostParts(space);
@@ -520,6 +521,7 @@ folly::SemiFuture<StorageRpcResponse<cpp2::LookupIndexResp>> GraphStorageClient:
     spec.set_schema_id(schemaId);
     req.set_indices(spec);
     req.set_common(common);
+    req.set_limit(limit);
   }
 
   return collectResponse(

src/clients/storage/GraphStorageClient.h

@@ -141,6 +141,7 @@ class GraphStorageClient : public StorageClientBase<cpp2::GraphStorageServiceAsy
       bool isEdge,
       int32_t tagOrEdge,
       const std::vector<std::string>& returnCols,
+      int64_t limit,
       folly::EventBase* evb = nullptr);
 
   folly::SemiFuture<StorageRpcResponse<cpp2::GetNeighborsResponse>> lookupAndTraverse(

src/graph/context/ast/QueryAstContext.h

@@ -111,7 +111,7 @@ struct LookupContext final : public AstContext {
   bool dedup{false};
   bool isEmptyResultSet{false};
   int32_t schemaId{-1};
-  int32_t limit{-1};
+  int64_t limit{-1};
   Expression* filter{nullptr};
   YieldColumns* yieldExpr{nullptr};
   std::vector<std::string> idxReturnCols;

src/graph/executor/query/IndexScanExecutor.cpp

@@ -43,7 +43,8 @@ folly::Future<Status> IndexScanExecutor::indexScan() {
                     ictxs,
                     lookup->isEdge(),
                     lookup->schemaId(),
-                    lookup->returnColumns())
+                    lookup->returnColumns(),
+                    lookup->limit())
       .via(runner())
       .thenValue([this](StorageRpcResponse<LookupIndexResp> &&rpcResp) {
         addStats(rpcResp, otherStats_);

src/graph/optimizer/CMakeLists.txt

@@ -20,7 +20,7 @@ nebula_add_library(
     rule/MergeGetNbrsAndDedupRule.cpp
     rule/MergeGetNbrsAndProjectRule.cpp
     rule/IndexScanRule.cpp
-    rule/LimitPushDownRule.cpp
+    rule/PushLimitDownGetNeighborsRule.cpp
     rule/TopNRule.cpp
     rule/PushFilterDownAggregateRule.cpp
     rule/PushFilterDownProjectRule.cpp
@@ -33,6 +33,7 @@ nebula_add_library(
     rule/IndexFullScanBaseRule.cpp
     rule/TagIndexFullScanRule.cpp
     rule/EdgeIndexFullScanRule.cpp
+    rule/PushLimitDownIndexScanRule.cpp
 )
 
 nebula_add_subdirectory(test)

src/graph/optimizer/OptRule.cpp

@@ -32,17 +32,20 @@ const PlanNode *MatchedResult::planNode(const std::vector<int32_t> &pos) const {
   return DCHECK_NOTNULL(result->node)->node();
 }
 
-Pattern Pattern::create(graph::PlanNode::Kind kind, std::initializer_list<Pattern> patterns) {
+Pattern Pattern::create(graph::PlanNode::Kind kind,
+                        std::initializer_list<Pattern> patterns,
+                        bool polymorphism) {
   Pattern pattern;
   pattern.kind_ = kind;
   for (auto &p : patterns) {
     pattern.dependencies_.emplace_back(p);
   }
+  pattern.polymorphism_ = polymorphism;
   return pattern;
 }
 
 StatusOr<MatchedResult> Pattern::match(const OptGroupNode *groupNode) const {
-  if (groupNode->node()->kind() != kind_) {
+  if (!(polymorphism_ ? groupNode->node()->isA(kind_) : groupNode->node()->kind() == kind_)) {
     return Status::Error();
   }
 

src/graph/optimizer/OptRule.h

@@ -45,7 +45,9 @@ struct MatchedResult {
 
 class Pattern final {
  public:
-  static Pattern create(graph::PlanNode::Kind kind, std::initializer_list<Pattern> patterns = {});
+  static Pattern create(graph::PlanNode::Kind kind,
+                        std::initializer_list<Pattern> patterns = {},
+                        bool polymorphism = false);
 
   StatusOr<MatchedResult> match(const OptGroupNode *groupNode) const;
 
@@ -55,6 +57,8 @@ class Pattern final {
 
   graph::PlanNode::Kind kind_;
   std::vector<Pattern> dependencies_;
+  // Whether enable plymorphsim plan node match
+  bool polymorphism_{false};
 };
 
 class OptRule {

src/graph/optimizer/rule/LimitPushDownRule.cpp → src/graph/optimizer/rule/PushLimitDownGetNeighborsRule.cpp

@@ -4,7 +4,7 @@
  * attached with Common Clause Condition 1.0, found in the LICENSES directory.
  */
 
-#include "graph/optimizer/rule/LimitPushDownRule.h"
+#include "graph/optimizer/rule/PushLimitDownGetNeighborsRule.h"
 
 #include "common/expression/BinaryExpression.h"
 #include "common/expression/ConstantExpression.h"
@@ -27,20 +27,22 @@ using nebula::graph::QueryContext;
 namespace nebula {
 namespace opt {
 
-std::unique_ptr<OptRule> LimitPushDownRule::kInstance =
-    std::unique_ptr<LimitPushDownRule>(new LimitPushDownRule());
+std::unique_ptr<OptRule> PushLimitDownGetNeighborsRule::kInstance =
+    std::unique_ptr<PushLimitDownGetNeighborsRule>(new PushLimitDownGetNeighborsRule());
 
-LimitPushDownRule::LimitPushDownRule() { RuleSet::QueryRules().addRule(this); }
+PushLimitDownGetNeighborsRule::PushLimitDownGetNeighborsRule() {
+  RuleSet::QueryRules().addRule(this);
+}
 
-const Pattern &LimitPushDownRule::pattern() const {
+const Pattern &PushLimitDownGetNeighborsRule::pattern() const {
   static Pattern pattern =
       Pattern::create(graph::PlanNode::Kind::kLimit,
                       {Pattern::create(graph::PlanNode::Kind::kProject,
                                        {Pattern::create(graph::PlanNode::Kind::kGetNeighbors)})});
   return pattern;
 }
 
-StatusOr<OptRule::TransformResult> LimitPushDownRule::transform(
+StatusOr<OptRule::TransformResult> PushLimitDownGetNeighborsRule::transform(
     OptContext *octx, const MatchedResult &matched) const {
   auto limitGroupNode = matched.node;
   auto projGroupNode = matched.dependencies.front().node;
@@ -79,7 +81,9 @@ StatusOr<OptRule::TransformResult> LimitPushDownRule::transform(
   return result;
 }
 
-std::string LimitPushDownRule::toString() const { return "LimitPushDownRule"; }
+std::string PushLimitDownGetNeighborsRule::toString() const {
+  return "PushLimitDownGetNeighborsRule";
+}
 
 }  // namespace opt
 }  // namespace nebula

src/graph/optimizer/rule/LimitPushDownRule.h → src/graph/optimizer/rule/PushLimitDownGetNeighborsRule.h

@@ -4,8 +4,7 @@
  * attached with Common Clause Condition 1.0, found in the LICENSES directory.
  */
 
-#ifndef GRAPH_OPTIMIZER_RULE_LIMITPUSHDOWN_H_
-#define GRAPH_OPTIMIZER_RULE_LIMITPUSHDOWN_H_
+#pragma once
 
 #include <memory>
 
@@ -14,7 +13,7 @@
 namespace nebula {
 namespace opt {
 
-class LimitPushDownRule final : public OptRule {
+class PushLimitDownGetNeighborsRule final : public OptRule {
  public:
   const Pattern &pattern() const override;
 
@@ -24,12 +23,10 @@ class LimitPushDownRule final : public OptRule {
   std::string toString() const override;
 
  private:
-  LimitPushDownRule();
+  PushLimitDownGetNeighborsRule();
 
   static std::unique_ptr<OptRule> kInstance;
 };
 
 }  // namespace opt
 }  // namespace nebula
-
-#endif  // GRAPH_OPTIMIZER_RULE_LIMITPUSHDOWN_H_

src/graph/optimizer/rule/PushLimitDownIndexScanRule.cpp

@@ -0,0 +1,85 @@
+/* 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.
+ */
+
+#include "graph/optimizer/rule/PushLimitDownIndexScanRule.h"
+
+#include "common/expression/BinaryExpression.h"
+#include "common/expression/ConstantExpression.h"
+#include "common/expression/Expression.h"
+#include "common/expression/FunctionCallExpression.h"
+#include "common/expression/LogicalExpression.h"
+#include "common/expression/UnaryExpression.h"
+#include "graph/optimizer/OptContext.h"
+#include "graph/optimizer/OptGroup.h"
+#include "graph/planner/plan/PlanNode.h"
+#include "graph/planner/plan/Query.h"
+#include "graph/visitor/ExtractFilterExprVisitor.h"
+
+using nebula::graph::IndexScan;
+using nebula::graph::Limit;
+using nebula::graph::PlanNode;
+using nebula::graph::Project;
+using nebula::graph::QueryContext;
+
+namespace nebula {
+namespace opt {
+
+std::unique_ptr<OptRule> PushLimitDownIndexScanRule::kInstance =
+    std::unique_ptr<PushLimitDownIndexScanRule>(new PushLimitDownIndexScanRule());
+
+PushLimitDownIndexScanRule::PushLimitDownIndexScanRule() { RuleSet::QueryRules().addRule(this); }
+
+const Pattern &PushLimitDownIndexScanRule::pattern() const {
+  static Pattern pattern = Pattern::create(
+      graph::PlanNode::Kind::kLimit,
+      {Pattern::create(graph::PlanNode::Kind::kProject,
+                       {Pattern::create(graph::PlanNode::Kind::kIndexScan, {}, true)})});
+  return pattern;
+}
+
+StatusOr<OptRule::TransformResult> PushLimitDownIndexScanRule::transform(
+    OptContext *octx, const MatchedResult &matched) const {
+  auto limitGroupNode = matched.node;
+  auto projGroupNode = matched.dependencies.front().node;
+  auto indexScanGroupNode = matched.dependencies.front().dependencies.front().node;
+
+  const auto limit = limitGroupNode->node()->asNode<const Limit>();
+  const auto proj = projGroupNode->node()->asNode<const Project>();
+  const auto indexScan = indexScanGroupNode->node()->asNode<const IndexScan>();
+
+  int64_t limitRows = limit->offset() + limit->count();
+  if (indexScan->limit() >= 0 && limitRows >= indexScan->limit()) {
+    return TransformResult::noTransform();
+  }
+
+  auto newLimit = static_cast<Limit *>(limit->clone());
+  auto newLimitGroupNode = OptGroupNode::create(octx, newLimit, limitGroupNode->group());
+
+  auto newProj = static_cast<Project *>(proj->clone());
+  auto newProjGroup = OptGroup::create(octx);
+  auto newProjGroupNode = newProjGroup->makeGroupNode(newProj);
+
+  auto newIndexScan = static_cast<IndexScan *>(indexScan->clone());
+  newIndexScan->setLimit(limitRows);
+  auto newIndexScanGroup = OptGroup::create(octx);
+  auto newIndexScanGroupNode = newIndexScanGroup->makeGroupNode(newIndexScan);
+
+  newLimitGroupNode->dependsOn(newProjGroup);
+  newProjGroupNode->dependsOn(newIndexScanGroup);
+  for (auto dep : indexScanGroupNode->dependencies()) {
+    newIndexScanGroupNode->dependsOn(dep);
+  }
+
+  TransformResult result;
+  result.eraseAll = true;
+  result.newGroupNodes.emplace_back(newLimitGroupNode);
+  return result;
+}
+
+std::string PushLimitDownIndexScanRule::toString() const { return "PushLimitDownIndexScanRule"; }
+
+}  // namespace opt
+}  // namespace nebula

src/graph/optimizer/rule/PushLimitDownIndexScanRule.h

@@ -0,0 +1,32 @@
+/* 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.
+ */
+
+#pragma once
+
+#include <memory>
+
+#include "graph/optimizer/OptRule.h"
+
+namespace nebula {
+namespace opt {
+
+class PushLimitDownIndexScanRule final : public OptRule {
+ public:
+  const Pattern &pattern() const override;
+
+  StatusOr<OptRule::TransformResult> transform(OptContext *ctx,
+                                               const MatchedResult &matched) const override;
+
+  std::string toString() const override;
+
+ private:
+  PushLimitDownIndexScanRule();
+
+  static std::unique_ptr<OptRule> kInstance;
+};
+
+}  // namespace opt
+}  // namespace nebula

src/graph/planner/ngql/LookupPlanner.cpp

@@ -6,7 +6,6 @@
 
 #include "graph/planner/ngql/LookupPlanner.h"
 
-
 #include "common/base/Base.h"
 #include "graph/context/ast/QueryAstContext.h"
 #include "graph/planner/Planner.h"
@@ -31,23 +30,31 @@ StatusOr<SubPlan> LookupPlanner::transform(AstContext* astCtx) {
   auto from = static_cast<const LookupSentence*>(lookupCtx->sentence)->from();
   SubPlan plan;
   if (lookupCtx->isEdge) {
-    plan.tail = EdgeIndexFullScan::make(qctx,
-                                        nullptr,
-                                        from,
-                                        lookupCtx->space.id,
-                                        {},
-                                        lookupCtx->idxReturnCols,
-                                        lookupCtx->schemaId,
-                                        lookupCtx->isEmptyResultSet);
+    auto* edgeIndexFullScan = EdgeIndexFullScan::make(qctx,
+                                                      nullptr,
+                                                      from,
+                                                      lookupCtx->space.id,
+                                                      {},
+                                                      lookupCtx->idxReturnCols,
+                                                      lookupCtx->schemaId,
+                                                      lookupCtx->isEmptyResultSet);
+    if (lookupCtx->limit >= 0) {
+      edgeIndexFullScan->setLimit(lookupCtx->limit);
+    }
+    plan.tail = edgeIndexFullScan;
   } else {
-    plan.tail = TagIndexFullScan::make(qctx,
-                                       nullptr,
-                                       from,
-                                       lookupCtx->space.id,
-                                       {},
-                                       lookupCtx->idxReturnCols,
-                                       lookupCtx->schemaId,
-                                       lookupCtx->isEmptyResultSet);
+    auto* tagIndexFullScan = TagIndexFullScan::make(qctx,
+                                                    nullptr,
+                                                    from,
+                                                    lookupCtx->space.id,
+                                                    {},
+                                                    lookupCtx->idxReturnCols,
+                                                    lookupCtx->schemaId,
+                                                    lookupCtx->isEmptyResultSet);
+    if (lookupCtx->limit >= 0) {
+      tagIndexFullScan->setLimit(lookupCtx->limit);
+    }
+    plan.tail = tagIndexFullScan;
   }
   plan.tail->setColNames(lookupCtx->idxColNames);