Improve calculateQ efficiency in Louvain (#614)

alibaba · Aug 10, 2021 · 35b2fdf · 35b2fdf
1 parent 79ec800
commit 35b2fdf
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
diff --git a/analytical_engine/apps/pregel/louvain/louvain.h b/analytical_engine/apps/pregel/louvain/louvain.h
@@ -17,6 +17,7 @@
 #define ANALYTICAL_ENGINE_APPS_PREGEL_LOUVAIN_LOUVAIN_H_
 
 #include <map>
+#include <set>
 #include <utility>
 #include <vector>
 
@@ -336,11 +337,13 @@ class PregelLouvain
                                 const grape::IteratorPair<md_t*>& messages) {
     auto& state = vertex.state();
     edata_t k_i_in = state.internal_weight;
+    std::set<vid_t> source_ids;
     for (auto& m : messages) {
       if (m.community_id == state.community) {
-        k_i_in += vertex.get_edge_value(m.source_id);
+        source_ids.insert(m.source_id);
       }
     }
+    k_i_in += vertex.get_edge_values(source_ids);
     edata_t sigma_tot = state.community_sigma_total;
     edata_t m2 = getTotalEdgeWeight(context, vertex);
     edata_t k_i = state.node_weight + state.internal_weight;

diff --git a/analytical_engine/apps/pregel/louvain/louvain_vertex.h b/analytical_engine/apps/pregel/louvain/louvain_vertex.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define ANALYTICAL_ENGINE_APPS_PREGEL_LOUVAIN_LOUVAIN_VERTEX_H_
 
 #include <map>
+#include <set>
 #include <string>
 #include <utility>
 #include <vector>
@@ -99,6 +100,35 @@ class LouvainVertex : public PregelVertex<FRAG_T, VD_T, MD_T> {
     return edata_t();
   }
 
+  edata_t get_edge_values(const std::set<vid_t>& dst_ids) {
+    edata_t ret = 0;
+    if (!this->use_fake_edges()) {
+      for (auto& edge : this->incoming_edges()) {
+        auto gid = this->fragment_->Vertex2Gid(edge.get_neighbor());
+        if (dst_ids.find(gid) != dst_ids.end()) {
+          ret += static_cast<edata_t>(edge.get_data());
+        }
+      }
+      for (auto& edge : this->outgoing_edges()) {
+        auto gid = this->fragment_->Vertex2Gid(edge.get_neighbor());
+        if (dst_ids.find(gid) != dst_ids.end()) {
+          ret += static_cast<edata_t>(edge.get_data());
+        }
+      }
+    } else {
+      auto edges = this->fake_edges();
+      for (auto gid : dst_ids) {
+        if (edges.find(gid) != edges.end()) {
+          ret += edges.at(gid);
+        } else {
+          LOG(ERROR) << "Warning: Cannot find a edge from " << this->id()
+                     << " to " << gid;
+        }
+      }
+    }
+    return ret;
+  }
+
   void set_fake_edges(std::map<vid_t, edata_t>&& edges) {
     state_t& ref_state = this->state();
     ref_state.fake_edges = std::move(edges);