change sample return type to pair

paddle/fluid/distributed/service/graph_brpc_client.cc

@@ -12,15 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/fluid/distributed/service/graph_brpc_client.h"
 #include <algorithm>
 #include <memory>
 #include <sstream>
 #include <string>
+#include <utility>
 #include <vector>
 #include "Eigen/Dense"
-
 #include "paddle/fluid/distributed/service/brpc_ps_client.h"
-#include "paddle/fluid/distributed/service/graph_brpc_client.h"
 #include "paddle/fluid/distributed/table/table.h"
 #include "paddle/fluid/framework/archive.h"
 #include "paddle/fluid/string/string_helper.h"
@@ -35,29 +35,26 @@ int GraphBrpcClient::get_server_index_by_id(uint64_t id) {
   return id % shard_num / shard_per_server;
 }
 // char* &buffer,int &actual_size
-std::future<int32_t> GraphBrpcClient::sample(uint32_t table_id,
-                                             uint64_t node_id, int sample_size,
-                                             std::vector<GraphNode> &res) {
+std::future<int32_t> GraphBrpcClient::sample(
+    uint32_t table_id, uint64_t node_id, int sample_size,
+    std::vector<std::pair<uint64_t, float>> &res) {
   int server_index = get_server_index_by_id(node_id);
   DownpourBrpcClosure *closure = new DownpourBrpcClosure(1, [&](void *done) {
     int ret = 0;
     auto *closure = (DownpourBrpcClosure *)done;
     if (closure->check_response(0, PS_GRAPH_SAMPLE) != 0) {
       ret = -1;
     } else {
-      VLOG(0) << "check sample response: "
-              << " " << closure->check_response(0, PS_GRAPH_SAMPLE);
       auto &res_io_buffer = closure->cntl(0)->response_attachment();
       butil::IOBufBytesIterator io_buffer_itr(res_io_buffer);
       size_t bytes_size = io_buffer_itr.bytes_left();
       char *buffer = new char[bytes_size];
       io_buffer_itr.copy_and_forward((void *)(buffer), bytes_size);
-      int start = 0;
-      while (start < bytes_size) {
-        GraphNode node;
-        node.recover_from_buffer(buffer + start);
-        start += node.get_size();
-        res.push_back(node);
+      int offset = 0;
+      while (offset < bytes_size) {
+        res.push_back({*(uint64_t *)(buffer + offset),
+                       *(float *)(buffer + offset + GraphNode::id_size)});
+        offset += GraphNode::id_size + GraphNode::weight_size;
       }
     }
     closure->set_promise_value(ret);
@@ -69,9 +66,7 @@ std::future<int32_t> GraphBrpcClient::sample(uint32_t table_id,
   closure->request(0)->set_cmd_id(PS_GRAPH_SAMPLE);
   closure->request(0)->set_table_id(table_id);
   closure->request(0)->set_client_id(_client_id);
-  // std::string type_str = GraphNode::node_type_to_string(type);
   closure->request(0)->add_params((char *)&node_id, sizeof(uint64_t));
-  // closure->request(0)->add_params(type_str.c_str(), type_str.size());
   closure->request(0)->add_params((char *)&sample_size, sizeof(int));
   PsService_Stub rpc_stub(get_cmd_channel(server_index));
   closure->cntl(0)->set_log_id(butil::gettimeofday_ms());

paddle/fluid/distributed/service/graph_brpc_client.h

@@ -18,6 +18,7 @@
 #include <string>
 #include <vector>
 
+#include <utility>
 #include "brpc/channel.h"
 #include "brpc/controller.h"
 #include "brpc/server.h"
@@ -35,9 +36,9 @@ class GraphBrpcClient : public BrpcPsClient {
  public:
   GraphBrpcClient() {}
   virtual ~GraphBrpcClient() {}
-  virtual std::future<int32_t> sample(uint32_t table_id, uint64_t node_id,
-                                      int sample_size,
-                                      std::vector<GraphNode> &res);
+  virtual std::future<int32_t> sample(
+      uint32_t table_id, uint64_t node_id, int sample_size,
+      std::vector<std::pair<uint64_t, float>> &res);
   virtual std::future<int32_t> pull_graph_list(uint32_t table_id,
                                                int server_index, int start,
                                                int size,

paddle/fluid/distributed/service/graph_py_service.h

@@ -21,8 +21,8 @@
 #include <sstream>
 #include <string>
 #include <thread>  // NOLINT
-#include <vector>
 #include <unordered_map>
+#include <vector>
 #include "google/protobuf/text_format.h"
 
 #include "gtest/gtest.h"
@@ -47,7 +47,7 @@ class GraphPyService {
   std::vector<int> keys;
   std::vector<std::string> server_list, port_list, host_sign_list;
   int server_size, shard_num, rank, client_id;
-  std::unordered_map<std::string, uint32_t > table_id_map;
+  std::unordered_map<std::string, uint32_t> table_id_map;
   std::thread *server_thread, *client_thread;
 
   std::shared_ptr<paddle::distributed::PSServer> pserver_ptr;
@@ -68,7 +68,8 @@ class GraphPyService {
   int get_shard_num() { return shard_num; }
   void set_shard_num(int shard_num) { this->shard_num = shard_num; }
   void GetDownpourSparseTableProto(
-      ::paddle::distributed::TableParameter* sparse_table_proto, uint32_t table_id) {
+      ::paddle::distributed::TableParameter* sparse_table_proto,
+      uint32_t table_id) {
     sparse_table_proto->set_table_id(table_id);
     sparse_table_proto->set_table_class("GraphTable");
     sparse_table_proto->set_shard_num(shard_num);
@@ -97,14 +98,13 @@ class GraphPyService {
     server_service_proto->set_start_server_port(0);
     server_service_proto->set_server_thread_num(12);
 
-    for(auto& tuple : this -> table_id_map) {
-        ::paddle::distributed::TableParameter* sparse_table_proto =
-            downpour_server_proto->add_downpour_table_param();
-        GetDownpourSparseTableProto(sparse_table_proto, tuple.second);
+    for (auto& tuple : this->table_id_map) {
+      ::paddle::distributed::TableParameter* sparse_table_proto =
+          downpour_server_proto->add_downpour_table_param();
+      GetDownpourSparseTableProto(sparse_table_proto, tuple.second);
     }
 
     return server_fleet_desc;
-
   }
 
   ::paddle::distributed::PSParameter GetWorkerProto() {
@@ -116,10 +116,10 @@ class GraphPyService {
     ::paddle::distributed::DownpourWorkerParameter* downpour_worker_proto =
         worker_proto->mutable_downpour_worker_param();
 
-    for(auto& tuple : this -> table_id_map) {
-        ::paddle::distributed::TableParameter* worker_sparse_table_proto =
-            downpour_worker_proto->add_downpour_table_param();
-        GetDownpourSparseTableProto(worker_sparse_table_proto, tuple.second);
+    for (auto& tuple : this->table_id_map) {
+      ::paddle::distributed::TableParameter* worker_sparse_table_proto =
+          downpour_worker_proto->add_downpour_table_param();
+      GetDownpourSparseTableProto(worker_sparse_table_proto, tuple.second);
     }
 
     ::paddle::distributed::ServerParameter* server_proto =
@@ -134,10 +134,10 @@ class GraphPyService {
     server_service_proto->set_start_server_port(0);
     server_service_proto->set_server_thread_num(12);
 
-    for(auto& tuple : this -> table_id_map) {
-        ::paddle::distributed::TableParameter* sparse_table_proto =
-            downpour_server_proto->add_downpour_table_param();
-        GetDownpourSparseTableProto(sparse_table_proto, tuple.second);
+    for (auto& tuple : this->table_id_map) {
+      ::paddle::distributed::TableParameter* sparse_table_proto =
+          downpour_server_proto->add_downpour_table_param();
+      GetDownpourSparseTableProto(sparse_table_proto, tuple.second);
     }
 
     return worker_fleet_desc;
@@ -148,44 +148,46 @@ class GraphPyService {
 
   void load_edge_file(std::string name, std::string filepath, bool reverse) {
     std::string params = "edge";
-    if(reverse) {
-        params += "|reverse";
+    if (reverse) {
+      params += "|reverse";
     }
-    if (this -> table_id_map.count(name)) {
-        uint32_t table_id =  this -> table_id_map[name];
-        auto status =
-            get_ps_client()->load(table_id, std::string(filepath), params);
-        status.wait();
+    if (this->table_id_map.count(name)) {
+      uint32_t table_id = this->table_id_map[name];
+      auto status =
+          get_ps_client()->load(table_id, std::string(filepath), params);
+      status.wait();
     }
   }
 
   void load_node_file(std::string name, std::string filepath) {
     std::string params = "node";
-    if (this -> table_id_map.count(name)) {
-        uint32_t table_id =  this -> table_id_map[name];
-        auto status =
-            get_ps_client()->load(table_id, std::string(filepath), params);
-        status.wait();
+    if (this->table_id_map.count(name)) {
+      uint32_t table_id = this->table_id_map[name];
+      auto status =
+          get_ps_client()->load(table_id, std::string(filepath), params);
+      status.wait();
     }
   }
 
-  std::vector<GraphNode> sample_k(std::string name, uint64_t node_id, int sample_size) {
-    std::vector<GraphNode> v;
-    if (this -> table_id_map.count(name)) {
-        uint32_t table_id =  this -> table_id_map[name];
-        auto status = worker_ptr->sample(table_id, node_id, sample_size, v);
-        status.wait();
+  std::vector<std::pair<uint64_t, float>> sample_k(std::string name,
+                                                   uint64_t node_id,
+                                                   int sample_size) {
+    std::vector<std::pair<uint64_t, float>> v;
+    if (this->table_id_map.count(name)) {
+      uint32_t table_id = this->table_id_map[name];
+      auto status = worker_ptr->sample(table_id, node_id, sample_size, v);
+      status.wait();
     }
     return v;
   }
-  std::vector<GraphNode> pull_graph_list(std::string name, int server_index, int start,
-                                         int size) {
+  std::vector<GraphNode> pull_graph_list(std::string name, int server_index,
+                                         int start, int size) {
     std::vector<GraphNode> res;
-    if (this -> table_id_map.count(name)) {
-        uint32_t table_id =  this -> table_id_map[name];
-        auto status =
-            worker_ptr->pull_graph_list(table_id, server_index, start, size, res);
-        status.wait();
+    if (this->table_id_map.count(name)) {
+      uint32_t table_id = this->table_id_map[name];
+      auto status =
+          worker_ptr->pull_graph_list(table_id, server_index, start, size, res);
+      status.wait();
     }
     return res;
   }

paddle/fluid/distributed/service/ps_client.h

@@ -155,9 +155,9 @@ class PSClient {
     promise.set_value(-1);
     return fut;
   }
-  virtual std::future<int32_t> sample(uint32_t table_id, uint64_t node_id,
-                                      int sample_size,
-                                      std::vector<GraphNode> &res) {
+  virtual std::future<int32_t> sample(
+      uint32_t table_id, uint64_t node_id, int sample_size,
+      std::vector<std::pair<uint64_t, float>> &res) {
     LOG(FATAL) << "Did not implement";
     std::promise<int32_t> promise;
     std::future<int> fut = promise.get_future();