code: def create_chunked_dataset(
      root_dir,
      num_chunks,
      data_fmt="numpy",
      edges_fmt="csv",
      vector_rows=False,
      **kwargs,
  ):
      """
      This function creates a sample dataset, based on MAG240 dataset.
  
      Parameters:
      -----------
      root_dir : string
          directory in which all the files for the chunked dataset will be stored.
      """
      # Step0: prepare chunked graph data format.
      # A synthetic mini MAG240.
      num_institutions = 1200
      num_authors = 1200
      num_papers = 1200
  
      def rand_edges(num_src, num_dst, num_edges):
          eids = np.random.choice(
              num_src * num_dst, num_edges, replace=False
          )
          src = torch.from_numpy(eids // num_dst)
          dst = torch.from_numpy(eids % num_dst)
  
          return src, dst
  
      num_cite_edges = 24 * 1000
      num_write_edges = 12 * 1000
      num_affiliate_edges = 2400
  
      # Structure.
      data_dict = {
          ("paper", "cites", "paper"): rand_edges(
              num_papers, num_papers, num_cite_edges
          ),
          ("author", "writes", "paper"): rand_edges(
              num_authors, num_papers, num_write_edges
          ),
          ("author", "affiliated_with", "institution"): rand_edges(
              num_authors, num_institutions, num_affiliate_edges
          ),
          ("institution", "writes", "paper"): rand_edges(
              num_institutions, num_papers, num_write_edges
          ),
      }
      src, dst = data_dict[("author", "writes", "paper")]
      data_dict[("paper", "rev_writes", "author")] = (dst, src)
      g = dgl.heterograph(data_dict)
  
      # paper feat, label, year
      num_paper_feats = 3
      paper_feat = np.random.randn(num_papers, num_paper_feats)
      num_classes = 4
      paper_label = np.random.choice(num_classes, num_papers)
      paper_year = np.random.choice(2022, num_papers)
      paper_orig_ids = np.arange(0, num_papers)
      writes_orig_ids = np.arange(0, num_write_edges)
  
      # masks.
      paper_train_mask = np.random.choice([True, False], num_papers)
      paper_test_mask = np.random.choice([True, False], num_papers)
      paper_val_mask = np.random.choice([True, False], num_papers)
  
      author_train_mask = np.random.choice([True, False], num_authors)
      author_test_mask = np.random.choice([True, False], num_authors)
      author_val_mask = np.random.choice([True, False], num_authors)
  
      inst_train_mask = np.random.choice(
          [True, False], num_institutions
      )
      inst_test_mask = np.random.choice([True, False], num_institutions)
      inst_val_mask = np.random.choice([True, False], num_institutions)
  
      write_train_mask = np.random.choice(
          [True, False], num_write_edges
      )
      write_test_mask = np.random.choice([True, False], num_write_edges)
      write_val_mask = np.random.choice([True, False], num_write_edges)
  
      # Edge features.
      cite_count = np.random.choice(10, num_cite_edges)
      write_year = np.random.choice(2022, num_write_edges)
      write2_year = np.random.choice(2022, num_write_edges)
  
      # Save features.
      input_dir = os.path.join(root_dir, "data_test")
      os.makedirs(input_dir)
      for sub_d in ["paper", "cites", "writes", "writes2"]:
          os.makedirs(os.path.join(input_dir, sub_d))
  
      paper_feat_path = os.path.join(input_dir, "paper/feat.npy")
      with open(paper_feat_path, "wb") as f:
          np.save(f, paper_feat)
      g.nodes["paper"].data["feat"] = torch.from_numpy(paper_feat)
  
      paper_label_path = os.path.join(input_dir, "paper/label.npy")
      with open(paper_label_path, "wb") as f:
          np.save(f, paper_label)
      g.nodes["paper"].data["label"] = torch.from_numpy(paper_label)
  
      paper_year_path = os.path.join(input_dir, "paper/year.npy")
      with open(paper_year_path, "wb") as f:
          np.save(f, paper_year)
      g.nodes["paper"].data["year"] = torch.from_numpy(paper_year)
  
      paper_orig_ids_path = os.path.join(
          input_dir, "paper/orig_ids.npy"
      )
      with open(paper_orig_ids_path, "wb") as f:
          np.save(f, paper_orig_ids)
      g.nodes["paper"].data["orig_ids"] = torch.from_numpy(
          paper_orig_ids
      )
  
      cite_count_path = os.path.join(input_dir, "cites/count.npy")
      with open(cite_count_path, "wb") as f:
          np.save(f, cite_count)
      g.edges["cites"].data["count"] = torch.from_numpy(cite_count)
  
      write_year_path = os.path.join(input_dir, "writes/year.npy")
      with open(write_year_path, "wb") as f:
          np.save(f, write_year)
      g.edges[("author", "writes", "paper")].data[
          "year"
      ] = torch.from_numpy(write_year)
      g.edges["rev_writes"].data["year"] = torch.from_numpy(write_year)
  
      writes_orig_ids_path = os.path.join(
          input_dir, "writes/orig_ids.npy"
      )
      with open(writes_orig_ids_path, "wb") as f:
          np.save(f, writes_orig_ids)
      g.edges[("author", "writes", "paper")].data[
          "orig_ids"
      ] = torch.from_numpy(writes_orig_ids)
  
      write2_year_path = os.path.join(input_dir, "writes2/year.npy")
      with open(write2_year_path, "wb") as f:
          np.save(f, write2_year)
      g.edges[("institution", "writes", "paper")].data[
          "year"
      ] = torch.from_numpy(write2_year)
  
      etype = ("author", "writes", "paper")
      write_train_mask_path = os.path.join(
          input_dir, "writes/train_mask.npy"
      )
      with open(write_train_mask_path, "wb") as f:
          np.save(f, write_train_mask)
      g.edges[etype].data["train_mask"] = torch.from_numpy(
          write_train_mask
      )
  
      write_test_mask_path = os.path.join(
          input_dir, "writes/test_mask.npy"
      )
      with open(write_test_mask_path, "wb") as f:
          np.save(f, write_test_mask)
      g.edges[etype].data["test_mask"] = torch.from_numpy(
          write_test_mask
      )
  
      write_val_mask_path = os.path.join(
          input_dir, "writes/val_mask.npy"
      )
      with open(write_val_mask_path, "wb") as f:
          np.save(f, write_val_mask)
      g.edges[etype].data["val_mask"] = torch.from_numpy(write_val_mask)
  
      for sub_d in ["author", "institution"]:
          os.makedirs(os.path.join(input_dir, sub_d))
      paper_train_mask_path = os.path.join(
          input_dir, "paper/train_mask.npy"
      )
      with open(paper_train_mask_path, "wb") as f:
          np.save(f, paper_train_mask)
      g.nodes["paper"].data["train_mask"] = torch.from_numpy(
          paper_train_mask
      )
  
      paper_test_mask_path = os.path.join(
          input_dir, "paper/test_mask.npy"
      )
      with open(paper_test_mask_path, "wb") as f:
          np.save(f, paper_test_mask)
      g.nodes["paper"].data["test_mask"] = torch.from_numpy(
          paper_test_mask
      )
  
      paper_val_mask_path = os.path.join(
          input_dir, "paper/val_mask.npy"
      )
      with open(paper_val_mask_path, "wb") as f:
          np.save(f, paper_val_mask)
      g.nodes["paper"].data["val_mask"] = torch.from_numpy(
          paper_val_mask
      )
  
      author_train_mask_path = os.path.join(
          input_dir, "author/train_mask.npy"
      )
      with open(author_train_mask_path, "wb") as f:
          np.save(f, author_train_mask)
      g.nodes["author"].data["train_mask"] = torch.from_numpy(
          author_train_mask
      )
  
      author_test_mask_path = os.path.join(
          input_dir, "author/test_mask.npy"
      )
      with open(author_test_mask_path, "wb") as f:
          np.save(f, author_test_mask)
      g.nodes["author"].data["test_mask"] = torch.from_numpy(
          author_test_mask
      )
  
      author_val_mask_path = os.path.join(
          input_dir, "author/val_mask.npy"
      )
      with open(author_val_mask_path, "wb") as f:
          np.save(f, author_val_mask)
      g.nodes["author"].data["val_mask"] = torch.from_numpy(
          author_val_mask
      )
  
      inst_train_mask_path = os.path.join(
          input_dir, "institution/train_mask.npy"
      )
      with open(inst_train_mask_path, "wb") as f:
          np.save(f, inst_train_mask)
      g.nodes["institution"].data["train_mask"] = torch.from_numpy(
          inst_train_mask
      )
  
      inst_test_mask_path = os.path.join(
          input_dir, "institution/test_mask.npy"
      )
      with open(inst_test_mask_path, "wb") as f:
          np.save(f, inst_test_mask)
      g.nodes["institution"].data["test_mask"] = torch.from_numpy(
          inst_test_mask
      )
  
      inst_val_mask_path = os.path.join(
          input_dir, "institution/val_mask.npy"
      )
      with open(inst_val_mask_path, "wb") as f:
          np.save(f, inst_val_mask)
      g.nodes["institution"].data["val_mask"] = torch.from_numpy(
          inst_val_mask
      )
  
      node_data = {
          "paper": {
              "feat": paper_feat_path,
              "train_mask": paper_train_mask_path,
              "test_mask": paper_test_mask_path,
              "val_mask": paper_val_mask_path,
              "label": paper_label_path,
              "year": paper_year_path,
              "orig_ids": paper_orig_ids_path,
          },
          "author": {
              "train_mask": author_train_mask_path,
              "test_mask": author_test_mask_path,
              "val_mask": author_val_mask_path,
          },
          "institution": {
              "train_mask": inst_train_mask_path,
              "test_mask": inst_test_mask_path,
              "val_mask": inst_val_mask_path,
          },
      }
  
      edge_data = {
          "cites": {"count": cite_count_path},
          ("author", "writes", "paper"): {
              "year": write_year_path,
              "orig_ids": writes_orig_ids_path,
              "train_mask": write_train_mask_path,
              "test_mask": write_test_mask_path,
              "val_mask": write_val_mask_path,
          },
          "rev_writes": {"year": write_year_path},
          ("institution", "writes", "paper"): {
              "year": write2_year_path
          },
      }
  
      output_dir = os.path.join(root_dir, "chunked-data")
      chunk_graph(
          g,
          "mag240m",
          node_data,
          edge_data,
          num_chunks=num_chunks,
          output_path=output_dir,
          data_fmt=data_fmt,
          edges_fmt=edges_fmt,
          vector_rows=vector_rows,
          **kwargs,
      )
      logging.debug("Done with creating chunked graph")
  
      return g
  
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 5 for calling a global function with example usage: create_chunked_dataset(root_dir, num_chunks, data_fmt, edges_fmt, vector_rows) and returns: g, src,  dst