def exchange_graph_data(
    rank,
    world_size,
    num_parts,
    node_features,
    edge_features,
    node_feat_tids,
    edge_feat_tids,
    edge_data,
    id_lookup,
    ntypes_ntypeid_map,
    ntypes_gnid_range_map,
    etypes_geid_range_map,
    ntid_ntype_map,
    schema_map,
):
    """
    Wrapper function which is used to shuffle graph data on all the processes.

    Parameters:
    -----------
    rank : int
        rank of the current process
    world_size : int
        total no. of participating processes.
    num_parts : int
        total no. of graph partitions.
    node_feautres : dicitonary
        dictionry where node_features are stored and this information is read from the appropriate
        node features file which belongs to the current process
    edge_features : dictionary
        dictionary where edge_features are stored. This information is read from the appropriate
        edge feature files whose ownership is assigned to the current process
    node_feat_tids: 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.
    edge_feat_tids : dictionary
        a dictionary in which keys are edge-type names and values are triplets of the format
        <feat-name, start-per-type-idx, end-per-type-idx>. This triplet is used to identify
        the chunk of feature data for which current process is responsible for
    edge_data : dictionary
        dictionary which is used to store edge information as read from appropriate files assigned
        to each process.
    id_lookup : instance of class DistLookupService
       Distributed lookup service used to map global-nids to respective partition-ids and
       shuffle-global-nids
    ntypes_ntypeid_map : dictionary
        mappings between node type names and node type ids
    ntypes_gnid_range_map : dictionary
        mapping between node type names and global_nids which belong to the keys in this dictionary
    etypes_geid_range_map : dictionary
        mapping between edge type names and global_eids which are assigned to the edges of this
        edge_type
    ntid_ntype_map : dictionary
        mapping between node type id and no of nodes which belong to each node_type_id
    schema_map : dictionary
        is the data structure read from the metadata json file for the input graph

    Returns:
    --------
    dictionary :
        the input argument, node_data dictionary, is updated with the node data received from other processes
        in the world. The node data is received by each rank in the process of data shuffling.
    dictionary :
        node features dictionary which has node features for the nodes which are owned by the current
        process
    dictionary :
        list of global_nids for the nodes whose node features are received when node features shuffling was
        performed in the `exchange_features` function call
    dictionary :
        the input argument, edge_data dictionary, is updated with the edge data received from other processes
        in the world. The edge data is received by each rank in the process of data shuffling.
    dictionary :
        edge features dictionary which has edge features. These destination end points of these edges
        are owned by the current process
    dictionary :
        list of global_eids for the edges whose edge features are received when edge features shuffling
        was performed in the `exchange_features` function call
    """
    memory_snapshot("ShuffleNodeFeaturesBegin: ", rank)
    logging.debug(f"[Rank: {rank} - node_feat_tids - {node_feat_tids}")
    rcvd_node_features, rcvd_global_nids = exchange_features(
        rank,
        world_size,
        num_parts,
        node_feat_tids,
        ntypes_gnid_range_map,
        id_lookup,
        node_features,
        constants.STR_NODE_FEATURES,
        None,
    )
    dist.barrier()
    memory_snapshot("ShuffleNodeFeaturesComplete: ", rank)
    logging.debug(f"[Rank: {rank}] Done with node features exchange.")

    rcvd_edge_features, rcvd_global_eids = exchange_features(
        rank,
        world_size,
        num_parts,
        edge_feat_tids,
        etypes_geid_range_map,
        id_lookup,
        edge_features,
        constants.STR_EDGE_FEATURES,
        edge_data,
    )
    dist.barrier()
    logging.debug(f"[Rank: {rank}] Done with edge features exchange.")

    node_data = gen_node_data(
        rank, world_size, num_parts, id_lookup, ntid_ntype_map, schema_map
    )
    dist.barrier()
    memory_snapshot("NodeDataGenerationComplete: ", rank)

    edge_data = exchange_edge_data(
        rank, world_size, num_parts, edge_data, id_lookup
    )
    dist.barrier()
    memory_snapshot("ShuffleEdgeDataComplete: ", rank)
    return (
        node_data,
        rcvd_node_features,
        rcvd_global_nids,
        edge_data,
        rcvd_edge_features,
        rcvd_global_eids,
    )

    Wrapper function which is used to shuffle graph data on all the processes.

    Parameters:
    -----------
    rank : int
        rank of the current process
    world_size : int
        total no. of participating processes.
    num_parts : int
        total no. of graph partitions.
    node_feautres : dicitonary
        dictionry where node_features are stored and this information is read from the appropriate
        node features file which belongs to the current process
    edge_features : dictionary
        dictionary where edge_features are stored. This information is read from the appropriate
        edge feature files whose ownership is assigned to the current process
    node_feat_tids: 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.
    edge_feat_tids : dictionary
        a dictionary in which keys are edge-type names and values are triplets of the format
        <feat-name, start-per-type-idx, end-per-type-idx>. This triplet is used to identify
        the chunk of feature data for which current process is responsible for
    edge_data : dictionary
        dictionary which is used to store edge information as read from appropriate files assigned
        to each process.
    id_lookup : instance of class DistLookupService
       Distributed lookup service used to map global-nids to respective partition-ids and
       shuffle-global-nids
    ntypes_ntypeid_map : dictionary
        mappings between node type names and node type ids
    ntypes_gnid_range_map : dictionary
        mapping between node type names and global_nids which belong to the keys in this dictionary
    etypes_geid_range_map : dictionary
        mapping between edge type names and global_eids which are assigned to the edges of this
        edge_type
    ntid_ntype_map : dictionary
        mapping between node type id and no of nodes which belong to each node_type_id
    schema_map : dictionary
        is the data structure read from the metadata json file for the input graph

    Returns:
    --------
    dictionary :
        the input argument, node_data dictionary, is updated with the node data received from other processes
        in the world. The node data is received by each rank in the process of data shuffling.
    dictionary :
        node features dictionary which has node features for the nodes which are owned by the current
        process
    dictionary :
        list of global_nids for the nodes whose node features are received when node features shuffling was
        performed in the `exchange_features` function call
    dictionary :
        the input argument, edge_data dictionary, is updated with the edge data received from other processes
        in the world. The edge data is received by each rank in the process of data shuffling.
    dictionary :
        edge features dictionary which has edge features. These destination end points of these edges
        are owned by the current process
    dictionary :
        list of global_eids for the edges whose edge features are received when edge features shuffling
        was performed in the `exchange_features` function call
    
tools/distpartitioning/data_shuffle.py