Add clone function for GetNeighbors/IndexScan plan node (#311)

* Add clone function for GetNeighbors plan node

* Move index scan clone to plan node

Co-authored-by: dutor <440396+dutor@users.noreply.github.com>
Co-authored-by: cpw <13495049+CPWstatic@users.noreply.github.com>

src/optimizer/rule/IndexScanRule.cpp

@@ -10,6 +10,8 @@
 #include "planner/PlanNode.h"
 #include "planner/Query.h"
 
+using nebula::graph::IndexScan;
+
 namespace nebula {
 namespace opt {
 std::unique_ptr<OptRule> IndexScanRule::kInstance =
@@ -39,7 +41,7 @@ Status IndexScanRule::transform(graph::QueryContext *qctx,
     ret = createIndexQueryCtx(iqctx, kind, items, qctx, groupExpr);
     NG_RETURN_IF_ERROR(ret);
 
-    auto newIN = cloneIndexScan(qctx, groupExpr);
+    auto newIN = static_cast<const IndexScan*>(groupExpr->node())->clone(qctx);
     newIN->setIndexQueryContext(std::move(iqctx));
     auto newGroupExpr = OptGroupExpr::create(qctx, newIN, groupExpr->group());
     if (groupExpr->dependencies().size() != 1) {
@@ -234,21 +236,6 @@ Status IndexScanRule::boundValue(const FilterItem& item,
     return Status::OK();
 }
 
-IndexScan* IndexScanRule::cloneIndexScan(graph::QueryContext *qctx,
-                                         const OptGroupExpr *groupExpr) const {
-    auto in = static_cast<const IndexScan *>(groupExpr->node());
-    auto ctx = std::make_unique<std::vector<storage::cpp2::IndexQueryContext>>();
-    auto returnCols = std::make_unique<std::vector<std::string>>(*in->returnColumns());
-    auto indexScan = IndexScan::make(qctx,
-                                     nullptr,
-                                     in->space(),
-                                     std::move(ctx),
-                                     std::move(returnCols),
-                                     in->isEdge(),
-                                     in->schemaId());
-    return indexScan;
-}
-
 bool IndexScanRule::isEdge(const OptGroupExpr *groupExpr) const {
     auto in = static_cast<const IndexScan *>(groupExpr->node());
     return in->isEdge();

src/optimizer/rule/IndexScanRule.h

@@ -130,8 +130,6 @@ private:
                       const meta::cpp2::ColumnDef& col,
                       Value& begin, Value& end) const;
 
-    IndexScan* cloneIndexScan(graph::QueryContext *qctx, const OptGroupExpr *groupExpr) const;
-
     bool isEdge(const OptGroupExpr *groupExpr) const;
 
     int32_t schemaId(const OptGroupExpr *groupExpr) const;

src/optimizer/rule/PushFilterDownGetNbrsRule.cpp

@@ -75,7 +75,7 @@ Status PushFilterDownGetNbrsRule::transform(QueryContext *qctx,
         newGNFilter = logicExpr.encode();
     }
 
-    auto newGN = cloneGetNbrs(qctx, gn);
+    auto newGN = gn->clone(qctx);
     newGN->setFilter(newGNFilter);
 
     OptGroupExpr *newGroupExpr = nullptr;
@@ -121,43 +121,5 @@ std::pair<bool, const OptGroupExpr *> PushFilterDownGetNbrsRule::findMatchedGrou
     return std::make_pair(false, nullptr);
 }
 
-GetNeighbors *PushFilterDownGetNbrsRule::cloneGetNbrs(QueryContext *qctx,
-                                                      const GetNeighbors *gn) const {
-    auto newGN = GetNeighbors::make(qctx, nullptr, gn->space());
-    newGN->setSrc(gn->src());
-    newGN->setEdgeTypes(gn->edgeTypes());
-    newGN->setEdgeDirection(gn->edgeDirection());
-    newGN->setDedup(gn->dedup());
-    newGN->setRandom(gn->random());
-    newGN->setLimit(gn->limit());
-    newGN->setInputVar(gn->inputVar());
-    newGN->setOutputVar(gn->outputVar());
-
-    if (gn->vertexProps()) {
-        auto vertexProps = *gn->vertexProps();
-        auto vertexPropsPtr = std::make_unique<decltype(vertexProps)>(std::move(vertexProps));
-        newGN->setVertexProps(std::move(vertexPropsPtr));
-    }
-
-    if (gn->edgeProps()) {
-        auto edgeProps = *gn->edgeProps();
-        auto edgePropsPtr = std::make_unique<decltype(edgeProps)>(std::move(edgeProps));
-        newGN->setEdgeProps(std::move(edgePropsPtr));
-    }
-
-    if (gn->statProps()) {
-        auto statProps = *gn->statProps();
-        auto statPropsPtr = std::make_unique<decltype(statProps)>(std::move(statProps));
-        newGN->setStatProps(std::move(statPropsPtr));
-    }
-
-    if (gn->exprs()) {
-        auto exprs = *gn->exprs();
-        auto exprsPtr = std::make_unique<decltype(exprs)>(std::move(exprs));
-        newGN->setExprs(std::move(exprsPtr));
-    }
-    return newGN;
-}
-
 }   // namespace opt
 }   // namespace nebula

src/optimizer/rule/PushFilterDownGetNbrsRule.h

@@ -31,9 +31,6 @@ public:
 private:
     PushFilterDownGetNbrsRule();
 
-    graph::GetNeighbors *cloneGetNbrs(graph::QueryContext *qctx,
-                                      const graph::GetNeighbors *getNbrs) const;
-
     std::pair<bool, const OptGroupExpr *> findMatchedGroupExpr(const OptGroupExpr *groupExpr) const;
 };
 

src/planner/Query.cpp

@@ -27,6 +27,44 @@ std::unique_ptr<cpp2::PlanNodeDescription> Explore::explain() const {
     return desc;
 }
 
+GetNeighbors* GetNeighbors::clone(QueryContext* qctx) const {
+    auto newGN = GetNeighbors::make(qctx, nullptr, space_);
+    newGN->setSrc(qctx->objPool()->add(src_->clone().release()));
+    newGN->setEdgeTypes(edgeTypes_);
+    newGN->setEdgeDirection(edgeDirection_);
+    newGN->setDedup(dedup_);
+    newGN->setRandom(random_);
+    newGN->setLimit(limit_);
+    newGN->setOrderBy(orderBy_);
+    newGN->setInputVar(inputVar());
+    newGN->setOutputVar(outputVar());
+
+    if (vertexProps_) {
+        auto vertexProps = *vertexProps_;
+        auto vertexPropsPtr = std::make_unique<decltype(vertexProps)>(vertexProps);
+        newGN->setVertexProps(std::move(vertexPropsPtr));
+    }
+
+    if (edgeProps_) {
+        auto edgeProps = *edgeProps_;
+        auto edgePropsPtr = std::make_unique<decltype(edgeProps)>(std::move(edgeProps));
+        newGN->setEdgeProps(std::move(edgePropsPtr));
+    }
+
+    if (statProps_) {
+        auto statProps = *statProps_;
+        auto statPropsPtr = std::make_unique<decltype(statProps)>(std::move(statProps));
+        newGN->setStatProps(std::move(statPropsPtr));
+    }
+
+    if (exprs_) {
+        auto exprs = *exprs_;
+        auto exprsPtr = std::make_unique<decltype(exprs)>(exprs);
+        newGN->setExprs(std::move(exprsPtr));
+    }
+    return newGN;
+}
+
 std::unique_ptr<cpp2::PlanNodeDescription> GetNeighbors::explain() const {
     auto desc = Explore::explain();
     addDescription("src", src_ ? src_->toString() : "", desc.get());
@@ -64,6 +102,13 @@ std::unique_ptr<cpp2::PlanNodeDescription> GetEdges::explain() const {
     return desc;
 }
 
+IndexScan* IndexScan::clone(QueryContext* qctx) const {
+    auto ctx = std::make_unique<std::vector<storage::cpp2::IndexQueryContext>>();
+    auto returnCols = std::make_unique<std::vector<std::string>>(*returnColumns());
+    return IndexScan::make(
+        qctx, nullptr, space(), std::move(ctx), std::move(returnCols), isEdge(), schemaId());
+}
+
 std::unique_ptr<cpp2::PlanNodeDescription> IndexScan::explain() const {
     auto desc = Explore::explain();
     // TODO
@@ -107,7 +152,7 @@ std::unique_ptr<cpp2::PlanNodeDescription> Aggregate::explain() const {
     auto desc = SingleInputNode::explain();
     addDescription("groupKeys", folly::toJson(util::toJson(groupKeys_)), desc.get());
     folly::dynamic itemArr = folly::dynamic::array();
-    for (const auto &item : groupItems_) {
+    for (const auto& item : groupItems_) {
         folly::dynamic itemObj = folly::dynamic::object();
         itemObj.insert("distinct", util::toJson(item.distinct));
         itemObj.insert("funcType", static_cast<uint8_t>(item.func));

src/planner/Query.h

@@ -144,6 +144,8 @@ public:
 
     std::unique_ptr<cpp2::PlanNodeDescription> explain() const override;
 
+    GetNeighbors* clone(QueryContext* qctx) const;
+
     Expression* src() const {
         return src_;
     }
@@ -432,6 +434,8 @@ public:
 
     std::unique_ptr<cpp2::PlanNodeDescription> explain() const override;
 
+    IndexScan* clone(QueryContext* qctx) const;
+
     const std::vector<storage::cpp2::IndexQueryContext>* queryContext() const {
         return contexts_.get();
     }