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/_sqlbackend.py b/grand/backends/_sqlbackend.py index 6e68e7a..271c63f 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( 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()