chore: import upstream snapshot with attribution

python/dgl/data/heterograph_serialize.py

@@ -0,0 +1,79 @@
+"""For HeteroGraph Serialization"""
+from __future__ import absolute_import
+
+from .. import backend as F
+from .._ffi.function import _init_api
+from .._ffi.object import ObjectBase, register_object
+from ..container import convert_to_strmap
+from ..frame import Frame
+from ..heterograph import DGLGraph
+
+_init_api("dgl.data.heterograph_serialize")
+
+
+def tensor_dict_to_ndarray_dict(tensor_dict):
+    """Convert dict[str, tensor] to StrMap[NDArray]"""
+    ndarray_dict = {}
+    for key, value in tensor_dict.items():
+        ndarray_dict[key] = F.zerocopy_to_dgl_ndarray(value)
+    return convert_to_strmap(ndarray_dict)
+
+
+def save_heterographs(filename, g_list, labels, formats):
+    """Save heterographs into file"""
+    if labels is None:
+        labels = {}
+    if isinstance(g_list, DGLGraph):
+        g_list = [g_list]
+    assert all(
+        [type(g) == DGLGraph for g in g_list]
+    ), "Invalid DGLGraph in g_list argument"
+    gdata_list = [HeteroGraphData.create(g) for g in g_list]
+    if formats is None:
+        formats = []
+    elif isinstance(formats, str):
+        formats = [formats]
+    _CAPI_SaveHeteroGraphData(
+        filename, gdata_list, tensor_dict_to_ndarray_dict(labels), formats
+    )
+
+
+@register_object("heterograph_serialize.HeteroGraphData")
+class HeteroGraphData(ObjectBase):
+    """Object to hold the data to be stored for DGLGraph"""
+
+    @staticmethod
+    def create(g):
+        edata_list = []
+        ndata_list = []
+        for etype in g.canonical_etypes:
+            edata_list.append(tensor_dict_to_ndarray_dict(g.edges[etype].data))
+        for ntype in g.ntypes:
+            ndata_list.append(tensor_dict_to_ndarray_dict(g.nodes[ntype].data))
+        return _CAPI_MakeHeteroGraphData(
+            g._graph, ndata_list, edata_list, g.ntypes, g.etypes
+        )
+
+    def get_graph(self):
+        ntensor_list = list(_CAPI_GetNDataFromHeteroGraphData(self))
+        etensor_list = list(_CAPI_GetEDataFromHeteroGraphData(self))
+        ntype_names = list(_CAPI_GetNtypesFromHeteroGraphData(self))
+        etype_names = list(_CAPI_GetEtypesFromHeteroGraphData(self))
+        gidx = _CAPI_GetGindexFromHeteroGraphData(self)
+        nframes = []
+        eframes = []
+        for ntid, ntensor in enumerate(ntensor_list):
+            ndict = {
+                ntensor[i]: F.zerocopy_from_dgl_ndarray(ntensor[i + 1])
+                for i in range(0, len(ntensor), 2)
+            }
+            nframes.append(Frame(ndict, num_rows=gidx.num_nodes(ntid)))
+
+        for etid, etensor in enumerate(etensor_list):
+            edict = {
+                etensor[i]: F.zerocopy_from_dgl_ndarray(etensor[i + 1])
+                for i in range(0, len(etensor), 2)
+            }
+            eframes.append(Frame(edict, num_rows=gidx.num_edges(etid)))
+
+        return DGLGraph(gidx, ntype_names, etype_names, nframes, eframes)