nx-cugraph: add G.__networkx_cache__ to enable graph conversion caching (#4567)

Adding this now, because this was an oversight.

By having a `G.__networkx_cache__` MutableMapping, we give NetworkX the ability to cache graph conversions. This will be most useful in dev and future versions of NetworkX. As such, it would be nice (but not strictly essential) to get this in 24.08.

Authors:
  - Erik Welch (https://github.com/eriknw)

Approvers:
  - Rick Ratzel (https://github.com/rlratzel)

URL: #4567

Author: eriknw

python/nx-cugraph/nx_cugraph/classes/graph.py

@@ -50,6 +50,10 @@ class Graph:
     __networkx_backend__: ClassVar[str] = "cugraph"  # nx >=3.2
     __networkx_plugin__: ClassVar[str] = "cugraph"  # nx <3.2
 
+    # Allow networkx dispatch machinery to cache conversions.
+    # This means we should clear the cache if we ever mutate the object!
+    __networkx_cache__: dict | None
+
     # networkx properties
     graph: dict
     graph_attr_dict_factory: ClassVar[type] = dict
@@ -108,6 +112,7 @@ def from_coo(
         **attr,
     ) -> Graph:
         new_graph = object.__new__(cls)
+        new_graph.__networkx_cache__ = {}
         new_graph.src_indices = src_indices
         new_graph.dst_indices = dst_indices
         new_graph.edge_values = {} if edge_values is None else dict(edge_values)
@@ -420,13 +425,17 @@ def clear(self) -> None:
         self._node_ids = None
         self.key_to_id = None
         self._id_to_key = None
+        if cache := self.__networkx_cache__:
+            cache.clear()
 
     @networkx_api
     def clear_edges(self) -> None:
         self.edge_values.clear()
         self.edge_masks.clear()
         self.src_indices = cp.empty(0, self.src_indices.dtype)
         self.dst_indices = cp.empty(0, self.dst_indices.dtype)
+        if cache := self.__networkx_cache__:
+            cache.clear()
 
     @networkx_api
     def copy(self, as_view: bool = False) -> Graph:
@@ -553,6 +562,12 @@ def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False):
         node_masks = self.node_masks
         key_to_id = self.key_to_id
         id_to_key = None if key_to_id is None else self._id_to_key
+        if self.__networkx_cache__ is None:
+            __networkx_cache__ = None
+        elif not reverse and cls is self.__class__:
+            __networkx_cache__ = self.__networkx_cache__
+        else:
+            __networkx_cache__ = {}
         if not as_view:
             src_indices = src_indices.copy()
             dst_indices = dst_indices.copy()
@@ -564,6 +579,8 @@ def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False):
                 key_to_id = key_to_id.copy()
                 if id_to_key is not None:
                     id_to_key = id_to_key.copy()
+            if __networkx_cache__ is not None:
+                __networkx_cache__ = __networkx_cache__.copy()
         if reverse:
             src_indices, dst_indices = dst_indices, src_indices
         rv = cls.from_coo(
@@ -581,6 +598,7 @@ def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False):
             rv.graph = self.graph
         else:
             rv.graph.update(deepcopy(self.graph))
+        rv.__networkx_cache__ = __networkx_cache__
         return rv
 
     def _get_plc_graph(
@@ -719,18 +737,26 @@ def _become(self, other: Graph):
         edge_masks = self.edge_masks
         node_values = self.node_values
         node_masks = self.node_masks
+        __networkx_cache__ = self.__networkx_cache__
         graph = self.graph
         edge_values.update(other.edge_values)
         edge_masks.update(other.edge_masks)
         node_values.update(other.node_values)
         node_masks.update(other.node_masks)
         graph.update(other.graph)
+        if other.__networkx_cache__ is None:
+            __networkx_cache__ = None
+        else:
+            if __networkx_cache__ is None:
+                __networkx_cache__ = {}
+            __networkx_cache__.update(other.__networkx_cache__)
         self.__dict__.update(other.__dict__)
         self.edge_values = edge_values
         self.edge_masks = edge_masks
         self.node_values = node_values
         self.node_masks = node_masks
         self.graph = graph
+        self.__networkx_cache__ = __networkx_cache__
         return self
 
     def _degrees_array(self, *, ignore_selfloops=False):

python/nx-cugraph/nx_cugraph/classes/multigraph.py

@@ -415,6 +415,12 @@ def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False):
         key_to_id = self.key_to_id
         id_to_key = None if key_to_id is None else self._id_to_key
         edge_keys = self.edge_keys
+        if self.__networkx_cache__ is None:
+            __networkx_cache__ = None
+        elif not reverse and cls is self.__class__:
+            __networkx_cache__ = self.__networkx_cache__
+        else:
+            __networkx_cache__ = {}
         if not as_view:
             src_indices = src_indices.copy()
             dst_indices = dst_indices.copy()
@@ -429,6 +435,8 @@ def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False):
                     id_to_key = id_to_key.copy()
             if edge_keys is not None:
                 edge_keys = edge_keys.copy()
+            if __networkx_cache__ is not None:
+                __networkx_cache__ = __networkx_cache__.copy()
         if reverse:
             src_indices, dst_indices = dst_indices, src_indices
         rv = cls.from_coo(
@@ -448,6 +456,7 @@ def _copy(self, as_view: bool, cls: type[Graph], reverse: bool = False):
             rv.graph = self.graph
         else:
             rv.graph.update(deepcopy(self.graph))
+        rv.__networkx_cache__ = __networkx_cache__
         return rv
 
     def _sort_edge_indices(self, primary="src"):