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,
    )

    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.
    
tools/distpartitioning/data_shuffle.py