(
g,
name,
ndata_paths,
edata_paths,
num_chunks,
data_fmt,
edges_format,
vector_rows=False,
**kwargs,
)
| 87 | |
| 88 | |
| 89 | def _chunk_graph( |
| 90 | g, |
| 91 | name, |
| 92 | ndata_paths, |
| 93 | edata_paths, |
| 94 | num_chunks, |
| 95 | data_fmt, |
| 96 | edges_format, |
| 97 | vector_rows=False, |
| 98 | **kwargs, |
| 99 | ): |
| 100 | # First deal with ndata and edata that are homogeneous |
| 101 | # (i.e. not a dict-of-dict) |
| 102 | if len(g.ntypes) == 1 and not isinstance( |
| 103 | next(iter(ndata_paths.values())), dict |
| 104 | ): |
| 105 | ndata_paths = {g.ntypes[0]: ndata_paths} |
| 106 | if len(g.etypes) == 1 and not isinstance( |
| 107 | next(iter(edata_paths.values())), dict |
| 108 | ): |
| 109 | edata_paths = {g.etypes[0]: ndata_paths} |
| 110 | # Then convert all edge types to canonical edge types |
| 111 | etypestrs = {etype: ":".join(etype) for etype in g.canonical_etypes} |
| 112 | edata_paths = { |
| 113 | ":".join(g.to_canonical_etype(k)): v for k, v in edata_paths.items() |
| 114 | } |
| 115 | |
| 116 | metadata = {} |
| 117 | |
| 118 | metadata["graph_name"] = name |
| 119 | metadata["node_type"] = g.ntypes |
| 120 | |
| 121 | # add node_type_counts |
| 122 | metadata["num_nodes_per_type"] = [g.num_nodes(ntype) for ntype in g.ntypes] |
| 123 | |
| 124 | # Initialize num_chunks for each node/edge. |
| 125 | num_chunks_details = _initialize_num_chunks(g, num_chunks, kwargs=kwargs) |
| 126 | |
| 127 | # Compute the number of nodes per chunk per node type |
| 128 | metadata["num_nodes_per_chunk"] = num_nodes_per_chunk = [] |
| 129 | num_chunks_nodes = num_chunks_details["num_chunks_nodes"] |
| 130 | for ntype in g.ntypes: |
| 131 | num_nodes = g.num_nodes(ntype) |
| 132 | num_nodes_list = [] |
| 133 | n_chunks = num_chunks_nodes[ntype] |
| 134 | for i in range(n_chunks): |
| 135 | n = num_nodes // n_chunks + (i < num_nodes % n_chunks) |
| 136 | num_nodes_list.append(n) |
| 137 | num_nodes_per_chunk.append(num_nodes_list) |
| 138 | |
| 139 | metadata["edge_type"] = [etypestrs[etype] for etype in g.canonical_etypes] |
| 140 | metadata["num_edges_per_type"] = [ |
| 141 | g.num_edges(etype) for etype in g.canonical_etypes |
| 142 | ] |
| 143 | |
| 144 | # Compute the number of edges per chunk per edge type |
| 145 | metadata["num_edges_per_chunk"] = num_edges_per_chunk = [] |
| 146 | num_chunks_edges = num_chunks_details["num_chunks_edges"] |
no test coverage detected