diff --git a/grand/backends/_dynamodb.py b/grand/backends/_dynamodb.py index 98e8eae..73f67c1 100644 --- a/grand/backends/_dynamodb.py +++ b/grand/backends/_dynamodb.py @@ -1,6 +1,5 @@ from typing import Collection, Hashable, Optional import time -import concurrent.futures import pandas as pd import boto3 @@ -10,9 +9,6 @@ _DEFAULT_DYNAMODB_URL = "http://localhost:4566" -_N_PARALLEL_REQUESTS = 16 - - def _dynamo_table_exists(table_name: str, client: boto3.client): """ Check to see if the DynamoDB table already exists. @@ -468,61 +464,38 @@ def ingest_from_edgelist_dataframe( edge_column_names = [ c for c in edgelist.columns if c not in [source_column, target_column] ] + sources = edgelist[source_column].tolist() + targets = edgelist[target_column].tolist() + edge_metadata = ( + edgelist[edge_column_names].to_dict("records") + if edge_column_names + else [{} for _ in range(len(edgelist))] + ) tic = time.time() with self._edge_table.batch_writer() as batch_writer: - - with concurrent.futures.ThreadPoolExecutor( - _N_PARALLEL_REQUESTS - ) as executor: - result_futures = list( - map( - lambda i: executor.submit( - lambda: batch_writer.put_item( - Item={ - self._primary_key: f"__{edgelist._get_value(i, source_column)}__{edgelist._get_value(i, target_column)}", - self._edge_source_key: edgelist._get_value( - i, source_column - ), - self._edge_target_key: edgelist._get_value( - i, target_column - ), - **{ - col: edgelist._get_value(i, col) - for col in edge_column_names - }, - } - ) - ), - edgelist.index, - ) + for source, target, metadata in zip(sources, targets, edge_metadata): + batch_writer.put_item( + Item={ + self._primary_key: f"__{source}__{target}", + self._edge_source_key: source, + self._edge_target_key: target, + **metadata, + } ) - for future in concurrent.futures.as_completed(result_futures): - future.result() - - # for edge in edges: - # batch_writer.put_item(Item=edge) edge_toc = time.time() - tic tic = time.time() # Construct a unique set of nodes: - nodes = edgelist[source_column].append(edgelist[target_column]).unique() + nodes = pd.unique( + pd.concat( + [edgelist[source_column], edgelist[target_column]], + ignore_index=True, + ) + ) with self._node_table.batch_writer() as batch_writer: - with concurrent.futures.ThreadPoolExecutor( - _N_PARALLEL_REQUESTS - ) as executor: - result_futures = list( - map( - lambda x: executor.submit( - lambda: batch_writer.put_item( - Item={self._primary_key: str(x)} - ) - ), - nodes, - ) - ) - for future in concurrent.futures.as_completed(result_futures): - future.result() + for node in nodes: + batch_writer.put_item(Item={self._primary_key: str(node)}) return { "node_count": len(nodes), diff --git a/grand/backends/_igraph.py b/grand/backends/_igraph.py index bd7e82d..abd0b4f 100644 --- a/grand/backends/_igraph.py +++ b/grand/backends/_igraph.py @@ -1,6 +1,6 @@ from typing import Hashable, Collection -from igraph import Graph, InternalError +from igraph import Graph import pandas as pd from .backend import Backend @@ -35,6 +35,13 @@ def __init__(self, directed: bool = False): """ self._directed = directed self._ig = Graph(directed=self._directed) + self._node_ids_by_name = {} + self._edge_ids_by_key = {} + + def _edge_key(self, u: Hashable, v: Hashable): + if self._directed: + return (u, v) + return frozenset((u, v)) def ingest_from_edgelist_dataframe( self, edgelist: pd.DataFrame, source_column: str, target_column: str @@ -70,12 +77,15 @@ def add_node(self, node_name: Hashable, metadata: dict): Hashable: The ID of this node, as inserted """ + metadata = metadata or {} + if self.has_node(node_name): # Update metadata - m = self._ig.vs.find(name=node_name) - m.update_attributes(metadata) + self._ig.vs[self._node_ids_by_name[node_name]].update_attributes(metadata) return node_name + self._ig.add_vertex(name=node_name, **metadata) + self._node_ids_by_name[node_name] = self._ig.vcount() - 1 return node_name def get_node_by_id(self, node_name: Hashable): @@ -90,7 +100,7 @@ def get_node_by_id(self, node_name: Hashable): """ return _remove_name_from_attributes( - self._ig.vs.find(name=node_name).attributes() + self._ig.vs[self._node_ids_by_name[node_name]].attributes() ) def all_nodes_as_iterable(self, include_metadata: bool = False) -> Collection: @@ -122,11 +132,7 @@ def has_node(self, u: Hashable) -> bool: Returns: bool: True if the node exists """ - try: - self._ig.vs.find(name=u) - return True - except: - return False + return u in self._node_ids_by_name def add_edge(self, u: Hashable, v: Hashable, metadata: dict): """ @@ -144,16 +150,24 @@ def add_edge(self, u: Hashable, v: Hashable, metadata: dict): Hashable: The edge ID, as inserted. """ - if self.has_edge(u, v): - # Update metadata - e = self._ig.get_eid(u, v) - self._ig.es[e].update_attributes(metadata) + metadata = metadata or {} + edge_key = self._edge_key(u, v) + + if edge_key in self._edge_ids_by_key: + self._ig.es[self._edge_ids_by_key[edge_key]].update_attributes(metadata) return if not self.has_node(u): self.add_node(u, {}) if not self.has_node(v): self.add_node(v, {}) - return self._ig.add_edge(source=u, target=v, **metadata) + + new_edge = self._ig.add_edge( + source=self._node_ids_by_name[u], + target=self._node_ids_by_name[v], + **metadata, + ) + self._edge_ids_by_key[edge_key] = self._ig.ecount() - 1 + return new_edge def all_edges_as_iterable(self, include_metadata: bool = False) -> Collection: """ @@ -174,13 +188,7 @@ def all_edges_as_iterable(self, include_metadata: bool = False) -> Collection: yield e.source_vertex["name"], e.target_vertex["name"] def has_edge(self, u, v): - try: - self._ig.get_eid(u, v) - return True - except (InternalError, ValueError): - # InternalError means no such edge - # ValueError means one vertex doesn't exist - return False + return self._edge_key(u, v) in self._edge_ids_by_key def get_edge_by_id(self, u: Hashable, v: Hashable): """ @@ -194,10 +202,10 @@ def get_edge_by_id(self, u: Hashable, v: Hashable): dict: Metadata associated with this edge """ - try: - return self._ig.es[self._ig.get_eid(u, v)].attributes() - except InternalError: + edge_id = self._edge_ids_by_key.get(self._edge_key(u, v)) + if edge_id is None: raise IndexError(f"The edge ({u}, {v}) is not in the graph.") + return self._ig.es[edge_id].attributes() def get_node_successors( self, u: Hashable, include_metadata: bool = False @@ -217,13 +225,15 @@ def get_node_neighbors( Generator """ + u_id = self._node_ids_by_name[u] + if include_metadata: return { - self._ig.vs[s]["name"]: self.get_edge_by_id(u, s) - for s in self._ig.successors(u) + self._ig.vs[s]["name"]: self.get_edge_by_id(u, self._ig.vs[s]["name"]) + for s in self._ig.successors(u_id) } else: - return iter([self._ig.vs[s]["name"] for s in self._ig.successors(u)]) + return iter([self._ig.vs[s]["name"] for s in self._ig.successors(u_id)]) def get_node_predecessors( self, u: Hashable, include_metadata: bool = False @@ -238,13 +248,15 @@ def get_node_predecessors( Generator """ + u_id = self._node_ids_by_name[u] + if include_metadata: return { - self._ig.vs[s]["name"]: self.get_edge_by_id(s, u) - for s in self._ig.predecessors(u) + self._ig.vs[s]["name"]: self.get_edge_by_id(self._ig.vs[s]["name"], u) + for s in self._ig.predecessors(u_id) } else: - return iter([self._ig.vs[s]["name"] for s in self._ig.predecessors(u)]) + return iter([self._ig.vs[s]["name"] for s in self._ig.predecessors(u_id)]) def get_node_count(self) -> int: """ diff --git a/grand/backends/_networkx.py b/grand/backends/_networkx.py index b012fef..702bb96 100644 --- a/grand/backends/_networkx.py +++ b/grand/backends/_networkx.py @@ -190,23 +190,26 @@ def ingest_from_edgelist_dataframe( """ tic = time.time() - self._nx_graph.add_edges_from( - [ - ( - e[source_column], - e[target_column], - { - k: v - for k, v in dict(e).items() - if k not in [source_column, target_column] - }, - ) - for _, e in edgelist.iterrows() - ] + edge_columns = [ + c for c in edgelist.columns if c not in [source_column, target_column] + ] + sources = edgelist[source_column].tolist() + targets = edgelist[target_column].tolist() + + if edge_columns: + self._nx_graph.add_edges_from( + zip(sources, targets, edgelist[edge_columns].to_dict("records")) + ) + else: + self._nx_graph.add_edges_from(zip(sources, targets)) + + nodes = pd.unique( + pd.concat( + [edgelist[source_column], edgelist[target_column]], + ignore_index=True, + ) ) - nodes = edgelist[source_column].append(edgelist[target_column]).unique() - return { "node_count": len(nodes), "node_duration": 0, diff --git a/grand/backends/_sqlbackend.py b/grand/backends/_sqlbackend.py index 6e68e7a..b086f2e 100644 --- a/grand/backends/_sqlbackend.py +++ b/grand/backends/_sqlbackend.py @@ -154,6 +154,32 @@ def add_node(self, node_name: Hashable, metadata: dict) -> Hashable: ) return node_name + def _insert_empty_node_if_missing(self, node_name: Hashable) -> None: + node_name = str(node_name) + insert_statement = self._node_table.insert() + + # Use dialect-native "ignore duplicate key" behavior when available to + # avoid the extra existence query on every edge insertion. + if self._engine.dialect.name == "sqlite": + self._connection.execute( + insert_statement.prefix_with("OR IGNORE"), + parameters={self._primary_key: node_name, "_metadata": {}}, + ) + return + + if self._engine.dialect.name in {"mysql", "mariadb"}: + self._connection.execute( + insert_statement.prefix_with("IGNORE"), + parameters={self._primary_key: node_name, "_metadata": {}}, + ) + return + + if not self.has_node(node_name): + self._connection.execute( + insert_statement, + parameters={self._primary_key: node_name, "_metadata": {}}, + ) + def add_nodes_from(self, nodes_for_adding, **attr): nodes = [ { @@ -249,13 +275,14 @@ def has_node(self, u: Hashable) -> bool: Returns: bool: True if the node exists """ - return len( + return ( self._connection.execute( - self._node_table.select().where( + select(self._node_table.c[self._primary_key]).where( self._node_table.c[self._primary_key] == str(u) ) - ).fetchall() - ) > 0 + ).fetchone() + is not None + ) def add_edge(self, u: Hashable, v: Hashable, metadata: dict): """ @@ -275,10 +302,8 @@ def add_edge(self, u: Hashable, v: Hashable, metadata: dict): """ pk = f"__{u}__{v}" - if not self.has_node(u): - self.add_node(u, {}) - if not self.has_node(v): - self.add_node(v, {}) + self._insert_empty_node_if_missing(u) + self._insert_empty_node_if_missing(v) try: self._connection.execute( @@ -625,7 +650,7 @@ def in_degrees(self, nbunch=None): def ingest_from_edgelist_dataframe( self, edgelist: pd.DataFrame, source_column: str, target_column: str - ) -> None: + ) -> dict: """ Ingest an edgelist from a Pandas DataFrame. @@ -633,60 +658,61 @@ def ingest_from_edgelist_dataframe( # Produce edge list: edge_tic = time.time() - newlist = edgelist.rename( - columns={ - source_column: self._edge_source_key, - target_column: self._edge_target_key, - } + edge_columns = [ + c for c in edgelist.columns if c not in [source_column, target_column] + ] + sources = [str(value) for value in edgelist[source_column].tolist()] + targets = [str(value) for value in edgelist[target_column].tolist()] + edge_metadata = ( + edgelist[edge_columns].to_dict("records") + if edge_columns + else [{} for _ in range(len(edgelist))] ) - newlist[self._primary_key] = edgelist.apply( - lambda x: f"__{x[source_column]}__{x[target_column]}", axis="columns" - ) - newlist["_metadata"] = edgelist.apply( - lambda x: { - k: v for k, v in x.items() if k not in [source_column, target_column] - }, - axis="columns", - ) + edge_rows = [ + { + self._edge_source_key: source, + self._edge_target_key: target, + self._primary_key: f"__{source}__{target}", + "_metadata": metadata, + } + for source, target, metadata in zip(sources, targets, edge_metadata) + ] - newlist[ - [ - self._edge_source_key, - self._edge_target_key, - self._primary_key, - "_metadata", - ] - ].to_sql( - self._edge_table_name, - self._engine, - index=False, - if_exists="append", - dtype={"_metadata": sqlalchemy.JSON}, - ) + if edge_rows: + self._connection.execute(self._edge_table.insert(), edge_rows) edge_toc = time.time() - edge_tic # now ingest nodes: node_tic = time.time() - nodes = edgelist[source_column].append(edgelist[target_column]).unique() - pd.DataFrame( - [ - { - self._primary_key: node, - # no metadata: - "_metadata": {}, - } - for node in nodes - ] - ).to_sql( - self._node_table_name, - self._engine, - index=False, - if_exists="replace", - dtype={"_metadata": sqlalchemy.JSON}, + nodes = pd.unique( + pd.concat( + [edgelist[source_column], edgelist[target_column]], + ignore_index=True, + ) ) + node_rows = [ + { + self._primary_key: str(node), + "_metadata": {}, + } + for node in nodes + ] + + if node_rows: + node_insert = self._node_table.insert() + if self._engine.dialect.name == "sqlite": + node_insert = node_insert.prefix_with("OR IGNORE") + self._connection.execute(node_insert, node_rows) + elif self._engine.dialect.name in {"mysql", "mariadb"}: + node_insert = node_insert.prefix_with("IGNORE") + self._connection.execute(node_insert, node_rows) + else: + for node in nodes: + self._insert_empty_node_if_missing(node) + return { "node_count": len(nodes), "node_duration": time.time() - node_tic, diff --git a/grand/backends/test_backends.py b/grand/backends/test_backends.py index f821794..dc88c9b 100644 --- a/grand/backends/test_backends.py +++ b/grand/backends/test_backends.py @@ -455,6 +455,49 @@ def test_get_density_performance(backend): assert nx.density(G.nx) <= 0.005 +@pytest.mark.benchmark +def test_networkx_ingest_performance(): + backend = NetworkXBackend(directed=True) + edgelist = pd.DataFrame( + { + "source": range(1000), + "target": range(1, 1001), + "weight": range(1000), + "kind": ["edge"] * 1000, + } + ) + + result = backend.ingest_from_edgelist_dataframe(edgelist, "source", "target") + + assert result["node_count"] == 1001 + assert result["edge_count"] == 1000 + assert backend.get_node_count() == 1001 + assert backend.get_edge_count() == 1000 + + +@pytest.mark.benchmark +def test_sql_ingest_performance(): + if not _CAN_IMPORT_SQL: + pytest.skip("sqlalchemy is not installed.") + + backend = SQLBackend(directed=True, db_url="sqlite:///:memory:") + edgelist = pd.DataFrame( + { + "source": range(1000), + "target": range(1, 1001), + "weight": range(1000), + "kind": ["edge"] * 1000, + } + ) + + result = backend.ingest_from_edgelist_dataframe(edgelist, "source", "target") + + assert result["node_count"] == 1001 + assert result["edge_count"] == 1000 + assert backend.get_node_count() == 1001 + assert backend.get_edge_count() == 1000 + + class TestDataFrameBackend: def test_can_create_empty(self): b = DataFrameBackend() @@ -485,3 +528,46 @@ def test_can_create_from_int_dataframes(self): b = DataFrameBackend(edge_df=edges, node_df=nodes) assert b.get_edge_count() == 5 assert b.get_node_count() == 5 + + +def test_networkx_can_ingest_edgelist_dataframe(): + backend = NetworkXBackend(directed=True) + edgelist = pd.DataFrame( + { + "source": ["A", "B"], + "target": ["B", "C"], + "weight": [1, 2], + } + ) + + result = backend.ingest_from_edgelist_dataframe(edgelist, "source", "target") + + assert result["node_count"] == 3 + assert result["edge_count"] == 2 + assert backend.get_node_count() == 3 + assert backend.get_edge_count() == 2 + assert backend.get_edge_by_id("A", "B")["weight"] == 1 + assert backend.get_edge_by_id("B", "C")["weight"] == 2 + + +def test_sql_can_ingest_edgelist_dataframe(): + if not _CAN_IMPORT_SQL: + pytest.skip("sqlalchemy is not installed.") + + backend = SQLBackend(directed=True, db_url="sqlite:///:memory:") + edgelist = pd.DataFrame( + { + "source": ["A", "B"], + "target": ["B", "C"], + "weight": [1, 2], + } + ) + + result = backend.ingest_from_edgelist_dataframe(edgelist, "source", "target") + + assert result["node_count"] == 3 + assert result["edge_count"] == 2 + assert backend.get_node_count() == 3 + assert backend.get_edge_count() == 2 + assert backend.get_edge_by_id("A", "B")["weight"] == 1 + assert backend.get_edge_by_id("B", "C")["weight"] == 2 diff --git a/grand/dialects/__init__.py b/grand/dialects/__init__.py index d5038d7..8883b8a 100644 --- a/grand/dialects/__init__.py +++ b/grand/dialects/__init__.py @@ -177,50 +177,34 @@ def out_degree(self, nbunch=None): return self.parent.backend.out_degrees(nbunch) def degree(self, nbunch=None): - if self.parent.backend.is_directed(): - # For directed graphs, degree = in_degree + out_degree - if nbunch is None: - # Return DegreeView-like object for all nodes - from networkx.classes.reportviews import DegreeView - - combined_degrees = {} - for node in self.parent.backend.all_nodes_as_iterable(): - in_deg = self.parent.backend.in_degree(node) - out_deg = self.parent.backend.out_degree(node) - combined_degrees[node] = in_deg + out_deg - return DegreeView(combined_degrees) - elif hasattr(nbunch, "__iter__") and not isinstance(nbunch, str): - # nbunch is a list/iterable of nodes - from networkx.classes.reportviews import DegreeView - - result = {} - for node in nbunch: - in_deg = self.parent.backend.in_degree(node) - out_deg = self.parent.backend.out_degree(node) - result[node] = in_deg + out_deg - return DegreeView(result) - else: - # nbunch is a single node + if nbunch is None: + nodes = list(self.parent.backend.all_nodes_as_iterable()) + elif hasattr(nbunch, "__iter__") and not isinstance(nbunch, str): + nodes = list(nbunch) + else: + if self.parent.backend.is_directed(): in_deg = self.parent.backend.in_degree(nbunch) out_deg = self.parent.backend.out_degree(nbunch) return in_deg + out_deg - else: - # For undirected graphs, use the backend's degree method directly - if nbunch is None: - # Return DegreeView for all nodes - from networkx.classes.reportviews import DegreeView - - degrees_dict = self.parent.backend.degrees(nbunch) - return DegreeView(degrees_dict) - elif hasattr(nbunch, "__iter__") and not isinstance(nbunch, str): - # nbunch is a list/iterable of nodes - from networkx.classes.reportviews import DegreeView - - degrees_dict = self.parent.backend.degrees(nbunch) - return DegreeView(degrees_dict) - else: - # nbunch is a single node - return self.parent.backend.degree(nbunch) + return self.parent.backend.degree(nbunch) + + if not nodes: + return {} + + if self.parent.backend.is_directed(): + in_degrees = self.parent.backend.in_degrees(nodes) + out_degrees = self.parent.backend.out_degrees(nodes) + + def _lookup_degree(degrees: dict, node: Hashable) -> int: + return degrees.get(node, degrees.get(str(node), 0)) + + return { + node: _lookup_degree(in_degrees, node) + + _lookup_degree(out_degrees, node) + for node in nodes + } + + return self.parent.backend.degrees(nodes) def is_directed(self): return self.parent.backend.is_directed() diff --git a/grand/dialects/test_dialect.py b/grand/dialects/test_dialect.py index cafcffe..ca9e164 100644 --- a/grand/dialects/test_dialect.py +++ b/grand/dialects/test_dialect.py @@ -121,6 +121,18 @@ def test_degree_directed(self): assert G1.nx.degree(0) == G2.degree(0), f"{G1.nx.degree(0)} != {G2.degree(0)}" assert G1.nx.degree(1) == G2.degree(1), f"{G1.nx.degree(1)} != {G2.degree(1)}" + def test_degree_directed_iterable(self): + G1 = Graph(backend=NetworkXBackend(directed=True)) + G2 = nx.DiGraph() + + G1.nx.add_edge(0, 1) + G1.nx.add_edge(0, 2) + G2.add_edge(0, 1) + G2.add_edge(0, 2) + + self.assertEqual(dict(G1.nx.degree()), dict(G2.degree())) + self.assertEqual(dict(G1.nx.degree([0, 2])), dict(G2.degree([0, 2]))) + def test_nx_export(self): gg = Graph() f = io.BytesIO()