code: def _test_chunk_graph(
      num_chunks,
      data_fmt="numpy",
      edges_fmt="csv",
      vector_rows=False,
      num_chunks_nodes=None,
      num_chunks_edges=None,
      num_chunks_node_data=None,
      num_chunks_edge_data=None,
  ):
      with tempfile.TemporaryDirectory() as root_dir:
          g = create_chunked_dataset(
              root_dir,
              num_chunks,
              data_fmt=data_fmt,
              edges_fmt=edges_fmt,
              vector_rows=vector_rows,
              num_chunks_nodes=num_chunks_nodes,
              num_chunks_edges=num_chunks_edges,
              num_chunks_node_data=num_chunks_node_data,
              num_chunks_edge_data=num_chunks_edge_data,
          )
  
          # check metadata.json
          output_dir = os.path.join(root_dir, "chunked-data")
          json_file = os.path.join(output_dir, "metadata.json")
          assert os.path.isfile(json_file)
          with open(json_file, "rb") as f:
              meta_data = json.load(f)
          assert meta_data["graph_name"] == "mag240m"
          assert len(meta_data["num_nodes_per_chunk"][0]) == num_chunks
  
          # check edge_index
          output_edge_index_dir = os.path.join(output_dir, "edge_index")
          for c_etype in g.canonical_etypes:
              c_etype_str = _etype_tuple_to_str(c_etype)
              if num_chunks_edges is None:
                  n_chunks = num_chunks
              else:
                  n_chunks = num_chunks_edges
              for i in range(n_chunks):
                  fname = os.path.join(
                      output_edge_index_dir, f"{c_etype_str}{i}.txt"
                  )
                  assert os.path.isfile(fname)
                  if edges_fmt == "csv":
                      with open(fname, "r") as f:
                          header = f.readline()
                          num1, num2 = header.rstrip().split(" ")
                          assert isinstance(int(num1), int)
                          assert isinstance(int(num2), int)
                  elif edges_fmt == "parquet":
                      metadata = pq.read_metadata(fname)
                      assert metadata.num_columns == 2
                  else:
                      assert False, f"Invalid edges_fmt: {edges_fmt}"
  
          # check node/edge_data
          suffix = "npy" if data_fmt == "numpy" else "parquet"
          reader_fmt_meta = {"name": data_fmt}
  
          def test_data(
              sub_dir, feat, expected_data, expected_shape, num_chunks
          ):
              data = []
              for i in range(num_chunks):
                  fname = os.path.join(sub_dir, f"{feat}-{i}.{suffix}")
                  assert os.path.isfile(
                      fname
                  ), f"{fname} cannot be found."
                  feat_array = array_readwriter.get_array_parser(
                      **reader_fmt_meta
                  ).read(fname)
                  assert feat_array.shape[0] == expected_shape
                  data.append(feat_array)
              data = np.concatenate(data, 0)
              assert torch.equal(torch.from_numpy(data), expected_data)
  
          output_node_data_dir = os.path.join(output_dir, "node_data")
          for ntype in g.ntypes:
              sub_dir = os.path.join(output_node_data_dir, ntype)
              if isinstance(num_chunks_node_data, int):
                  chunks_data = num_chunks_node_data
              elif isinstance(num_chunks_node_data, dict):
                  chunks_data = num_chunks_node_data.get(
                      ntype, num_chunks
                  )
              else:
                  chunks_data = num_chunks
              for feat, data in g.nodes[ntype].data.items():
                  if isinstance(chunks_data, dict):
                      n_chunks = chunks_data.get(feat, num_chunks)
                  else:
                      n_chunks = chunks_data
                  test_data(
                      sub_dir,
                      feat,
                      data,
                      g.num_nodes(ntype) // n_chunks,
                      n_chunks,
                  )
  
          output_edge_data_dir = os.path.join(output_dir, "edge_data")
          for c_etype in g.canonical_etypes:
              c_etype_str = _etype_tuple_to_str(c_etype)
              sub_dir = os.path.join(output_edge_data_dir, c_etype_str)
              if isinstance(num_chunks_edge_data, int):
                  chunks_data = num_chunks_edge_data
              elif isinstance(num_chunks_edge_data, dict):
                  chunks_data = num_chunks_edge_data.get(
                      c_etype, num_chunks
                  )
              else:
                  chunks_data = num_chunks
              for feat, data in g.edges[c_etype].data.items():
                  if isinstance(chunks_data, dict):
                      n_chunks = chunks_data.get(feat, num_chunks)
                  else:
                      n_chunks = chunks_data
                  test_data(
                      sub_dir,
                      feat,
                      data,
                      g.num_edges(c_etype) // n_chunks,
                      n_chunks,
                  )
  
imports: import json
  import os
  import tempfile
  import dgl
  import numpy
  import pyarrow.parquet
  import pytest
  import torch
  from dgl.data.utils import load_graphs, load_tensors
  from dgl.distributed.partition import _etype_tuple_to_str, _get_inner_edge_mask, _get_inner_node_mask, load_partition, RESERVED_FIELD_DTYPE
  from distpartitioning import array_readwriter
  from distpartitioning.utils import generate_read_list
  from pytest_utils import create_chunked_dataset
  from tools.verification_utils import verify_graph_feats, verify_partition_data_types, verify_partition_formats
label_desc: Use image node_id 1 for calling a global function with example usage: _test_chunk_graph(num_chunks, data_fmt, edges_fmt, vector_rows, num_chunks_nodes, num_chunks_edges, num_chunks_node_data, num_chunks_edge_data) without return types