random_sample return future

paddle/fluid/distributed/service/graph_brpc_server.cc

@@ -283,28 +283,22 @@ int32_t GraphBrpcService::graph_random_sample(Table *table,
         "graph_random_sample request requires at least 2 arguments");
     return 0;
   }
-  size_t num_nodes = request.params(0).size() / sizeof(uint64_t);
+  size_t node_num = request.params(0).size() / sizeof(uint64_t);
   uint64_t *node_data = (uint64_t *)(request.params(0).c_str());
   int sample_size = *(uint64_t *)(request.params(1).c_str());
 
-  std::vector<std::future<int>*> tasks;
-  std::vector<char*> buffers(num_nodes);
-  std::vector<int> actual_sizes(num_nodes);
+  std::vector<char*> buffers(node_num, nullptr);
+  std::vector<int> actual_sizes(node_num, 0);
+  table->random_sample(node_data, sample_size, buffers, actual_sizes);
 
-  for (size_t idx = 0; idx < num_nodes; ++idx){
-    //std::future<int> task = table->random_sample(node_data[idx], sample_size, 
-            //buffers[idx], actual_sizes[idx]);
-    table->random_sample(node_data[idx], sample_size, 
-            buffers[idx], actual_sizes[idx]);
-    //tasks.push_back(&task);
-  }
-  //for (size_t idx = 0; idx < num_nodes; ++idx){
-    //tasks[idx]->get();
-  //}
-  cntl->response_attachment().append(&num_nodes, sizeof(size_t));
-  cntl->response_attachment().append(actual_sizes.data(), sizeof(int)*num_nodes);
-  for (size_t idx = 0; idx < num_nodes; ++idx){
+  cntl->response_attachment().append(&node_num, sizeof(size_t));
+  cntl->response_attachment().append(actual_sizes.data(), sizeof(int)*node_num);
+  for (size_t idx = 0; idx < node_num; ++idx){
     cntl->response_attachment().append(buffers[idx], actual_sizes[idx]);
+    if (buffers[idx] != nullptr){
+      delete buffers[idx];
+      buffers[idx] = nullptr;
+    }
   }
   return 0;
 }

paddle/fluid/distributed/table/common_graph_table.cc

@@ -200,13 +200,16 @@ GraphNode *GraphTable::find_node(uint64_t id) {
 uint32_t GraphTable::get_thread_pool_index(uint64_t node_id) {
   return node_id % shard_num_per_table % task_pool_size_;
 }
-//std::future<int> GraphTable::random_sample(uint64_t node_id, int sample_size,
-                                  //char *&buffer, int &actual_size) {
-int32_t GraphTable::random_sample(uint64_t node_id, int sample_size,
-                                  char *&buffer, int &actual_size) {
-  return _shards_task_pool[get_thread_pool_index(node_id)]
+int GraphTable::random_sample(uint64_t* node_ids, int sample_size,
+                                  std::vector<char*>& buffers, std::vector<int> &actual_sizes) {
+  size_t node_num = buffers.size();
+  std::vector<std::future<int>> tasks;
+  for (size_t idx = 0; idx < node_num; ++idx){
+    uint64_t node_id = node_ids[idx];
+    char* & buffer = buffers[idx];
+    int& actual_size = actual_sizes[idx];
+    tasks.push_back(_shards_task_pool[get_thread_pool_index(node_id)]
       ->enqueue([&]() -> int {
-
         GraphNode *node = find_node(node_id);
         if (node == NULL) {
           actual_size = 0;
@@ -229,8 +232,12 @@ int32_t GraphTable::random_sample(uint64_t node_id, int sample_size,
           offset += GraphNode::weight_size;
         }
         return 0;
-      })
-      .get();
+      }));
+  }
+  for (size_t idx = 0; idx < node_num; ++idx){
+    tasks[idx].get();
+  }
+  return 0;
 }
 int32_t GraphTable::pull_graph_list(int start, int total_size, char *&buffer,
                                     int &actual_size) {

paddle/fluid/distributed/table/common_graph_table.h

@@ -71,10 +71,8 @@ class GraphTable : public SparseTable {
   virtual ~GraphTable() {}
   virtual int32_t pull_graph_list(int start, int size, char *&buffer,
                                   int &actual_size);
-  //virtual std::future<int> random_sample(uint64_t node_id, int sampe_size, char *&buffer,
-                                //int &actual_size);
-  virtual int32_t random_sample(uint64_t node_id, int sampe_size, char *&buffer,
-                                int &actual_size);
+  virtual int random_sample(uint64_t* node_ids, int sampe_size, std::vector<char *>&buffers,
+                                std::vector<int> &actual_sizes);
   virtual int32_t initialize();
 
   int32_t load(const std::string &path, const std::string &param);

paddle/fluid/distributed/table/table.h

@@ -93,14 +93,10 @@ class Table {
     return 0;
   }
   // only for graph table
-  virtual int32_t random_sample(uint64_t node_id, int sampe_size, char *&buffer,
-                                int &actual_size) {
+  virtual int random_sample(uint64_t* node_ids, int sampe_size, std::vector<char *>&buffers,
+                                std::vector<int> &actual_sizes) {
     return 0;
   }
-  //virtual std::future<int> random_sample(uint64_t node_id, int sampe_size, char *&buffer,
-                                //int &actual_size) {
-    //return std::future<int>();
-  //}
   virtual int32_t pour() { return 0; }
 
   virtual void clear() = 0;

paddle/fluid/distributed/test/graph_node_test.cc

@@ -236,6 +236,11 @@ void RunBrpcPushSparse() {
     ASSERT_EQ(true, s.find(g.first) != s.end());
   }
   vs.clear();
+
+  pull_status = worker_ptr_->batch_sample(0, std::vector<uint64_t>(1, 10240001024), 4, vs);
+  pull_status.wait();
+  ASSERT_EQ(0, vs[0].size());
+
   std::vector<distributed::GraphNode> nodes;
   pull_status = worker_ptr_->pull_graph_list(0, 0, 0, 1, nodes);
   pull_status.wait();