Improve the traverse executor performance

src/clients/storage/StorageClientBase.h

@@ -97,6 +97,11 @@ class StorageRpcResponse final {
  return responses_;
  }
 
+ // Not thread-safe.
+ const std::vector<Response>& responses() const {
+ return responses_;
+ }
+
  // Not thread-safe.
  const std::vector<std::tuple<HostAddr, int32_t, int32_t>>& hostLatency() const {
  return hostLatency_;

src/common/base/Base.h

@@ -133,6 +133,9 @@ constexpr char kSrc[] = "_src";
 constexpr char kType[] = "_type";
 constexpr char kRank[] = "_rank";
 constexpr char kDst[] = "_dst";
+constexpr char kEdgePrefix[] = "_edge";
+constexpr char kStatsPrefix[] = "_stats";
+constexpr char kExprPrefix[] = "_expr";
 
 // Useful type traits
 

src/common/datatypes/List.h

@@ -80,6 +80,10 @@ struct List {
  return values[i];
  }
 
+ Value& operator[](size_t i) {
+ return values[i];
+ }
+
  bool contains(const Value& value) const {
  return std::find(values.begin(), values.end(), value) != values.end();
  }

src/common/datatypes/Value.cpp

@@ -3000,4 +3000,43 @@ Value operator^(const Value& lhs, const Value& rhs) {
  }
  }
 }
+
+std::size_t VertexHash::operator()(const Value& v) const {
+ switch (v.type()) {
+ case Value::Type::VERTEX: {
+ auto& vid = v.getVertex().vid;
+ if (vid.type() == Value::Type::STRING) {
+ return std::hash<std::string>()(vid.getStr());
+ } else {
+ return vid.getInt();
+ }
+ }
+ case Value::Type::STRING: {
+ return std::hash<std::string>()(v.getStr());
+ }
+ case Value::Type::INT: {
+ return v.getInt();
+ }
+ default: {
+ return v.hash();
+ }
+ }
+}
+
+bool VertexEqual::operator()(const Value& lhs, const Value& rhs) const {
+ if (lhs.type() == rhs.type()) {
+ if (lhs.isVertex()) {
+ return lhs.getVertex().vid == rhs.getVertex().vid;
+ }
+ return lhs == rhs;
+ }
+ if (lhs.type() == Value::Type::VERTEX) {
+ return lhs.getVertex().vid == rhs;
+ }
+ if (rhs.type() == Value::Type::VERTEX) {
+ return lhs == rhs.getVertex().vid;
+ }
+ return lhs == rhs;
+}
+
 } // namespace nebula

src/common/datatypes/Value.h

@@ -550,6 +550,16 @@ inline uint64_t operator&(const Value::Type& lhs, const uint64_t rhs) {
  return static_cast<uint64_t>(lhs) & rhs;
 }
 
+struct VertexHash {
+ std::size_t operator()(const Value& v) const;
+};
+
+struct VertexEqual {
+ bool operator()(const Value& lhs, const Value& rhs) const;
+};
+
+using VidHashSet = std::unordered_set<Value, VertexHash, VertexEqual>;
+
 } // namespace nebula
 
 namespace std {

src/common/datatypes/Vertex.h

@@ -24,6 +24,7 @@ struct Tag {
  Tag(const Tag& tag) : name(tag.name), props(tag.props) {}
  Tag(std::string tagName, std::unordered_map<std::string, Value> tagProps)
  : name(std::move(tagName)), props(std::move(tagProps)) {}
+ explicit Tag(const std::string& tagName) : name(tagName), props() {}
 
  void clear() {
  name.clear();

src/graph/context/CMakeLists.txt

@@ -10,6 +10,7 @@ nebula_add_library(
  Iterator.cpp
  Result.cpp
  Symbols.cpp
+ iterator/GetNbrsRespDataSetIter.cpp
 )
 
 

src/graph/context/iterator/GetNbrsRespDataSetIter.cpp

@@ -0,0 +1,154 @@
+/* Copyright (c) 2023 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License.
+ */
+
+#include "graph/context/iterator/GetNbrsRespDataSetIter.h"
+
+#include "common/base/Base.h"
+#include "common/datatypes/Edge.h"
+#include "common/datatypes/Vertex.h"
+
+namespace nebula {
+namespace graph {
+
+bool isDataSetInvalid(const DataSet* dataset) {
+ const auto& colNames = dataset->colNames;
+ return colNames.size() < 3 || // the dataset has vid/_tag/_edge 3 columns at least
+ colNames[0] != nebula::kVid || // the first column should be Vid
+ colNames[1].find(kStatsPrefix) != 0 || // the second column could not be _stats column now
+ colNames.back().find(kExprPrefix) != 0; // the last column could not be _expr column now
+}
+
+GetNbrsRespDataSetIter::GetNbrsRespDataSetIter(const DataSet* dataset)
+ : dataset_(DCHECK_NOTNULL(dataset)), curRowIdx_(0) {
+ DCHECK(!isDataSetInvalid(dataset));
+ for (size_t i = 0, e = dataset->colNames.size(); i < e; ++i) {
+ buildPropIndex(dataset->colNames[i], i);
+ }
+}
+
+void GetNbrsRespDataSetIter::buildPropIndex(const std::string& colName, size_t colIdx) {
+ std::vector<std::string> pieces;
+ folly::split(":", colName, pieces);
+ if (pieces.size() < 2) {
+ // Skip the column which is not tag or edge
+ return;
+ }
+
+ PropIndex propIdx;
+ propIdx.colIdx = colIdx;
+
+ propIdx.propIdxMap.reserve(pieces.size() - 2);
+ // if size == 2, it is the tag defined without props.
+ for (size_t i = 2; i < pieces.size(); ++i) {
+ const auto& name = pieces[i];
+ size_t idx = i - 2;
+ if (name == kType) {
+ propIdx.edgeTypeIdx = idx;
+ } else if (name == kRank) {
+ propIdx.edgeRankIdx = idx;
+ } else if (name == kDst) {
+ propIdx.edgeDstIdx = idx;
+ } else if (name == kSrc) {
+ // Always skip the src of edge since it is returned by the first column
+ } else if (name == kTag) {
+ // Skip the _tag prop of vertex since it is not used to create vertex
+ } else {
+ propIdx.propIdxMap.emplace(name, idx);
+ }
+ }
+
+ folly::StringPiece prefix(pieces[0]), name(pieces[1]);
+ DCHECK(!name.empty()) << "The name of tag/edge is empty";
+ if (prefix.startsWith(kEdgePrefix)) {
+ DCHECK(name.startsWith("-") || name.startsWith("+")) << "the edge name has to start with '-/+'";
+ edgePropsMap_.emplace(name, std::move(propIdx));
+ } else if (prefix.startsWith(kTag)) {
+ tagPropsMap_.emplace(name, std::move(propIdx));
+ }
+}
+
+Value GetNbrsRespDataSetIter::getVertex() const {
+ // Always check the valid() before getVertex
+ DCHECK(valid());
+ const Row& curRow = dataset_->rows[curRowIdx_];
+ Vertex vertex;
+ vertex.vid = curRow[0];
+ vertex.tags.reserve(tagPropsMap_.size());
+ for (const auto& [tagName, propIdx] : tagPropsMap_) {
+ DCHECK_LT(propIdx.colIdx, curRow.size());
+ const Value& propColumn = curRow[propIdx.colIdx];
+ if (propColumn.isList()) {
+ const List& propList = propColumn.getList();
+
+ Tag tag(tagName);
+ tag.props.reserve(propIdx.propIdxMap.size());
+ for (const auto& [propName, pIdx] : propIdx.propIdxMap) {
+ DCHECK_LT(pIdx, propList.size());
+ tag.props.emplace(propName, propList[pIdx]);
+ }
+
+ vertex.tags.emplace_back(std::move(tag));
+ }
+ }
+ return vertex;
+}
+
+Value GetNbrsRespDataSetIter::createEdgeByPropList(const PropIndex& propIdx,
+ const Value& edgeVal,
+ const Value& src,
+ const std::string& edgeName) const {
+ if (!edgeVal.isList() || edgeVal.getList().empty()) {
+ return Value::kEmpty;
+ }
+ const List& propList = edgeVal.getList();
+ DCHECK_LT(propIdx.edgeDstIdx, propList.size());
+ DCHECK_LT(propIdx.edgeTypeIdx, propList.size());
+ DCHECK_LT(propIdx.edgeRankIdx, propList.size());
+
+ Edge edge;
+ edge.name = edgeName;
+ edge.src = src;
+ edge.dst = propList[propIdx.edgeDstIdx];
+ const Value& typeVal = propList[propIdx.edgeTypeIdx];
+ edge.type = typeVal.isInt() ? typeVal.getInt() : 0;
+ const Value& rankVal = propList[propIdx.edgeRankIdx];
+ edge.ranking = rankVal.isInt() ? rankVal.getInt() : 0;
+
+ edge.props.reserve(propIdx.propIdxMap.size());
+ for (const auto& [propName, pIdx] : propIdx.propIdxMap) {
+ DCHECK_LT(pIdx, propList.size());
+ edge.props.emplace(propName, propList[pIdx]);
+ }
+
+ return edge;
+}
+
+std::vector<Value> GetNbrsRespDataSetIter::getAdjEdges(VidHashSet* dstSet) const {
+ DCHECK(valid());
+
+ std::vector<Value> adjEdges;
+ const Row& curRow = dataset_->rows[curRowIdx_];
+ for (const auto& [edgeName, propIdx] : edgePropsMap_) {
+ DCHECK_LT(propIdx.colIdx, curRow.size());
+ const Value& edgeColumn = curRow[propIdx.colIdx];
+ if (edgeColumn.isList()) {
+ for (const Value& edgeVal : edgeColumn.getList().values) {
+ // skip the edge direction symbol: `-/+`
+ auto name = edgeName.substr(1);
+ Value edge = createEdgeByPropList(propIdx, edgeVal, curRow[0], name);
+ if (!edge.empty()) {
+ if (dstSet) {
+ dstSet->emplace(edge.getEdge().dst);
+ }
+ adjEdges.emplace_back(std::move(edge));
+ }
+ }
+ }
+ }
+ return adjEdges;
+}
+
+} // namespace graph
+} // namespace nebula

src/graph/context/iterator/GetNbrsRespDataSetIter.h

@@ -0,0 +1,58 @@
+/* Copyright (c) 2023 vesoft inc. All rights reserved.
+ *
+ * This source code is licensed under Apache 2.0 License.
+ */
+
+#ifndef GRAPH_CONTEXT_ITERATOR_GETNBRSRESPDATASETITER_H_
+#define GRAPH_CONTEXT_ITERATOR_GETNBRSRESPDATASETITER_H_
+
+#include "common/datatypes/DataSet.h"
+#include "common/datatypes/Value.h"
+
+namespace nebula {
+namespace graph {
+
+class GetNbrsRespDataSetIter final {
+ public:
+ explicit GetNbrsRespDataSetIter(const DataSet* dataset);
+
+ bool valid() const {
+ return curRowIdx_ < dataset_->rowSize();
+ }
+ // Next row in dataset
+ void next() {
+ curRowIdx_++;
+ }
+
+ Value getVertex() const;
+ std::vector<Value> getAdjEdges(VidHashSet* dstSet) const;
+
+ private:
+ struct PropIndex {
+ size_t colIdx;
+ size_t edgeTypeIdx;
+ size_t edgeRankIdx;
+ size_t edgeDstIdx;
+ std::unordered_map<std::string, size_t> propIdxMap;
+ };
+
+ void buildPropIndex(const std::string& colName, size_t colIdx);
+ Value createEdgeByPropList(const PropIndex& propIdx,
+ const Value& edgeVal,
+ const Value& src,
+ const std::string& edgeName) const;
+
+ // my fields
+ const DataSet* dataset_;
+ size_t curRowIdx_;
+
+ // _tag:t1:p1:p2 -> {t1 : [column_idx, [p1, p2], {p1 : 0, p2 : 1}]}
+ std::unordered_map<std::string, PropIndex> tagPropsMap_;
+ // _edge:e1:p1:p2 -> {e1 : [column_idx, [p1, p2], {p1 : 0, p2 : 1}]}
+ std::unordered_map<std::string, PropIndex> edgePropsMap_;
+};
+
+} // namespace graph
+} // namespace nebula
+
+#endif // GRAPH_CONTEXT_ITERATOR_GETNBRSRESPDATASETITER_H_