diff --git a/query_modules/wcc.py b/query_modules/wcc.py new file mode 100644 index 000000000..4e16e1f36 --- /dev/null +++ b/query_modules/wcc.py @@ -0,0 +1,40 @@ +import mgp + +import networkx as nx + + +@mgp.read_proc +def get_components(vertices: mgp.List[mgp.Vertex], + edges: mgp.List[mgp.Edge] + ) -> mgp.Record(n_components=int, + components=mgp.List[mgp.List[mgp.Vertex]]): + ''' + This procedure finds weakly connected components of a given subgraph of a + directed graph. + + The subgraph is defined by a list of vertices and a list edges which are + passed as arguments of the procedure. More precisely, a set of vertices of + a subgraph contains all vertices provided in a list of vertices along with + all vertices that are endpoints of provided edges. Similarly, a set of + edges of a subgraph contains all edges from the list of provided edges. + + The procedure returns 2 fields: + * `n_components` is the number of weakly connected components of the + subgraph. + * `components` is a list of weakly connected components. Each component + is given as a list of `mgp.Vertex` objects from that component. + + For example, weakly connected components in a subgraph formed from all + vertices labeled `Person` and edges between such vertices can be obtained + using the following openCypher query: + + MATCH (n:Person)-[e]->(m:Person) WITH collect(n) AS nodes, collect(e) AS edges + CALL wcc.get_components(nodes, edges) YIELD * RETURN n_components, components; + ''' + g = nx.DiGraph() + g.add_nodes_from(vertices) + g.add_edges_from([(edge.from_vertex, edge.to_vertex) for edge in edges]) + + components = [list(wcc) for wcc in nx.weakly_connected_components(g)] + + return mgp.Record(n_components=len(components), components=components)