fix the unbalance issue

Author: lqhl

src/main.cpp

@@ -15,6 +15,7 @@ DEFINE_uint64(memsize, 4096, "memory size in megabytes");
 DEFINE_string(filename, "", "the file name of the input graph");
 DEFINE_string(filetype, "edgelist",
               "the type of input file (supports 'edgelist' and 'adjlist')");
+DEFINE_bool(inmem, false, "in-memory mode");
 DEFINE_double(sample_ratio, 2, "the sample size divided by num_vertices");
 DEFINE_string(method, "neighbor",
               "partition method: neighbor, random, and dbh");

src/neighbor_partitioner.cpp

@@ -29,14 +29,22 @@ NeighborPartitioner::NeighborPartitioner(std::string basefilename)
 
     LOG(INFO) << "num_vertices: " << num_vertices
               << ", num_edges: " << num_edges;
+    CHECK_EQ(sizeof(vid_t) + sizeof(size_t) + num_edges * sizeof(edge_t), filesize);
 
     p = FLAGS_p;
     average_degree = (double)num_edges * 2 / num_vertices;
     assigned_edges = 0;
-    LOG(INFO) << "sample_ratio: " << FLAGS_sample_ratio;
-    max_sample_size = num_vertices * FLAGS_sample_ratio;
+    LOG(INFO) << "inmem: " << FLAGS_inmem;
+    if (!FLAGS_inmem) {
+        LOG(INFO) << "sample_ratio: " << FLAGS_sample_ratio;
+        max_sample_size = num_vertices * FLAGS_sample_ratio;
+    } else
+        max_sample_size = num_edges;
     local_average_degree = 2 * (double)max_sample_size / num_vertices;
     capacity = (double)num_edges * 1.05 / p + 1;
+    BUFFER_SIZE = std::min(64 * 1024 / sizeof(edge_t),
+                           std::max((size_t)1, (size_t)(num_edges * 0.05 / p + 1)));
+    LOG(INFO) << "buffer size: " << BUFFER_SIZE;
     occupied.assign(p, 0);
     adj_out.resize(num_vertices);
     adj_in.resize(num_vertices);
@@ -122,27 +130,36 @@ void NeighborPartitioner::read_remaining()
     }
 }
 
-void NeighborPartitioner::clean_samples()
+void NeighborPartitioner::clean_buffer()
 {
-    repv (u, num_vertices)
-        for (auto &v : adj_out[u])
-            sample_edges.emplace_back(u, v);
-
-    results.resize(sample_edges.size());
+    results.resize(buffer.size());
 
 #pragma omp parallel for
-    for (size_t i = 0; i < sample_edges.size(); i++)
-        results[i] = check_edge(&sample_edges[i]);
+    for (size_t i = 0; i < buffer.size(); i++)
+        results[i] = check_edge(&buffer[i]);
 
-    for (size_t i = 0; i < sample_edges.size();)
+    for (size_t i = 0; i < buffer.size();)
         if (results[i] < p) {
-            assign_edge(results[i], sample_edges[i].first, sample_edges[i].second);
+            assign_edge(results[i], buffer[i].first, buffer[i].second);
             std::swap(results[i], results.back());
             results.pop_back();
-            std::swap(sample_edges[i], sample_edges.back());
-            sample_edges.pop_back();
+            std::swap(buffer[i], buffer.back());
+            buffer.pop_back();
         } else
             i++;
+    sample_edges.insert(sample_edges.end(), buffer.begin(), buffer.end());
+    buffer.clear();
+}
+
+void NeighborPartitioner::clean_samples()
+{
+    repv (u, num_vertices)
+        for (auto &v : adj_out[u]) {
+            buffer.emplace_back(u, v);
+            if (buffer.size() >= BUFFER_SIZE)
+                clean_buffer();
+        }
+    clean_buffer();
 }
 
 void NeighborPartitioner::assign_master()
@@ -205,7 +222,8 @@ void NeighborPartitioner::split()
         read_timer.stop();
         DLOG(INFO) << "sample size: " << adj_out.num_edges();
         compute_timer.start();
-        local_capacity = adj_out.num_edges() / (p - bucket);
+        local_capacity =
+            FLAGS_inmem ? capacity : adj_out.num_edges() / (p - bucket);
         while (occupied[bucket] < local_capacity) {
             vid_t d, vid;
             if (!min_heap.get_min(d, vid)) {

src/neighbor_partitioner.hpp

@@ -105,7 +105,7 @@ class graph_t
 class NeighborPartitioner : public Partitioner
 {
   private:
-    const size_t BUFFER_SIZE = 64 * 1024 / sizeof(edge_t);
+    size_t BUFFER_SIZE;
     std::string basefilename;
 
     vid_t num_vertices;
@@ -120,6 +120,7 @@ class NeighborPartitioner : public Partitioner
     off_t filesize;
     char *fin_map, *fin_ptr, *fin_end;
 
+    std::vector<edge_t> buffer;
     std::vector<edge_t> sample_edges;
     graph_t adj_out, adj_in;
     MinHeap<vid_t, vid_t> min_heap;
@@ -268,6 +269,7 @@ class NeighborPartitioner : public Partitioner
 
     void read_more();
     void read_remaining();
+    void clean_buffer();
     void clean_samples();
     void assign_master();
     size_t count_mirrors();

src/util.hpp

@@ -17,6 +17,7 @@ DECLARE_int32(p);
 DECLARE_uint64(memsize);
 DECLARE_string(filename);
 DECLARE_string(filetype);
+DECLARE_bool(inmem);
 DECLARE_double(sample_ratio);
 
 typedef uint32_t vid_t;