def exchange_features(
    rank,
    world_size,
    num_parts,
    feature_tids,
    type_id_map,
    id_lookup,
    feature_data,
    feat_type,
    data,
):
    """
    This function is used to shuffle node features so that each process will receive
    all the node features whose corresponding nodes are owned by the same process.
    The mapping procedure to identify the owner process is not straight forward. The
    following steps are used to identify the owner processes for the locally read node-
    features.
    a. Compute the global_nids for the locally read node features. Here metadata json file
        is used to identify the corresponding global_nids. Please note that initial graph input
        nodes.txt files are sorted based on node_types.
    b. Using global_nids and metis partitions owner processes can be easily identified.
    c. Now each process sends the global_nids for which shuffle_global_nids are needed to be
        retrieved.
    d. After receiving the corresponding shuffle_global_nids these ids are added to the
        node_data and edge_data dictionaries

    This pipeline assumes all the input data in numpy format, except node/edge features which
    are maintained as tensors throughout the various stages of the pipeline execution.

    Parameters:
    -----------
    rank : int
        rank of the current process
    world_size : int
        total no. of participating processes.
    feature_tids : dictionary
        dictionary with keys as node-type names with suffixes as feature names
        and value is a dictionary. This dictionary contains information about
        node-features associated with a given node-type and value is a list.
        This list contains a of indexes, like [starting-idx, ending-idx) which
        can be used to index into the node feature tensors read from
        corresponding input files.
    type_id_map : dictionary
        mapping between type names and global_ids, of either nodes or edges,
        which belong to the keys in this dictionary
    id_lookup : instance of class DistLookupService
       Distributed lookup service used to map global-nids to respective
       partition-ids and shuffle-global-nids
    feat_type : string
        this is used to distinguish which features are being exchanged. Please
        note that for nodes ownership is clearly defined and for edges it is
        always assumed that destination end point of the edge defines the
        ownership of that particular edge
    data: dicitonary
        dictionry in which node or edge features are stored and this information
        is read from the appropriate node features file which belongs to the
        current process

    Returns:
    --------
    dictionary :
        a dictionary is returned where keys are type names and
        feature data are the values
    list :
        a dictionary of global_ids either nodes or edges whose features are
        received during the data shuffle process
    """
    start = timer()
    own_features = {}
    own_global_ids = {}

    # To iterate over the node_types and associated node_features
    for feat_key, type_info in feature_tids.items():
        # To iterate over the feature data, of a given (node or edge )type
        # type_info is a list of 3 elements (as shown below):
        #   [feature-name, starting-idx, ending-idx]
        #       feature-name is the name given to the feature-data,
        #       read from the input metadata file
        #       [starting-idx, ending-idx) specifies the range of indexes
        #        associated with the features data
        # Determine the owner process for these features.
        # Note that the keys in the node features (and similarly edge features)
        # dictionary is of the following format:
        #   `node_type/feature_name/local_part_id`:
        #    where node_type and feature_name are self-explanatory and
        #    local_part_id denotes the partition-id, in the local process,
        #    which will be used a suffix to store all the information of a
        #    given partition which is processed by the current process. Its
        #    values start from 0 onwards, for instance 0, 1, 2 ... etc.
        #    local_part_id can be easily mapped to global partition id very
        #    easily, using cyclic ordering. All local_part_ids = 0 from all
        #    processes will form global partition-ids between 0 and world_size-1.
        #    Similarly all local_part_ids = 1 from all processes will form
        #    global partition ids in the range [world_size, 2*world_size-1] and
        #    so on.
        tokens = feat_key.split("/")
        assert len(tokens) == 3
        type_name = tokens[0]
        feat_name = tokens[1]
        logging.debug(f"[Rank: {rank}] processing feature: {feat_key}")

        for feat_info in type_info:
            # Compute the global_id range for this feature data
            type_id_start = int(feat_info[0])
            type_id_end = int(feat_info[1])
            begin_global_id = type_id_map[type_name][0]
            gid_start = begin_global_id + type_id_start
            gid_end = begin_global_id + type_id_end

            # Check if features exist for this type_name + feat_name.
            # This check should always pass, because feature_tids are built
            # by reading the input metadata json file for existing features.
            assert feat_key in feature_data

            for local_part_id in range(num_parts // world_size):
                featdata_key = feature_data[feat_key]

                # Synchronize for each feature
                dist.barrier()
                own_features, own_global_ids = exchange_feature(
                    rank,
                    data,
                    id_lookup,
                    feat_type,
                    feat_key,
                    featdata_key,
                    gid_start,
                    gid_end,
                    type_id_start,
                    type_id_end,
                    local_part_id,
                    world_size,
                    num_parts,
                    own_features,
                    own_global_ids,
                )

    end = timer()
    logging.info(
        f"[Rank: {rank}] Total time for feature exchange: {timedelta(seconds = end - start)}"
    )
    for k, v in own_features.items():
        logging.debug(f"Rank: {rank}] Key - {k} Value - {v.shape}")
    return own_features, own_global_ids

    This function is used to shuffle node features so that each process will receive
    all the node features whose corresponding nodes are owned by the same process.
    The mapping procedure to identify the owner process is not straight forward. The
    following steps are used to identify the owner processes for the locally read node-
    features.
    a. Compute the global_nids for the locally read node features. Here metadata json file
        is used to identify the corresponding global_nids. Please note that initial graph input
        nodes.txt files are sorted based on node_types.
    b. Using global_nids and metis partitions owner processes can be easily identified.
    c. Now each process sends the global_nids for which shuffle_global_nids are needed to be
        retrieved.
    d. After receiving the corresponding shuffle_global_nids these ids are added to the
        node_data and edge_data dictionaries

    This pipeline assumes all the input data in numpy format, except node/edge features which
    are maintained as tensors throughout the various stages of the pipeline execution.

    Parameters:
    -----------
    rank : int
        rank of the current process
    world_size : int
        total no. of participating processes.
    feature_tids : dictionary
        dictionary with keys as node-type names with suffixes as feature names
        and value is a dictionary. This dictionary contains information about
        node-features associated with a given node-type and value is a list.
        This list contains a of indexes, like [starting-idx, ending-idx) which
        can be used to index into the node feature tensors read from
        corresponding input files.
    type_id_map : dictionary
        mapping between type names and global_ids, of either nodes or edges,
        which belong to the keys in this dictionary
    id_lookup : instance of class DistLookupService
       Distributed lookup service used to map global-nids to respective
       partition-ids and shuffle-global-nids
    feat_type : string
        this is used to distinguish which features are being exchanged. Please
        note that for nodes ownership is clearly defined and for edges it is
        always assumed that destination end point of the edge defines the
        ownership of that particular edge
    data: dicitonary
        dictionry in which node or edge features are stored and this information
        is read from the appropriate node features file which belongs to the
        current process

    Returns:
    --------
    dictionary :
        a dictionary is returned where keys are type names and
        feature data are the values
    list :
        a dictionary of global_ids either nodes or edges whose features are
        received during the data shuffle process
    
tools/distpartitioning/data_shuffle.py