Expose modularity value in louvain query module

Reviewers: teon.banek

Reviewed By: teon.banek

Subscribers: mferencevic, pullbot

Differential Revision: https://phabricator.memgraph.io/D2591

query_modules/louvain/src/louvain_module.cpp

@@ -6,18 +6,18 @@
 #include "algorithms/algorithms.hpp"
 #include "data_structures/graph.hpp"
 
-static void communities(const mgp_list *args, const mgp_graph *graph,
+namespace {
-                        mgp_result *result, mgp_memory *memory) {
+
+std::optional<std::unordered_map<int64_t, uint32_t>> NormalizeVertexIds(
+    const mgp_graph *graph, mgp_result *result, mgp_memory *memory) {
+  std::unordered_map<int64_t, uint32_t> mem_to_louv_id;
   mgp_vertices_iterator *vertices_iterator =
       mgp_graph_iter_vertices(graph, memory);
   if (vertices_iterator == nullptr) {
     mgp_result_set_error_msg(result, "Not enough memory!");
-    return;
+    return std::nullopt;
   }
 
-  // Normalize vertex ids
-  std::unordered_map<int64_t, uint32_t> mem_to_louv_id;
-
   uint32_t louv_id = 0;
   for (const mgp_vertex *vertex = mgp_vertices_iterator_get(vertices_iterator);
        vertex != nullptr;
@@ -28,15 +28,20 @@ static void communities(const mgp_list *args, const mgp_graph *graph,
   }
 
   mgp_vertices_iterator_destroy(vertices_iterator);
+  return mem_to_louv_id;
+}
 
+std::optional<comdata::Graph> RunLouvain(
+    const mgp_graph *graph, mgp_result *result, mgp_memory *memory,
+    const std::unordered_map<int64_t, uint32_t> &mem_to_louv_id) {
+  comdata::Graph louvain_graph(mem_to_louv_id.size());
   // Extract the graph structure
   // TODO(ipaljak): consider filtering nodes and edges by labels.
-  comdata::Graph louvain_graph(louv_id);
   for (const auto &p : mem_to_louv_id) {
     mgp_vertex *vertex = mgp_graph_get_vertex_by_id(graph, {p.first}, memory);
     if (!vertex) {
       mgp_result_set_error_msg(result, "Not enough memory!");
-      return;
+      return std::nullopt;
     }
 
     // iterate over inbound edges. This is enough because we will eventually
@@ -46,14 +51,21 @@ static void communities(const mgp_list *args, const mgp_graph *graph,
     if (edges_iterator == nullptr) {
       mgp_vertex_destroy(vertex);
       mgp_result_set_error_msg(result, "Not enough memory!");
-      return;
+      return std::nullopt;
     }
 
     for (const mgp_edge *edge = mgp_edges_iterator_get(edges_iterator);
          edge != nullptr; edge = mgp_edges_iterator_next(edges_iterator)) {
       const mgp_vertex *next_vertex = mgp_edge_get_from(edge);
       mgp_vertex_id next_mem_id = mgp_vertex_get_id(next_vertex);
-      uint32_t next_louv_id = mem_to_louv_id[next_mem_id.as_int];
+      uint32_t next_louv_id;
+      try {
+        next_louv_id = mem_to_louv_id.at(next_mem_id.as_int);
+      } catch (const std::exception &e) {
+        const auto msg = std::string("[Internal error] ") + e.what();
+        mgp_result_set_error_msg(result, msg.c_str());
+        return std::nullopt;
+      }
 
       // retrieve edge weight (default to 1)
       mgp_value *weight_prop = mgp_edge_get_property(edge, "weight", memory);
@@ -77,7 +89,7 @@ static void communities(const mgp_list *args, const mgp_graph *graph,
         mgp_vertex_destroy(vertex);
         mgp_edges_iterator_destroy(edges_iterator);
         mgp_result_set_error_msg(result, e.what());
-        return;
+        return std::nullopt;
       }
     }
 
@@ -90,53 +102,124 @@ static void communities(const mgp_list *args, const mgp_graph *graph,
   } catch (const std::exception &e) {
     const auto msg = std::string("[Internal error] ") + e.what();
     mgp_result_set_error_msg(result, msg.c_str());
-    return;
+    return std::nullopt;
   }
 
-  // Return node ids and their corresponding communities.
+  return louvain_graph;
-  for (const auto &p : mem_to_louv_id) {
+}
+
+void communities(const mgp_list *args, const mgp_graph *graph,
+                        mgp_result *result, mgp_memory *memory) {
+  try {
+    // Normalize vertex ids
+    auto mem_to_louv_id = NormalizeVertexIds(graph, result, memory);
+    if (!mem_to_louv_id) return;
+
+    // Run louvain
+    auto louvain_graph = RunLouvain(graph, result, memory, *mem_to_louv_id);
+    if (!louvain_graph) return;
+
+    // Return node ids and their corresponding communities.
+    for (const auto &p : *mem_to_louv_id) {
+      mgp_result_record *record = mgp_result_new_record(result);
+      if (record == nullptr) {
+        mgp_result_set_error_msg(result, "Not enough memory!");
+        return;
+      }
+
+      mgp_value *mem_id_value = mgp_value_make_int(p.first, memory);
+      if (mem_id_value == nullptr) {
+        mgp_result_set_error_msg(result, "Not enough memory!");
+        return;
+      }
+
+      mgp_value *com_value =
+          mgp_value_make_int(louvain_graph->Community(p.second), memory);
+      if (com_value == nullptr) {
+        mgp_value_destroy(mem_id_value);
+        mgp_result_set_error_msg(result, "Not enough memory!");
+        return;
+      }
+
+      int mem_id_inserted =
+          mgp_result_record_insert(record, "id", mem_id_value);
+      int com_inserted =
+          mgp_result_record_insert(record, "community", com_value);
+
+      mgp_value_destroy(mem_id_value);
+      mgp_value_destroy(com_value);
+
+      if (!mem_id_inserted || !com_inserted) {
+        mgp_result_set_error_msg(result, "Not enough memory!");
+        return;
+      }
+    }
+  } catch (const std::exception &e) {
+    mgp_result_set_error_msg(result, e.what());
+    return;
+  }
+}
+
+void modularity(const mgp_list *args, const mgp_graph *graph,
+                       mgp_result *result, mgp_memory *memory) {
+  try {
+    // Normalize vertex ids
+    auto mem_to_louv_id = NormalizeVertexIds(graph, result, memory);
+    if (!mem_to_louv_id) return;
+
+    // Run louvain
+    auto louvain_graph = RunLouvain(graph, result, memory, *mem_to_louv_id);
+    if (!louvain_graph) return;
+
+    // Return graph modularity after Louvain
+    // TODO(ipaljak) - consider allowing the user to specify seed communities and
+    //                 yield modularity values both before and after running
+    //                 louvain.
     mgp_result_record *record = mgp_result_new_record(result);
     if (record == nullptr) {
       mgp_result_set_error_msg(result, "Not enough memory!");
       return;
     }
 
-    mgp_value *mem_id_value = mgp_value_make_int(p.first, memory);
+    mgp_value *modularity_value =
-    if (mem_id_value == nullptr) {
+        mgp_value_make_double(louvain_graph->Modularity(), memory);
+    if (modularity_value == nullptr) {
       mgp_result_set_error_msg(result, "Not enough memory!");
       return;
     }
 
-    mgp_value *com_value =
+    int value_inserted =
-        mgp_value_make_int(louvain_graph.Community(p.second), memory);
+        mgp_result_record_insert(record, "modularity", modularity_value);
-    if (com_value == nullptr) {
+
-      mgp_value_destroy(mem_id_value);
+    mgp_value_destroy(modularity_value);
-      mgp_result_set_error_msg(result, "Not enough memory!");
+
-      return;
+    if (!value_inserted) {
-    }
-
-    int mem_id_inserted =
-        mgp_result_record_insert(record, "id", mem_id_value);
-    int com_inserted =
-        mgp_result_record_insert(record, "community", com_value);
-
-    mgp_value_destroy(mem_id_value);
-    mgp_value_destroy(com_value);
-
-    if (!mem_id_inserted || !com_inserted) {
       mgp_result_set_error_msg(result, "Not enough memory!");
       return;
     }
+  } catch (const std::exception &e) {
+    mgp_result_set_error_msg(result, e.what());
+    return;
   }
 }
 
+} // namespace
+
 extern "C" int mgp_init_module(struct mgp_module *module,
                                struct mgp_memory *memory) {
-  struct mgp_proc *proc =
+  struct mgp_proc *community_proc =
       mgp_module_add_read_procedure(module, "communities", communities);
-  if (!proc) return 1;
+  if (!community_proc) return 1;
-  if (!mgp_proc_add_result(proc, "id", mgp_type_int())) return 1;
+  if (!mgp_proc_add_result(community_proc, "id", mgp_type_int())) return 1;
-  if (!mgp_proc_add_result(proc, "community", mgp_type_int())) return 1;
+  if (!mgp_proc_add_result(community_proc, "community", mgp_type_int()))
+    return 1;
+
+  struct mgp_proc *modularity_proc =
+    mgp_module_add_read_procedure(module, "modularity", modularity);
+  if (!modularity_proc) return 1;
+  if (!mgp_proc_add_result(modularity_proc, "modularity", mgp_type_float()))
+    return 1;
+
   return 0;
 }