refactor(interactive): Support running pagerank on edge triplet subgraph (#4537)

Preivously, builtin pagerank suppose that the src and dst vertex of the
edge label is of the same label. In this PR, we remove this constraint.

Author: zhanglei1949

flex/engines/graph_db/app/builtin/pagerank.cc

@@ -13,21 +13,55 @@
  * limitations under the License.
  */
 #include "flex/engines/graph_db/app/builtin/pagerank.h"
+#include "flex/engines/graph_db/runtime/common/graph_interface.h"
+#include "flex/engines/graph_db/runtime/common/rt_any.h"
 
 namespace gs {
 
-results::CollectiveResults PageRank::Query(const GraphDBSession& sess,
-                                           std::string vertex_label,
-                                           std::string edge_label,
-                                           double damping_factor,
-                                           int max_iterations, double epsilon) {
+void write_result(
+    const ReadTransaction& txn, results::CollectiveResults& results,
+    const std::vector<std::tuple<label_t, vid_t, double>>& pagerank,
+    int32_t result_limit) {
+  runtime::GraphReadInterface graph(txn);
+
+  for (int32_t j = 0; j < std::min((int32_t) pagerank.size(), result_limit);
+       ++j) {
+    auto vertex_label = std::get<0>(pagerank[j]);
+    auto vertex_label_name = txn.schema().get_vertex_label_name(vertex_label);
+    auto vid = std::get<1>(pagerank[j]);
+    runtime::RTAny any(txn.GetVertexId(vertex_label, vid));
+    auto result = results.add_results();
+    auto first_col = result->mutable_record()->add_columns();
+    first_col->mutable_name_or_id()->set_id(0);
+    first_col->mutable_entry()->mutable_element()->mutable_object()->set_str(
+        vertex_label_name);
+
+    auto oid_col = result->mutable_record()->add_columns();
+    any.sink(graph, 1, oid_col);
+
+    auto pagerank_col = result->mutable_record()->add_columns();
+    pagerank_col->mutable_name_or_id()->set_id(2);
+    pagerank_col->mutable_entry()->mutable_element()->mutable_object()->set_f64(
+        std::get<2>(pagerank[j]));
+  }
+}
+
+results::CollectiveResults PageRank::Query(
+    const GraphDBSession& sess, std::string src_vertex_label,
+    std::string dst_vertex_label, std::string edge_label, double damping_factor,
+    int32_t max_iterations, double epsilon, int32_t result_limit) {
   auto txn = sess.GetReadTransaction();
 
-  if (!sess.schema().has_vertex_label(vertex_label)) {
-    LOG(ERROR) << "The requested vertex label doesn't exits.";
+  if (!sess.schema().has_vertex_label(src_vertex_label)) {
+    LOG(ERROR) << "The requested src vertex label doesn't exits.";
+    return {};
+  }
+  if (!sess.schema().has_vertex_label(dst_vertex_label)) {
+    LOG(ERROR) << "The requested dst vertex label doesn't exits.";
     return {};
   }
-  if (!sess.schema().has_edge_label(vertex_label, vertex_label, edge_label)) {
+  if (!sess.schema().has_edge_label(src_vertex_label, dst_vertex_label,
+                                    edge_label)) {
     LOG(ERROR) << "The requested edge label doesn't exits.";
     return {};
   }
@@ -44,59 +78,117 @@ results::CollectiveResults PageRank::Query(const GraphDBSession& sess,
     return {};
   }
 
-  auto vertex_label_id = sess.schema().get_vertex_label_id(vertex_label);
+  auto src_vertex_label_id =
+      sess.schema().get_vertex_label_id(src_vertex_label);
+  auto dst_vertex_label_id =
+      sess.schema().get_vertex_label_id(dst_vertex_label);
   auto edge_label_id = sess.schema().get_edge_label_id(edge_label);
 
-  auto num_vertices = txn.GetVertexNum(vertex_label_id);
-
-  std::unordered_map<vid_t, double> pagerank;
-  std::unordered_map<vid_t, double> new_pagerank;
+  auto num_src_vertices = txn.GetVertexNum(src_vertex_label_id);
+  auto num_dst_vertices = txn.GetVertexNum(dst_vertex_label_id);
+  auto num_vertices = src_vertex_label_id == dst_vertex_label_id
+                          ? num_src_vertices
+                          : num_src_vertices + num_dst_vertices;
+
+  std::vector<std::vector<double>> pagerank;
+  std::vector<std::vector<double>> new_pagerank;
+  std::vector<std::vector<int32_t>> outdegree;
+
+  bool dst_to_src = src_vertex_label_id != dst_vertex_label_id &&
+                    txn.schema().exist(dst_vertex_label_id, src_vertex_label_id,
+                                       edge_label_id);
+
+  pagerank.emplace_back(std::vector<double>(num_src_vertices, 0.0));
+  new_pagerank.emplace_back(std::vector<double>(num_src_vertices, 0.0));
+  outdegree.emplace_back(std::vector<int32_t>(num_src_vertices, 0));
+  if (dst_to_src) {
+    pagerank.emplace_back(std::vector<double>(num_dst_vertices, 0.0));
+    new_pagerank.emplace_back(std::vector<double>(num_dst_vertices, 0.0));
+    outdegree.emplace_back(std::vector<int32_t>(num_dst_vertices, 0));
+  }
 
-  auto vertex_iter = txn.GetVertexIterator(vertex_label_id);
+  auto src_vertex_iter = txn.GetVertexIterator(src_vertex_label_id);
 
-  while (vertex_iter.IsValid()) {
-    vid_t vid = vertex_iter.GetIndex();
-    pagerank[vid] = 1.0 / num_vertices;
-    new_pagerank[vid] = 0.0;
-    vertex_iter.Next();
+  while (src_vertex_iter.IsValid()) {
+    vid_t vid = src_vertex_iter.GetIndex();
+    pagerank[0][vid] = 1.0 / num_vertices;
+    new_pagerank[0][vid] = 0.0;
+    src_vertex_iter.Next();
+    outdegree[0][vid] = txn.GetOutDegree(src_vertex_label_id, vid,
+                                         dst_vertex_label_id, edge_label_id);
+  }
+  if (dst_to_src) {
+    auto dst_vertex_iter = txn.GetVertexIterator(dst_vertex_label_id);
+    while (dst_vertex_iter.IsValid()) {
+      vid_t vid = dst_vertex_iter.GetIndex();
+      pagerank[1][vid] = 1.0 / num_vertices;
+      new_pagerank[1][vid] = 0.0;
+      dst_vertex_iter.Next();
+      outdegree[1][vid] = txn.GetOutDegree(
+          dst_vertex_label_id, src_vertex_label_id, vid, edge_label_id);
+    }
   }
-
-  std::unordered_map<vid_t, double> outdegree;
 
   for (int iter = 0; iter < max_iterations; ++iter) {
-    for (auto& kv : new_pagerank) {
-      kv.second = 0.0;
+    new_pagerank[0].assign(num_src_vertices, 0.0);
+    if (dst_to_src) {
+      new_pagerank[1].assign(num_dst_vertices, 0.0);
     }
 
-    auto vertex_iter = txn.GetVertexIterator(vertex_label_id);
-    while (vertex_iter.IsValid()) {
-      vid_t v = vertex_iter.GetIndex();
+    auto src_vertex_iter = txn.GetVertexIterator(src_vertex_label_id);
+    while (src_vertex_iter.IsValid()) {
+      vid_t v = src_vertex_iter.GetIndex();
 
       double sum = 0.0;
-      auto edges = txn.GetInEdgeIterator(vertex_label_id, v, vertex_label_id,
-                                         edge_label_id);
-      while (edges.IsValid()) {
-        auto neighbor = edges.GetNeighbor();
-        if (outdegree[neighbor] == 0) {
-          auto out_edges = txn.GetOutEdgeIterator(
-              vertex_label_id, neighbor, vertex_label_id, edge_label_id);
-          while (out_edges.IsValid()) {
-            outdegree[neighbor]++;
-            out_edges.Next();
-          }
+      {
+        auto edges = txn.GetInEdgeIterator(dst_vertex_label_id, v,
+                                           src_vertex_label_id, edge_label_id);
+        while (edges.IsValid()) {
+          auto neighbor = edges.GetNeighbor();
+          sum += pagerank[0][neighbor] / outdegree[0][neighbor];
+          edges.Next();
         }
-        sum += pagerank[neighbor] / outdegree[neighbor];
-        edges.Next();
       }
 
-      new_pagerank[v] =
+      new_pagerank[0][v] =
           damping_factor * sum + (1.0 - damping_factor) / num_vertices;
-      vertex_iter.Next();
+      src_vertex_iter.Next();
+    }
+
+    if (dst_to_src) {
+      auto dst_vertex_iter = txn.GetVertexIterator(dst_vertex_label_id);
+      while (dst_vertex_iter.IsValid()) {
+        vid_t v = dst_vertex_iter.GetIndex();
+
+        double sum = 0.0;
+        {
+          auto edges = txn.GetInEdgeIterator(
+              src_vertex_label_id, v, dst_vertex_label_id, edge_label_id);
+          while (edges.IsValid()) {
+            LOG(INFO) << "got edge, from " << edges.GetNeighbor() << " to " << v
+                      << " label: " << std::to_string(src_vertex_label_id)
+                      << " " << std::to_string(dst_vertex_label_id) << " "
+                      << std::to_string(edge_label_id);
+            auto neighbor = edges.GetNeighbor();
+            sum += pagerank[1][neighbor] / outdegree[1][neighbor];
+            edges.Next();
+          }
+        }
+
+        new_pagerank[1][v] =
+            damping_factor * sum + (1.0 - damping_factor) / num_vertices;
+        dst_vertex_iter.Next();
+      }
     }
 
     double diff = 0.0;
-    for (const auto& kv : pagerank) {
-      diff += std::abs(new_pagerank[kv.first] - kv.second);
+    for (size_t j = 0; j < new_pagerank[0].size(); ++j) {
+      diff += std::abs(new_pagerank[0][j] - pagerank[0][j]);
+    }
+    if (dst_to_src) {
+      for (size_t j = 0; j < new_pagerank[1].size(); ++j) {
+        diff += std::abs(new_pagerank[1][j] - pagerank[1][j]);
+      }
     }
 
     if (diff < epsilon) {
@@ -108,28 +200,23 @@ results::CollectiveResults PageRank::Query(const GraphDBSession& sess,
 
   results::CollectiveResults results;
 
-  for (auto kv : pagerank) {
-    int64_t oid_ = txn.GetVertexId(vertex_label_id, kv.first).AsInt64();
-    auto result = results.add_results();
-    result->mutable_record()
-        ->add_columns()
-        ->mutable_entry()
-        ->mutable_element()
-        ->mutable_object()
-        ->set_str(vertex_label);
-    result->mutable_record()
-        ->add_columns()
-        ->mutable_entry()
-        ->mutable_element()
-        ->mutable_object()
-        ->set_i64(oid_);
-    result->mutable_record()
-        ->add_columns()
-        ->mutable_entry()
-        ->mutable_element()
-        ->mutable_object()
-        ->set_f64(kv.second);
+  std::vector<std::tuple<label_t, vid_t, double>> final_pagerank(num_vertices);
+  for (size_t i = 0; i < pagerank[0].size(); ++i) {
+    final_pagerank[i] = std::make_tuple(src_vertex_label_id, i, pagerank[0][i]);
   }
+  if (dst_to_src) {
+    for (size_t i = 0; i < pagerank[1].size(); ++i) {
+      final_pagerank[i + num_src_vertices] =
+          std::make_tuple(dst_vertex_label_id, i, pagerank[1][i]);
+    }
+  }
+  std::sort(final_pagerank.begin(), final_pagerank.end(),
+            [](const std::tuple<label_t, vid_t, double>& a,
+               const std::tuple<label_t, vid_t, double>& b) {
+              return std::get<2>(a) > std::get<2>(b);
+            });
+
+  write_result(txn, results, final_pagerank, result_limit);
 
   txn.Commit();
   return results;

flex/engines/graph_db/app/builtin/pagerank.h

@@ -15,19 +15,22 @@
 
 #ifndef ENGINES_GRAPH_DB_APP_BUILDIN_PAGERANK_H_
 #define ENGINES_GRAPH_DB_APP_BUILDIN_PAGERANK_H_
+
 #include "flex/engines/graph_db/database/graph_db_session.h"
 #include "flex/engines/hqps_db/app/interactive_app_base.h"
 
 namespace gs {
-class PageRank
-    : public CypherReadAppBase<std::string, std::string, double, int, double> {
+class PageRank : public CypherReadAppBase<std::string, std::string, std::string,
+                                          double, int32_t, double, int32_t> {
  public:
   PageRank() {}
   results::CollectiveResults Query(const GraphDBSession& sess,
-                                   std::string vertex_label,
+                                   std::string src_vertex_label,
+                                   std::string dst_vertex_label,
                                    std::string edge_label,
-                                   double damping_factor, int max_iterations,
-                                   double epsilon);
+                                   double damping_factor,
+                                   int32_t max_iterations, double epsilon,
+                                   int32_t result_limit) override;
 };
 
 class PageRankFactory : public AppFactoryBase {

flex/engines/graph_db/database/read_transaction.cc

@@ -137,6 +137,20 @@ ReadTransaction::edge_iterator ReadTransaction::GetInEdgeIterator(
           graph_.get_incoming_edges(label, u, neighbor_label, edge_label)};
 }
 
+size_t ReadTransaction::GetOutDegree(label_t label, vid_t u,
+                                     label_t neighbor_label,
+                                     label_t edge_label) const {
+  return graph_.get_outgoing_edges(label, u, neighbor_label, edge_label)
+      ->size();
+}
+
+size_t ReadTransaction::GetInDegree(label_t label, vid_t u,
+                                    label_t neighbor_label,
+                                    label_t edge_label) const {
+  return graph_.get_incoming_edges(label, u, neighbor_label, edge_label)
+      ->size();
+}
+
 void ReadTransaction::release() {
   if (timestamp_ != std::numeric_limits<timestamp_t>::max()) {
     vm_.release_read_timestamp();

flex/engines/graph_db/database/read_transaction.h

@@ -491,6 +491,12 @@ class ReadTransaction {
                                   label_t neighbor_label,
                                   label_t edge_label) const;
 
+  size_t GetOutDegree(label_t label, vid_t u, label_t neighbor_label,
+                      label_t edge_label) const;
+
+  size_t GetInDegree(label_t label, vid_t u, label_t neighbor_label,
+                     label_t edge_label) const;
+
   template <typename EDATA_T>
   AdjListView<EDATA_T> GetOutgoingEdges(label_t v_label, vid_t v,
                                         label_t neighbor_label,

flex/interactive/sdk/python/gs_interactive/tests/test_robustness.py

@@ -260,10 +260,12 @@ def test_builtin_procedure(interactive_session, neo4j_session, create_modern_gra
         create_modern_graph,
         "pagerank",
         '"person"',
+        '"person"',
         '"knows"',
         "0.85",
         "100",
         "0.000001",
+        "10",
     )
 
     call_procedure(

flex/storages/metadata/graph_meta_store.cc

@@ -80,11 +80,15 @@ const std::vector<PluginMeta>& get_builtin_plugin_metas() {
     pagerank.type = "cypher";
     pagerank.creation_time = GetCurrentTimeStamp();
     pagerank.update_time = GetCurrentTimeStamp();
-    pagerank.params.push_back({"vertex_label", PropertyType::kString, true});
+    pagerank.params.push_back(
+        {"src_vertex_label", PropertyType::kString, true});
+    pagerank.params.push_back(
+        {"dst_vertex_label", PropertyType::kString, true});
     pagerank.params.push_back({"edge_label", PropertyType::kString, true});
     pagerank.params.push_back({"damping_factor", PropertyType::kDouble, false});
     pagerank.params.push_back({"max_iterations", PropertyType::kInt32, false});
     pagerank.params.push_back({"epsilon", PropertyType::kDouble, false});
+    pagerank.params.push_back({"result_limit", PropertyType::kInt32, false});
     pagerank.returns.push_back({"label_name", PropertyType::kString});
     pagerank.returns.push_back({"vertex_oid", PropertyType::kInt64});
     pagerank.returns.push_back({"pagerank", PropertyType::kDouble});