code: def read_dataset(
      rank, world_size, id_lookup, params, schema_map, ntype_counts
  ):
      """
      This function gets the dataset and performs post-processing on the data which is read from files.
      Additional information(columns) are added to nodes metadata like owner_process, global_nid which
      are later used in processing this information. For edge data, which is now a dictionary, we add new columns
      like global_edge_id and owner_process. Augmenting these data structure helps in processing these data structures
      when data shuffling is performed.
  
      Parameters:
      -----------
      rank : int
          rank of the current process
      world_size : int
          total no. of processes instantiated
      id_lookup : instance of class DistLookupService
         Distributed lookup service used to map global-nids to respective partition-ids and
         shuffle-global-nids
      params : argparser object
          argument parser object to access command line arguments
      schema_map : dictionary
          dictionary created by reading the input graph metadata json file
  
      Returns :
      ---------
      dictionary
          in which keys are node-type names and values are are tuples representing the range of ids
          for nodes to be read by the current process
      dictionary
          node features which is a dictionary where keys are feature names and values are feature
          data as multi-dimensional tensors
      dictionary
          in which keys are node-type names and values are triplets. Each triplet has node-feature name
          and the starting and ending type ids of the node-feature data read from the corresponding
          node feature data file read by current process. Each node type may have several features and
          hence each key may have several triplets.
      dictionary
          edge data information is read from edges.txt and additional columns are added such as
          owner process for each edge.
      dictionary
          edge features which is also a dictionary, similar to node features dictionary
      dictionary
          a dictionary in which keys are edge-type names and values are tuples indicating the range of ids
          for edges read by the current process.
      dictionary
          a dictionary in which keys are edge-type names and values are triplets,
          (edge-feature-name, start_type_id, end_type_id). These type_ids are indices in the edge-features
          read by the current process. Note that each edge-type may have several edge-features.
      """
      edge_features = {}
      (
          node_features,
          node_feat_tids,
          edge_data,
          edge_typecounts,
          edge_tids,
          edge_features,
          edge_feat_tids,
      ) = get_dataset(
          params.input_dir,
          params.graph_name,
          rank,
          world_size,
          params.num_parts,
          schema_map,
          ntype_counts,
      )
      # Synchronize so that everybody completes reading dataset from disk
      dist.barrier()
      logging.info(
          f"[Rank: {rank}] Done reading dataset {params.input_dir}"
      )
  
      edge_data = augment_edge_data(
          edge_data,
          id_lookup,
          edge_tids,
          rank,
          world_size,
          params.num_parts,
      )
      dist.barrier()  # SYNCH
      logging.debug(
          f"[Rank: {rank}] Done augmenting edge_data: {len(edge_data)}, {edge_data[constants.GLOBAL_SRC_ID].shape}"
      )
  
      return (
          node_features,
          node_feat_tids,
          edge_data,
          edge_typecounts,
          edge_features,
          edge_feat_tids,
      )
  
imports: import gc
  import logging
  import math
  import os
  import sys
  from datetime import timedelta
  from timeit import default_timer
  import constants
  import dgl
  import numpy
  import torch
  import torch.distributed
  import torch.multiprocessing
  from convert_partition import create_dgl_object, create_metadata_json
  from dataset_utils import get_dataset
  from dist_lookup import DistLookupService
  from globalids import assign_shuffle_global_nids_edges, assign_shuffle_global_nids_nodes, lookup_shuffle_global_nids_edges
  from gloo_wrapper import allgather_sizes, alltoallv_cpu, gather_metadata_json
  from utils import augment_edge_data, DATA_TYPE_ID, get_edge_types, get_etype_featnames, get_gid_offsets, get_gnid_range_map, get_idranges, get_node_types, get_ntype_counts_map, get_ntype_featnames, map_partid_rank, memory_snapshot, read_json, read_ntype_partition_files, REV_DATA_TYPE_ID, write_dgl_objects, write_metadata_json
label_desc: Use image node_id 6 for calling a global function with example usage: read_dataset(rank, world_size, id_lookup, params, schema_map, ntype_counts) and returns: node_features,  node_feat_tids,  edge_data,  edge_typecounts,  edge_features,  edge_feat_tids