code: def _chunk_graph(
      g,
      name,
      ndata_paths,
      edata_paths,
      num_chunks,
      data_fmt,
      edges_format,
      vector_rows=False,
      **kwargs,
  ):
      # First deal with ndata and edata that are homogeneous
      # (i.e. not a dict-of-dict)
      if len(g.ntypes) == 1 and not isinstance(
          next(iter(ndata_paths.values())), dict
      ):
          ndata_paths = {g.ntypes[0]: ndata_paths}
      if len(g.etypes) == 1 and not isinstance(
          next(iter(edata_paths.values())), dict
      ):
          edata_paths = {g.etypes[0]: ndata_paths}
      # Then convert all edge types to canonical edge types
      etypestrs = {
          etype: ":".join(etype) for etype in g.canonical_etypes
      }
      edata_paths = {
          ":".join(g.to_canonical_etype(k)): v
          for k, v in edata_paths.items()
      }
  
      metadata = {}
  
      metadata["graph_name"] = name
      metadata["node_type"] = g.ntypes
  
      # add node_type_counts
      metadata["num_nodes_per_type"] = [
          g.num_nodes(ntype) for ntype in g.ntypes
      ]
  
      # Initialize num_chunks for each node/edge.
      num_chunks_details = _initialize_num_chunks(
          g, num_chunks, kwargs=kwargs
      )
  
      # Compute the number of nodes per chunk per node type
      metadata["num_nodes_per_chunk"] = num_nodes_per_chunk = []
      num_chunks_nodes = num_chunks_details["num_chunks_nodes"]
      for ntype in g.ntypes:
          num_nodes = g.num_nodes(ntype)
          num_nodes_list = []
          n_chunks = num_chunks_nodes[ntype]
          for i in range(n_chunks):
              n = num_nodes // n_chunks + (i < num_nodes % n_chunks)
              num_nodes_list.append(n)
          num_nodes_per_chunk.append(num_nodes_list)
  
      metadata["edge_type"] = [
          etypestrs[etype] for etype in g.canonical_etypes
      ]
      metadata["num_edges_per_type"] = [
          g.num_edges(etype) for etype in g.canonical_etypes
      ]
  
      # Compute the number of edges per chunk per edge type
      metadata["num_edges_per_chunk"] = num_edges_per_chunk = []
      num_chunks_edges = num_chunks_details["num_chunks_edges"]
      for etype in g.canonical_etypes:
          num_edges = g.num_edges(etype)
          num_edges_list = []
          n_chunks = num_chunks_edges[etype]
          for i in range(n_chunks):
              n = num_edges // n_chunks + (i < num_edges % n_chunks)
              num_edges_list.append(n)
          num_edges_per_chunk.append(num_edges_list)
      num_edges_per_chunk_dict = {
          k: v for k, v in zip(g.canonical_etypes, num_edges_per_chunk)
      }
  
      idxes_etypestr = {
          idx: (etype, etypestrs[etype])
          for idx, etype in enumerate(g.canonical_etypes)
      }
      idxes = np.arange(len(idxes_etypestr))
  
      # Split edge index
      metadata["edges"] = {}
      with setdir("edge_index"):
          np.random.shuffle(idxes)
          for idx in idxes:
              etype = idxes_etypestr[idx][0]
              etypestr = idxes_etypestr[idx][1]
              logging.info("Chunking edge index for %s" % etypestr)
              edges_meta = {}
              if edges_format == "csv":
                  fmt_meta = {"name": edges_format, "delimiter": " "}
              elif edges_format == "parquet":
                  fmt_meta = {"name": edges_format}
              else:
                  raise RuntimeError(
                      f"Invalid edges_fmt: {edges_format}"
                  )
              edges_meta["format"] = fmt_meta
  
              srcdst = torch.stack(g.edges(etype=etype), 1)
              edges_meta["data"] = _chunk_numpy_array(
                  srcdst.numpy(),
                  fmt_meta,
                  num_edges_per_chunk_dict[etype],
                  etypestr + "%d.txt",
              )
              metadata["edges"][etypestr] = edges_meta
  
      # Chunk node data
      reader_fmt_meta, writer_fmt_meta = {"name": "numpy"}, {
          "name": data_fmt
      }
      file_suffix = "npy" if data_fmt == "numpy" else "parquet"
      metadata["node_data"] = {}
      num_chunks_node_data = num_chunks_details["num_chunks_node_data"]
      with setdir("node_data"):
          for ntype, ndata_per_type in ndata_paths.items():
              ndata_meta = {}
              with setdir(ntype):
                  for key, path in ndata_per_type.items():
                      logging.info(
                          "Chunking node data for type %s key %s"
                          % (ntype, key)
                      )
                      chunk_sizes = []
                      num_nodes = g.num_nodes(ntype)
                      n_chunks = num_chunks_node_data[ntype]
                      if isinstance(n_chunks, dict):
                          n_chunks = n_chunks.get(key, num_chunks)
                      assert isinstance(n_chunks, int), (
                          f"num_chunks for {ntype}/{key} should be int while "
                          f"{type(n_chunks)} is got."
                      )
                      for i in range(n_chunks):
                          n = num_nodes // n_chunks + (
                              i < num_nodes % n_chunks
                          )
                          chunk_sizes.append(n)
                      ndata_key_meta = {}
                      arr = array_readwriter.get_array_parser(
                          **reader_fmt_meta
                      ).read(path)
                      ndata_key_meta["format"] = writer_fmt_meta
                      ndata_key_meta["data"] = _chunk_numpy_array(
                          arr,
                          writer_fmt_meta,
                          chunk_sizes,
                          key + "-%d." + file_suffix,
                          vector_rows=vector_rows,
                      )
                      ndata_meta[key] = ndata_key_meta
  
              metadata["node_data"][ntype] = ndata_meta
  
      # Chunk edge data
      metadata["edge_data"] = {}
      num_chunks_edge_data = num_chunks_details["num_chunks_edge_data"]
      with setdir("edge_data"):
          for etypestr, edata_per_type in edata_paths.items():
              edata_meta = {}
              etype = tuple(etypestr.split(":"))
              with setdir(etypestr):
                  for key, path in edata_per_type.items():
                      logging.info(
                          "Chunking edge data for type %s key %s"
                          % (etypestr, key)
                      )
                      chunk_sizes = []
                      num_edges = g.num_edges(etype)
                      n_chunks = num_chunks_edge_data[etype]
                      if isinstance(n_chunks, dict):
                          n_chunks = n_chunks.get(key, num_chunks)
                      assert isinstance(n_chunks, int), (
                          f"num_chunks for {etype}/{key} should be int while "
                          f"{type(n_chunks)} is got."
                      )
                      for i in range(n_chunks):
                          n = num_edges // n_chunks + (
                              i < num_edges % n_chunks
                          )
                          chunk_sizes.append(n)
                      edata_key_meta = {}
                      arr = array_readwriter.get_array_parser(
                          **reader_fmt_meta
                      ).read(path)
                      edata_key_meta["format"] = writer_fmt_meta
                      edata_key_meta["data"] = _chunk_numpy_array(
                          arr,
                          writer_fmt_meta,
                          chunk_sizes,
                          key + "-%d." + file_suffix,
                          vector_rows=vector_rows,
                      )
                      edata_meta[key] = edata_key_meta
  
              metadata["edge_data"][etypestr] = edata_meta
  
      metadata_path = "metadata.json"
      with open(metadata_path, "w") as f:
          json.dump(metadata, f, sort_keys=True, indent=4)
      logging.info(
          "Saved metadata in %s" % os.path.abspath(metadata_path)
      )
  
file_path: dgl/tests/tools/pytest_utils.py
imports: import json
  import logging
  import os
  import dgl
  import numpy
  import torch
  from distpartitioning import array_readwriter
  from distpartitioning.array_readwriter.parquet import ParquetArrayParser
  from files import setdir
label_desc: Use image node_id 3 for calling a global function with example usage: _chunk_graph(g, name, ndata_paths, edata_paths, num_chunks, data_fmt, edges_format, vector_rows) without return types