code: def gen_node_weights_files(schema_map, params):
      """Function to create node weight files for ParMETIS along with the edge files.
  
      This function generates node-data files, which will be read by the ParMETIS
      executable for partitioning purposes. Each line in these files will be of the
      following format:
          <node_type_id> <node_weight_list> <type_wise_node_id>
      node_type_id -  is id assigned to the node-type to which a given particular
          node belongs to
      weight_list - this is a one-hot vector in which the number in the location of
          the current nodes' node-type will be set to `1` and other will be `0`
      type_node_id - this is the id assigned to the node (in the context of the current
          nodes` node-type). Meaning this id is unique across all the nodes which belong to
          the current nodes` node-type.
  
      Parameters:
      -----------
      schema_map : json dictionary
          Dictionary created by reading the metadata.json file for the input dataset.
      output : string
          Location of storing the node-weights and edge files for ParMETIS.
  
      Returns:
      --------
      list :
          List of filenames for nodes of the input graph.
      list :
          List o ffilenames for edges of the input graph.
      """
      rank = get_proc_info()
      ntypes_ntypeid_map, ntypes, ntid_ntype_map = get_node_types(
          schema_map
      )
      type_nid_dict, ntype_gnid_offset = get_idranges(
          schema_map[constants.STR_NODE_TYPE],
          dict(
              zip(
                  schema_map[constants.STR_NODE_TYPE],
                  schema_map[constants.STR_NUM_NODES_PER_TYPE],
              )
          ),
      )
  
      node_files = []
      outdir = Path(params.output_dir)
      os.makedirs(outdir, exist_ok=True)
  
      for ntype_id, ntype_name in ntid_ntype_map.items():
          # This ntype does not have any train/test/val masks...
          # Each rank will generate equal no. of rows for this node type.
          total_count = schema_map[constants.STR_NUM_NODES_PER_TYPE][
              ntype_id
          ]
          per_rank_range = np.ones(
              (params.num_parts,), dtype=np.int64
          ) * (total_count // params.num_parts)
          for i in range(total_count % params.num_parts):
              per_rank_range[i] += 1
  
          tid_start = np.cumsum([0] + list(per_rank_range[:-1]))
          tid_end = np.cumsum(list(per_rank_range))
          local_tid_start = tid_start[rank]
          local_tid_end = tid_end[rank]
          sz = local_tid_end - local_tid_start
  
          cols = []
          col_names = []
  
          # ntype-id
          cols.append(
              pyarrow.array(
                  np.ones(sz, dtype=np.int64) * np.int64(ntype_id)
              )
          )
          col_names.append("ntype")
  
          # one-hot vector for ntype-id here.
          for i in range(len(ntypes)):
              if i == ntype_id:
                  cols.append(
                      pyarrow.array(np.ones(sz, dtype=np.int64))
                  )
              else:
                  cols.append(
                      pyarrow.array(np.zeros(sz, dtype=np.int64))
                  )
              col_names.append("w{}".format(i))
  
          # `type_nid` should be the very last column in the node weights files.
          cols.append(
              pyarrow.array(
                  np.arange(
                      local_tid_start, local_tid_end, dtype=np.int64
                  )
              )
          )
          col_names.append("type_nid")
  
          out_file = os.path.join(
              outdir, "node_weights_{}_{}.txt".format(ntype_name, rank)
          )
          options = csv.WriteOptions(
              include_header=False, delimiter=" "
          )
          options.delimiter = " "
  
          csv.write_csv(
              pyarrow.Table.from_arrays(cols, names=col_names),
              out_file,
              options,
          )
          node_files.append(
              (
                  ntype_gnid_offset[ntype_name][0, 0] + local_tid_start,
                  ntype_gnid_offset[ntype_name][0, 0] + local_tid_end,
                  out_file,
              )
          )
  
      return node_files
  
file_path: dgl/tools/distpartitioning/parmetis_preprocess.py
imports: import argparse
  import logging
  import os
  import platform
  from pathlib import Path
  import array_readwriter
  import constants
  import numpy
  import pyarrow
  import pyarrow.csv
  from utils import generate_read_list, generate_roundrobin_read_list, get_idranges, get_node_types, read_json
label_desc: Use image node_id 4 for calling a global function with example usage: gen_node_weights_files(schema_map, params) and returns: node_files