extract attribute from properties function (#4604)

nevermore3 · web-flow · commit fc5735a133a2 · 2022-09-06T16:22:18.000+08:00
* extract attribute from properties function

* fix error

* fix subscript error

* add test case

* process scanEdges

* fix test error
diff --git a/src/graph/executor/query/ScanEdgesExecutor.cpp b/src/graph/executor/query/ScanEdgesExecutor.cpp
@@ -40,10 +40,10 @@ folly::Future<Status> ScanEdgesExecutor::scanEdges() {
         SCOPED_TIMER(&execTime_);
         otherStats_.emplace("total_rpc", folly::sformat("{}(us)", scanEdgesTime.elapsedInUSec()));
       })
-      .thenValue([this, se](StorageRpcResponse<ScanResponse> &&rpcResp) {
+      .thenValue([this](StorageRpcResponse<ScanResponse> &&rpcResp) {
         SCOPED_TIMER(&execTime_);
         addStats(rpcResp, otherStats_);
-        return handleResp(std::move(rpcResp), se->colNames());
+        return handleResp(std::move(rpcResp), {});
       });
 }
 
diff --git a/src/graph/optimizer/rule/GetEdgesTransformUtils.h b/src/graph/optimizer/rule/GetEdgesTransformUtils.h
@@ -48,6 +48,9 @@ class GetEdgesTransformUtils final {
         limit_count,
         {},
         traverse->filter() == nullptr ? nullptr : traverse->filter()->clone());
+    auto &colNames = traverse->colNames();
+    DCHECK_EQ(colNames.size(), 2);
+    scanEdges->setColNames({colNames.back()});
     return scanEdges;
   }
 
diff --git a/src/graph/planner/plan/PlanNodeVisitor.h b/src/graph/planner/plan/PlanNodeVisitor.h
@@ -22,6 +22,7 @@ class PlanNodeVisitor {
   virtual void visit(Project *node) = 0;
   virtual void visit(Aggregate *node) = 0;
   virtual void visit(Traverse *node) = 0;
+  virtual void visit(ScanEdges *node) = 0;
   virtual void visit(AppendVertices *node) = 0;
   virtual void visit(BiJoin *node) = 0;
   virtual void visit(Union *node) = 0;
diff --git a/src/graph/planner/plan/Query.cpp b/src/graph/planner/plan/Query.cpp
@@ -260,6 +260,10 @@ PlanNode* ScanEdges::clone() const {
   return newGE;
 }
 
+void ScanEdges::accept(PlanNodeVisitor* visitor) {
+  visitor->visit(this);
+}
+
 void ScanEdges::cloneMembers(const ScanEdges& ge) {
   Explore::cloneMembers(ge);
 
diff --git a/src/graph/planner/plan/Query.h b/src/graph/planner/plan/Query.h
@@ -719,6 +719,8 @@ class ScanEdges final : public Explore {
   PlanNode* clone() const override;
   std::unique_ptr<PlanNodeDescription> explain() const override;
 
+  void accept(PlanNodeVisitor* visitor) override;
+
  private:
   friend ObjectPool;
   ScanEdges(QueryContext* qctx,
diff --git a/src/graph/visitor/PropertyTrackerVisitor.cpp b/src/graph/visitor/PropertyTrackerVisitor.cpp
@@ -5,9 +5,6 @@
 
 #include "graph/visitor/PropertyTrackerVisitor.h"
 
-#include <sstream>
-#include <unordered_set>
-
 #include "common/expression/Expression.h"
 #include "graph/context/QueryContext.h"
 
@@ -174,27 +171,78 @@ void PropertyTrackerVisitor::visit(AttributeExpression *expr) {
     return;
   }
   auto &propName = constVal.getStr();
-  static const int kUnknownEdgeType = 0;
   switch (lhs->kind()) {
     case Expression::Kind::kInputProperty:
     case Expression::Kind::kVarProperty: {  // $e.name
       auto *varPropExpr = static_cast<PropertyExpression *>(lhs);
       auto &edgeAlias = varPropExpr->prop();
-      propsUsed_.insertEdgeProp(edgeAlias, kUnknownEdgeType, propName);
+      propsUsed_.insertEdgeProp(edgeAlias, unKnowType_, propName);
       break;
     }
     case Expression::Kind::kSubscript: {  // $-.e[0].name
       auto *subscriptExpr = static_cast<SubscriptExpression *>(lhs);
       auto *subLeftExpr = subscriptExpr->left();
-      if (subLeftExpr->kind() == Expression::Kind::kInputProperty) {
-        auto *inputPropExpr = static_cast<InputPropertyExpression *>(subLeftExpr);
-        auto &edgeAlias = inputPropExpr->prop();
-        propsUsed_.insertEdgeProp(edgeAlias, kUnknownEdgeType, propName);
+      auto kind = subLeftExpr->kind();
+      if (kind == Expression::Kind::kInputProperty || kind == Expression::Kind::kVarProperty) {
+        auto *propExpr = static_cast<PropertyExpression *>(subLeftExpr);
+        auto &edgeAlias = propExpr->prop();
+        propsUsed_.insertEdgeProp(edgeAlias, unKnowType_, propName);
+      } else if (kind == Expression::Kind::kListComprehension) {
+        //  match (src_v:player{name:"Manu Ginobili"})-[e*2]-(dst_v) return e[0].start_year
+        auto *listExpr = static_cast<ListComprehensionExpression *>(subLeftExpr);
+        auto *collectExpr = listExpr->collection();
+        if (collectExpr->kind() == Expression::Kind::kInputProperty) {
+          auto *inputPropExpr = static_cast<InputPropertyExpression *>(collectExpr);
+          auto &aliasName = inputPropExpr->prop();
+          propsUsed_.insertEdgeProp(aliasName, unKnowType_, propName);
+        }
       }
       break;
     }
     case Expression::Kind::kFunctionCall: {  // properties(t3).name
-      // TODO(jmq) determine whether it is a vertex or edge
+      auto *funCallExpr = static_cast<FunctionCallExpression *>(lhs);
+      auto funName = funCallExpr->name();
+      std::transform(funName.begin(), funName.end(), funName.begin(), ::tolower);
+      if (funName != "properties") {
+        break;
+      }
+      auto *argExpr = funCallExpr->args()->args().front();
+      auto kind = argExpr->kind();
+      switch (kind) {
+        case Expression::Kind::kVarProperty:
+        case Expression::Kind::kInputProperty: {
+          //  match (v) return properties(v).name
+          auto *inputPropExpr = static_cast<InputPropertyExpression *>(argExpr);
+          auto &aliasName = inputPropExpr->prop();
+          propsUsed_.insertVertexProp(aliasName, unKnowType_, propName);
+          break;
+        }
+        case Expression::Kind::kSubscript: {
+          auto *subscriptExpr = static_cast<SubscriptExpression *>(argExpr);
+          auto *subLeftExpr = subscriptExpr->left();
+          auto leftKind = subLeftExpr->kind();
+          if (leftKind == Expression::Kind::kInputProperty ||
+              leftKind == Expression::Kind::kVarProperty) {
+            //  match (v)-[e]->(b) return properties(e).start_year
+            auto *propExpr = static_cast<PropertyExpression *>(subLeftExpr);
+            auto &aliasName = propExpr->prop();
+            propsUsed_.insertEdgeProp(aliasName, unKnowType_, propName);
+          } else if (leftKind == Expression::Kind::kListComprehension) {
+            //  match (v)-[c*2]->(b) retrun properties(c[0]).start_year
+            //  properties([e IN $-.c WHERE is_edge($e)][0]).start_year
+            auto *listExpr = static_cast<ListComprehensionExpression *>(subLeftExpr);
+            auto *collectExpr = listExpr->collection();
+            if (collectExpr->kind() == Expression::Kind::kInputProperty) {
+              auto *inputPropExpr = static_cast<InputPropertyExpression *>(collectExpr);
+              auto &aliasName = inputPropExpr->prop();
+              propsUsed_.insertEdgeProp(aliasName, unKnowType_, propName);
+            }
+          }
+          break;
+        }
+        default:
+          break;
+      }
       break;
     }
     default:
@@ -294,15 +342,5 @@ void PropertyTrackerVisitor::visit(EdgeExpression *expr) {
   UNUSED(expr);
 }
 
-std::string PropertyTrackerVisitor::extractColNameFromInputPropOrVarPropExpr(
-    const Expression *expr) {
-  if (expr->kind() == Expression::Kind::kInputProperty ||
-      expr->kind() == Expression::Kind::kVarProperty) {
-    auto *propExpr = static_cast<const PropertyExpression *>(expr);
-    return propExpr->prop();
-  }
-  return "";
-}
-
 }  // namespace graph
 }  // namespace nebula
diff --git a/src/graph/visitor/PropertyTrackerVisitor.h b/src/graph/visitor/PropertyTrackerVisitor.h
@@ -74,9 +74,8 @@ class PropertyTrackerVisitor : public ExprVisitorImpl {
   void visit(ColumnExpression* expr) override;
   void visit(AggregateExpression* expr) override;
 
-  std::string extractColNameFromInputPropOrVarPropExpr(const Expression* expr);
-
   const QueryContext* qctx_{nullptr};
+  const int unKnowType_{0};
   GraphSpaceID space_;
   PropertyTracker& propsUsed_;
   std::string entityAlias_;
diff --git a/src/graph/visitor/PrunePropertiesVisitor.cpp b/src/graph/visitor/PrunePropertiesVisitor.cpp
@@ -129,6 +129,61 @@ void PrunePropertiesVisitor::visitCurrent(Aggregate *node) {
   }
 }
 
+void PrunePropertiesVisitor::visit(ScanEdges *node) {
+  rootNode_ = false;
+  pruneCurrent(node);
+  status_ = depsPruneProperties(node->dependencies());
+}
+
+void PrunePropertiesVisitor::pruneCurrent(ScanEdges *node) {
+  auto &colNames = node->colNames();
+  DCHECK(!colNames.empty());
+  auto &edgeAlias = colNames.back();
+  auto it = propsUsed_.colsSet.find(edgeAlias);
+  if (it != propsUsed_.colsSet.end()) {
+    // all properties are used
+    return;
+  }
+  auto *edgeProps = node->props();
+  auto &edgePropsMap = propsUsed_.edgePropsMap;
+
+  auto prunedEdgeProps = std::make_unique<std::vector<EdgeProp>>();
+  prunedEdgeProps->reserve(edgeProps->size());
+  auto edgeAliasIter = edgePropsMap.find(edgeAlias);
+
+  for (auto &edgeProp : *edgeProps) {
+    auto edgeType = edgeProp.type_ref().value();
+    auto &props = edgeProp.props_ref().value();
+    EdgeProp newEdgeProp;
+    newEdgeProp.type_ref() = edgeType;
+    if (edgeAliasIter == edgePropsMap.end()) {
+      // only type, dst are used
+      newEdgeProp.props_ref() = {nebula::kSrc, nebula::kDst, nebula::kType, nebula::kRank};
+    } else {
+      std::unordered_set<std::string> uniqueProps{
+          nebula::kSrc, nebula::kDst, nebula::kType, nebula::kRank};
+      std::vector<std::string> newProps;
+      auto &usedEdgeProps = edgeAliasIter->second;
+      auto edgeTypeIter = usedEdgeProps.find(std::abs(edgeType));
+      if (edgeTypeIter != usedEdgeProps.end()) {
+        uniqueProps.insert(edgeTypeIter->second.begin(), edgeTypeIter->second.end());
+      }
+      auto unKnowEdgeIter = usedEdgeProps.find(unKnowType_);
+      if (unKnowEdgeIter != usedEdgeProps.end()) {
+        uniqueProps.insert(unKnowEdgeIter->second.begin(), unKnowEdgeIter->second.end());
+      }
+      for (auto &prop : props) {
+        if (uniqueProps.find(prop) != uniqueProps.end()) {
+          newProps.emplace_back(prop);
+        }
+      }
+      newEdgeProp.props_ref() = std::move(newProps);
+    }
+    prunedEdgeProps->emplace_back(std::move(newEdgeProp));
+  }
+  node->setEdgeProps(std::move(prunedEdgeProps));
+}
+
 void PrunePropertiesVisitor::visit(Traverse *node) {
   rootNode_ = false;
   visitCurrent(node);
@@ -176,24 +231,34 @@ void PrunePropertiesVisitor::pruneCurrent(Traverse *node) {
       if (usedVertexProps.empty()) {
         node->setVertexProps(nullptr);
       } else {
+        auto unknowIter = usedVertexProps.find(unKnowType_);
         auto prunedVertexProps = std::make_unique<std::vector<VertexProp>>();
         prunedVertexProps->reserve(usedVertexProps.size());
         for (auto &vertexProp : *vertexProps) {
           auto tagId = vertexProp.tag_ref().value();
           auto &props = vertexProp.props_ref().value();
+          std::unordered_set<std::string> usedProps;
+          if (unknowIter != usedVertexProps.end()) {
+            usedProps.insert(unknowIter->second.begin(), unknowIter->second.end());
+          }
           auto tagIter = usedVertexProps.find(tagId);
-          if (tagIter == usedVertexProps.end()) {
+          if (tagIter != usedVertexProps.end()) {
+            usedProps.insert(tagIter->second.begin(), tagIter->second.end());
+          }
+          if (usedProps.empty()) {
             continue;
           }
-          auto &usedProps = tagIter->second;
-          VertexProp newVProp;
-          newVProp.tag_ref() = tagId;
           std::vector<std::string> newProps;
           for (auto &prop : props) {
             if (usedProps.find(prop) != usedProps.end()) {
               newProps.emplace_back(prop);
             }
           }
+          if (newProps.empty()) {
+            continue;
+          }
+          VertexProp newVProp;
+          newVProp.tag_ref() = tagId;
           newVProp.props_ref() = std::move(newProps);
           prunedVertexProps->emplace_back(std::move(newVProp));
         }
@@ -227,8 +292,7 @@ void PrunePropertiesVisitor::pruneCurrent(Traverse *node) {
       if (edgeTypeIter != usedEdgeProps.end()) {
         uniqueProps.insert(edgeTypeIter->second.begin(), edgeTypeIter->second.end());
       }
-      int kUnknowEdgeType = 0;
-      auto unKnowEdgeIter = usedEdgeProps.find(kUnknowEdgeType);
+      auto unKnowEdgeIter = usedEdgeProps.find(unKnowType_);
       if (unKnowEdgeIter != usedEdgeProps.end()) {
         uniqueProps.insert(unKnowEdgeIter->second.begin(), unKnowEdgeIter->second.end());
       }
@@ -318,24 +382,33 @@ void PrunePropertiesVisitor::pruneCurrent(AppendVertices *node) {
     }
     return;
   }
-
+  auto unknowIter = usedVertexProps.find(unKnowType_);
   prunedVertexProps->reserve(usedVertexProps.size());
   for (auto &vertexProp : *vertexProps) {
     auto tagId = vertexProp.tag_ref().value();
     auto &props = vertexProp.props_ref().value();
     auto tagIter = usedVertexProps.find(tagId);
-    if (tagIter == usedVertexProps.end()) {
+    std::unordered_set<std::string> usedProps;
+    if (unknowIter != usedVertexProps.end()) {
+      usedProps.insert(unknowIter->second.begin(), unknowIter->second.end());
+    }
+    if (tagIter != usedVertexProps.end()) {
+      usedProps.insert(tagIter->second.begin(), tagIter->second.end());
+    }
+    if (usedProps.empty()) {
       continue;
     }
-    auto &usedProps = tagIter->second;
-    VertexProp newVProp;
-    newVProp.tag_ref() = tagId;
     std::vector<std::string> newProps;
     for (auto &prop : props) {
       if (usedProps.find(prop) != usedProps.end()) {
         newProps.emplace_back(prop);
       }
     }
+    if (newProps.empty()) {
+      continue;
+    }
+    VertexProp newVProp;
+    newVProp.tag_ref() = tagId;
     newVProp.props_ref() = std::move(newProps);
     prunedVertexProps->emplace_back(std::move(newVProp));
   }
diff --git a/src/graph/visitor/PrunePropertiesVisitor.h b/src/graph/visitor/PrunePropertiesVisitor.h
@@ -55,6 +55,10 @@ class PrunePropertiesVisitor final : public PlanNodeVisitor {
   // prune properties in Traverse according to the used properties collected previous
   void pruneCurrent(Traverse *node);
 
+  void visit(ScanEdges *node) override;
+
+  void pruneCurrent(ScanEdges *node);
+
   void visit(AppendVertices *node) override;
   // \param node, the current node to visit
   // \param used, whether properties in current node are used
@@ -80,6 +84,7 @@ class PrunePropertiesVisitor final : public PlanNodeVisitor {
   GraphSpaceID spaceID_;
   Status status_;
   bool rootNode_{true};
+  const int unKnowType_ = 0;
 };
 
 }  // namespace graph
diff --git a/tests/tck/features/optimizer/PrunePropertiesRule.feature b/tests/tck/features/optimizer/PrunePropertiesRule.feature

Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,9 @@ class GetEdgesTransformUtils final {`
`48`	`48`	`limit_count,`
`49`	`49`	`{},`
`50`	`50`	`traverse->filter() == nullptr ? nullptr : traverse->filter()->clone());`
	`51`	`+ auto &colNames = traverse->colNames();`
	`52`	`+ DCHECK_EQ(colNames.size(), 2);`
	`53`	`+ scanEdges->setColNames({colNames.back()});`
`51`	`54`	`return scanEdges;`
`52`	`55`	`}`
`53`	`56`
Original file line number	Diff line number	Diff line change
`@@ -260,6 +260,10 @@ PlanNode* ScanEdges::clone() const {`
`260`	`260`	`return newGE;`
`261`	`261`	`}`
`262`	`262`
	`263`	`+void ScanEdges::accept(PlanNodeVisitor* visitor) {`
	`264`	`+ visitor->visit(this);`
	`265`	`+}`
	`266`	`+`
`263`	`267`	`void ScanEdges::cloneMembers(const ScanEdges& ge) {`
`264`	`268`	`Explore::cloneMembers(ge);`
`265`	`269`