topN push down for lookup

src/clients/storage/StorageClient.cpp

@@ -485,6 +485,7 @@ StorageRpcRespFuture<cpp2::LookupIndexResp> StorageClient::lookupIndex(
     bool isEdge,
     int32_t tagOrEdge,
     const std::vector<std::string>& returnCols,
+    std::vector<storage::cpp2::OrderBy> orderBy,
     int64_t limit) {
   // TODO(sky) : instead of isEdge and tagOrEdge to nebula::cpp2::SchemaID for graph layer.
   auto space = param.space;
@@ -516,6 +517,7 @@ StorageRpcRespFuture<cpp2::LookupIndexResp> StorageClient::lookupIndex(
     req.indices_ref() = spec;
     req.common_ref() = common;
     req.limit_ref() = limit;
+    req.order_by_ref() = orderBy;
   }
 
   return collectResponse(param.evb,

src/clients/storage/StorageClient.h

@@ -141,6 +141,7 @@ class StorageClient
       bool isEdge,
       int32_t tagOrEdge,
       const std::vector<std::string>& returnCols,
+      std::vector<storage::cpp2::OrderBy> orderBy,
       int64_t limit);
 
   StorageRpcRespFuture<cpp2::GetNeighborsResponse> lookupAndTraverse(

src/graph/executor/query/IndexScanExecutor.cpp

@@ -47,6 +47,7 @@ folly::Future<Status> IndexScanExecutor::indexScan() {
                     lookup->isEdge(),
                     lookup->schemaId(),
                     lookup->returnColumns(),
+                    lookup->orderBy(),
                     lookup->limit(qctx_))
       .via(runner())
       .thenValue([this](StorageRpcResponse<LookupIndexResp> &&rpcResp) {

src/graph/optimizer/CMakeLists.txt

@@ -50,6 +50,13 @@ nebula_add_library(
     rule/PushLimitDownScanAppendVerticesRule.cpp
     rule/GetEdgesTransformRule.cpp
     rule/PushLimitDownScanEdgesAppendVerticesRule.cpp
+    rule/PushTopNDownIndexScanRule.cpp
+    rule/PushTopNDownTagIndexFullScanRule.cpp
+    rule/PushTopNDownTagIndexPrefixScanRule.cpp
+    rule/PushTopNDownTagIndexRangeScanRule.cpp
+    rule/PushTopNDownEdgeIndexFullScanRule.cpp
+    rule/PushTopNDownEdgeIndexPrefixScanRule.cpp
+    rule/PushTopNDownEdgeIndexRangeScanRule.cpp
 )
 
 nebula_add_subdirectory(test)

src/graph/optimizer/OptimizerUtils.cpp

@@ -427,6 +427,7 @@ void OptimizerUtils::copyIndexScanData(const nebula::graph::IndexScan* from,
   to->setOrderBy(from->orderBy());
   to->setLimit(from->limit(qctx));
   to->setFilter(from->filter() == nullptr ? nullptr : from->filter()->clone());
+  to->setYieldColumns(from->yieldColumns());
 }
 
 Status OptimizerUtils::compareAndSwapBound(std::pair<Value, bool>& a, std::pair<Value, bool>& b) {

src/graph/optimizer/rule/PushTopNDownEdgeIndexFullScanRule.cpp

@@ -0,0 +1,110 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License.
+ */
+
+#include "graph/optimizer/rule/PushTopNDownEdgeIndexFullScanRule.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/planner/plan/Scan.h"
+
+using nebula::graph::EdgeIndexFullScan;
+using nebula::graph::PlanNode;
+using nebula::graph::Project;
+using nebula::graph::QueryContext;
+using nebula::graph::TopN;
+
+namespace nebula {
+namespace opt {
+
+std::unique_ptr<OptRule> PushTopNDownEdgeIndexFullScanRule::kInstance =
+    std::unique_ptr<PushTopNDownEdgeIndexFullScanRule>(new PushTopNDownEdgeIndexFullScanRule());
+
+PushTopNDownEdgeIndexFullScanRule::PushTopNDownEdgeIndexFullScanRule() {
+  RuleSet::QueryRules().addRule(this);
+}
+
+const Pattern &PushTopNDownEdgeIndexFullScanRule::pattern() const {
+  static Pattern pattern = Pattern::create(
+      graph::PlanNode::Kind::kTopN,
+      {Pattern::create(graph::PlanNode::Kind::kProject,
+                       {Pattern::create(graph::PlanNode::Kind::kEdgeIndexFullScan)})});
+  return pattern;
+}
+
+StatusOr<OptRule::TransformResult> PushTopNDownEdgeIndexFullScanRule::transform(
+    OptContext *octx, const MatchedResult &matched) const {
+  auto topNGroupNode = matched.node;
+  auto projectGroupNode = matched.dependencies.front().node;
+  auto indexScanGroupNode = matched.dependencies.front().dependencies.front().node;
+
+  const auto topN = static_cast<const TopN *>(topNGroupNode->node());
+  const auto project = static_cast<const Project *>(projectGroupNode->node());
+  const auto indexScan = static_cast<const EdgeIndexFullScan *>(indexScanGroupNode->node());
+
+  int64_t limitRows = topN->offset() + topN->count();
+
+  auto &factors = topN->factors();
+  auto projColNames = project->colNames();
+  const auto &yieldColumns = indexScan->yieldColumns();
+
+  std::unordered_map<std::string, std::string> namesMap;
+  for (auto yieldColumn : yieldColumns->columns()) {
+    if (yieldColumn->expr()->kind() == Expression::Kind::kEdgeProperty) {
+      const auto &propName = static_cast<EdgePropertyExpression *>(yieldColumn->expr())->prop();
+      namesMap[yieldColumn->name()] = propName;
+      continue;
+    }
+    return TransformResult::noTransform();
+  }
+  std::vector<storage::cpp2::OrderBy> orderBys;
+  orderBys.reserve(factors.size());
+
+  for (auto factor : factors) {
+    auto colName = projColNames[factor.first];
+    auto found = namesMap.find(colName);
+    if (found == namesMap.end()) {
+      return Status::Error();
+    }
+    storage::cpp2::OrderBy orderBy;
+    orderBy.prop_ref() = found->second;
+    orderBy.direction_ref() = factor.second == OrderFactor::OrderType::ASCEND
+                                  ? storage::cpp2::OrderDirection::ASCENDING
+                                  : storage::cpp2::OrderDirection::DESCENDING;
+    orderBys.emplace_back(orderBy);
+  }
+
+  auto newTopN = static_cast<TopN *>(topN->clone());
+  auto newtopNGroupNode = OptGroupNode::create(octx, newTopN, topNGroupNode->group());
+
+  auto newProject = static_cast<Project *>(project->clone());
+  auto newProjectGroup = OptGroup::create(octx);
+  auto newProjectGroupNode = newProjectGroup->makeGroupNode(newProject);
+
+  auto newIndexScan = static_cast<EdgeIndexFullScan *>(indexScan->clone());
+  newIndexScan->setLimit(limitRows);
+  newIndexScan->setOrderBy(orderBys);
+  auto newIndexScanGroup = OptGroup::create(octx);
+  auto newIndexScanGroupNode = newIndexScanGroup->makeGroupNode(newIndexScan);
+
+  newtopNGroupNode->dependsOn(newProjectGroup);
+  newProjectGroupNode->dependsOn(newIndexScanGroup);
+  for (auto dep : indexScanGroupNode->dependencies()) {
+    newIndexScanGroupNode->dependsOn(dep);
+  }
+
+  TransformResult result;
+  result.eraseAll = true;
+  result.newGroupNodes.emplace_back(newtopNGroupNode);
+  return result;
+}
+
+std::string PushTopNDownEdgeIndexFullScanRule::toString() const {
+  return "PushTopNDownEdgeIndexFullScanRule";
+}
+
+}  // namespace opt
+}  // namespace nebula

src/graph/optimizer/rule/PushTopNDownEdgeIndexFullScanRule.h

@@ -0,0 +1,29 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License.
+ */
+
+#pragma once
+
+#include "graph/optimizer/OptRule.h"
+
+namespace nebula {
+namespace opt {
+
+class PushTopNDownEdgeIndexFullScanRule 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:
+  PushTopNDownEdgeIndexFullScanRule();
+
+  static std::unique_ptr<OptRule> kInstance;
+};
+
+}  // namespace opt
+}  // namespace nebula

src/graph/optimizer/rule/PushTopNDownEdgeIndexPrefixScanRule.cpp

@@ -0,0 +1,110 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License.
+ */
+
+#include "graph/optimizer/rule/PushTopNDownEdgeIndexPrefixScanRule.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/planner/plan/Scan.h"
+
+using nebula::graph::EdgeIndexPrefixScan;
+using nebula::graph::PlanNode;
+using nebula::graph::Project;
+using nebula::graph::QueryContext;
+using nebula::graph::TopN;
+
+namespace nebula {
+namespace opt {
+
+std::unique_ptr<OptRule> PushTopNDownEdgeIndexPrefixScanRule::kInstance =
+    std::unique_ptr<PushTopNDownEdgeIndexPrefixScanRule>(new PushTopNDownEdgeIndexPrefixScanRule());
+
+PushTopNDownEdgeIndexPrefixScanRule::PushTopNDownEdgeIndexPrefixScanRule() {
+  RuleSet::QueryRules().addRule(this);
+}
+
+const Pattern &PushTopNDownEdgeIndexPrefixScanRule::pattern() const {
+  static Pattern pattern = Pattern::create(
+      graph::PlanNode::Kind::kTopN,
+      {Pattern::create(graph::PlanNode::Kind::kProject,
+                       {Pattern::create(graph::PlanNode::Kind::kEdgeIndexPrefixScan)})});
+  return pattern;
+}
+
+StatusOr<OptRule::TransformResult> PushTopNDownEdgeIndexPrefixScanRule::transform(
+    OptContext *octx, const MatchedResult &matched) const {
+  auto topNGroupNode = matched.node;
+  auto projectGroupNode = matched.dependencies.front().node;
+  auto indexScanGroupNode = matched.dependencies.front().dependencies.front().node;
+
+  const auto topN = static_cast<const TopN *>(topNGroupNode->node());
+  const auto project = static_cast<const Project *>(projectGroupNode->node());
+  const auto indexScan = static_cast<const EdgeIndexPrefixScan *>(indexScanGroupNode->node());
+
+  int64_t limitRows = topN->offset() + topN->count();
+
+  auto &factors = topN->factors();
+  auto projColNames = project->colNames();
+  const auto &yieldColumns = indexScan->yieldColumns();
+
+  std::unordered_map<std::string, std::string> namesMap;
+  for (auto yieldColumn : yieldColumns->columns()) {
+    if (yieldColumn->expr()->kind() == Expression::Kind::kEdgeProperty) {
+      const auto &propName = static_cast<EdgePropertyExpression *>(yieldColumn->expr())->prop();
+      namesMap[yieldColumn->name()] = propName;
+      continue;
+    }
+    return TransformResult::noTransform();
+  }
+  std::vector<storage::cpp2::OrderBy> orderBys;
+  orderBys.reserve(factors.size());
+
+  for (auto factor : factors) {
+    auto colName = projColNames[factor.first];
+    auto found = namesMap.find(colName);
+    if (found == namesMap.end()) {
+      return Status::Error();
+    }
+    storage::cpp2::OrderBy orderBy;
+    orderBy.prop_ref() = found->second;
+    orderBy.direction_ref() = factor.second == OrderFactor::OrderType::ASCEND
+                                  ? storage::cpp2::OrderDirection::ASCENDING
+                                  : storage::cpp2::OrderDirection::DESCENDING;
+    orderBys.emplace_back(orderBy);
+  }
+
+  auto newTopN = static_cast<TopN *>(topN->clone());
+  auto newtopNGroupNode = OptGroupNode::create(octx, newTopN, topNGroupNode->group());
+
+  auto newProject = static_cast<Project *>(project->clone());
+  auto newProjectGroup = OptGroup::create(octx);
+  auto newProjectGroupNode = newProjectGroup->makeGroupNode(newProject);
+
+  auto newIndexScan = static_cast<EdgeIndexPrefixScan *>(indexScan->clone());
+  newIndexScan->setLimit(limitRows);
+  newIndexScan->setOrderBy(orderBys);
+  auto newIndexScanGroup = OptGroup::create(octx);
+  auto newIndexScanGroupNode = newIndexScanGroup->makeGroupNode(newIndexScan);
+
+  newtopNGroupNode->dependsOn(newProjectGroup);
+  newProjectGroupNode->dependsOn(newIndexScanGroup);
+  for (auto dep : indexScanGroupNode->dependencies()) {
+    newIndexScanGroupNode->dependsOn(dep);
+  }
+
+  TransformResult result;
+  result.eraseAll = true;
+  result.newGroupNodes.emplace_back(newtopNGroupNode);
+  return result;
+}
+
+std::string PushTopNDownEdgeIndexPrefixScanRule::toString() const {
+  return "PushTopNDownEdgeIndexPrefixScanRule";
+}
+
+}  // namespace opt
+}  // namespace nebula

src/graph/optimizer/rule/PushTopNDownEdgeIndexPrefixScanRule.h

@@ -0,0 +1,29 @@
+/* Copyright (c) 2021 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License.
+ */
+
+#pragma once
+
+#include "graph/optimizer/OptRule.h"
+
+namespace nebula {
+namespace opt {
+
+class PushTopNDownEdgeIndexPrefixScanRule 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:
+  PushTopNDownEdgeIndexPrefixScanRule();
+
+  static std::unique_ptr<OptRule> kInstance;
+};
+
+}  // namespace opt
+}  // namespace nebula