memgraph/query_modules/wcc.py
Marko Budiselić 48587d6d5e
Improve NetworkX module import (#21)
* Improve NetworkX module import
* Add Networkx dependencies to Dockerfiles
2020-10-15 09:14:50 +02:00

52 lines
1.9 KiB
Python

import sys
import mgp
try:
import networkx as nx
except ImportError as import_error:
sys.stderr.write(
'\n'
'NOTE: Please install networkx to be able to use wcc module.\n'
'Using Python:\n'
+ sys.version +
'\n')
raise import_error
@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)