def exchange_edge_data(rank, world_size, num_parts, edge_data, id_lookup):
    """
    Exchange edge_data among processes in the world.
    Prepare list of sliced data targeting each process and trigger
    alltoallv_cpu to trigger messaging api

    Parameters:
    -----------
    rank : int
        rank of the process
    world_size : int
        total no. of processes
    edge_data : dictionary
        edge information, as a dicitonary which stores column names as keys and values
        as column data. This information is read from the edges.txt file.
    id_lookup : DistLookupService instance
        this object will be used to retrieve ownership information of nodes

    Returns:
    --------
    dictionary :
        the input argument, edge_data, is updated with the edge data received by other processes
        in the world.
    """

    # Synchronize at the beginning of this function
    dist.barrier()

    # Prepare data for each rank in the cluster.
    timer_start = timer()

    CHUNK_SIZE = 100 * 1000 * 1000  # 100 * 8 * 5 = 1 * 4 = 8 GB/message/node
    num_edges = edge_data[constants.GLOBAL_SRC_ID].shape[0]
    all_counts = allgather_sizes(
        [num_edges], world_size, num_parts, return_sizes=True
    )
    max_edges = np.amax(all_counts)
    all_edges = np.sum(all_counts)
    num_chunks = (max_edges // CHUNK_SIZE) + (
        0 if (max_edges % CHUNK_SIZE == 0) else 1
    )
    LOCAL_CHUNK_SIZE = (num_edges // num_chunks) + (
        0 if (num_edges % num_chunks == 0) else 1
    )
    logging.debug(
        f"[Rank: {rank} Edge Data Shuffle - max_edges: {max_edges}, \
                        local_edges: {num_edges} and num_chunks: {num_chunks} \
                        Total edges: {all_edges} Local_CHUNK_SIZE: {LOCAL_CHUNK_SIZE}"
    )

    for local_part_id in range(num_parts // world_size):
        local_src_ids = []
        local_dst_ids = []
        local_type_eids = []
        local_etype_ids = []
        local_eids = []

        for chunk in range(num_chunks):
            chunk_start = chunk * LOCAL_CHUNK_SIZE
            chunk_end = (chunk + 1) * LOCAL_CHUNK_SIZE

            logging.debug(
                f"[Rank: {rank}] EdgeData Shuffle: processing \
                    local_part_id: {local_part_id} and chunkid: {chunk}"
            )
            cur_src_id = edge_data[constants.GLOBAL_SRC_ID][
                chunk_start:chunk_end
            ]
            cur_dst_id = edge_data[constants.GLOBAL_DST_ID][
                chunk_start:chunk_end
            ]
            cur_type_eid = edge_data[constants.GLOBAL_TYPE_EID][
                chunk_start:chunk_end
            ]
            cur_etype_id = edge_data[constants.ETYPE_ID][chunk_start:chunk_end]
            cur_eid = edge_data[constants.GLOBAL_EID][chunk_start:chunk_end]

            input_list = []
            owner_ids = id_lookup.get_partition_ids(cur_dst_id)
            for idx in range(world_size):
                send_idx = owner_ids == (idx + local_part_id * world_size)
                send_idx = send_idx.reshape(cur_src_id.shape[0])
                filt_data = np.column_stack(
                    (
                        cur_src_id[send_idx == 1],
                        cur_dst_id[send_idx == 1],
                        cur_type_eid[send_idx == 1],
                        cur_etype_id[send_idx == 1],
                        cur_eid[send_idx == 1],
                    )
                )
                if filt_data.shape[0] <= 0:
                    input_list.append(torch.empty((0, 5), dtype=torch.int64))
                else:
                    input_list.append(torch.from_numpy(filt_data))

            # Now send newly formed chunk to others.
            dist.barrier()
            output_list = alltoallv_cpu(
                rank, world_size, input_list, retain_nones=False
            )

            # Replace the values of the edge_data, with the received data from all the other processes.
            rcvd_edge_data = torch.cat(output_list).numpy()
            local_src_ids.append(rcvd_edge_data[:, 0])
            local_dst_ids.append(rcvd_edge_data[:, 1])
            local_type_eids.append(rcvd_edge_data[:, 2])
            local_etype_ids.append(rcvd_edge_data[:, 3])
            local_eids.append(rcvd_edge_data[:, 4])

        edge_data[
            constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
        ] = np.concatenate(local_src_ids)
        edge_data[
            constants.GLOBAL_DST_ID + "/" + str(local_part_id)
        ] = np.concatenate(local_dst_ids)
        edge_data[
            constants.GLOBAL_TYPE_EID + "/" + str(local_part_id)
        ] = np.concatenate(local_type_eids)
        edge_data[
            constants.ETYPE_ID + "/" + str(local_part_id)
        ] = np.concatenate(local_etype_ids)
        edge_data[
            constants.GLOBAL_EID + "/" + str(local_part_id)
        ] = np.concatenate(local_eids)

    # Check if the data was exchanged correctly
    local_edge_count = 0
    for local_part_id in range(num_parts // world_size):
        local_edge_count += edge_data[
            constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
        ].shape[0]
    shuffle_edge_counts = allgather_sizes(
        [local_edge_count], world_size, num_parts, return_sizes=True
    )
    shuffle_edge_total = np.sum(shuffle_edge_counts)
    assert shuffle_edge_total == all_edges

    timer_end = timer()
    logging.info(
        f"[Rank: {rank}] Time to send/rcv edge data: {timedelta(seconds=timer_end-timer_start)}"
    )

    # Clean up.
    edge_data.pop(constants.GLOBAL_SRC_ID)
    edge_data.pop(constants.GLOBAL_DST_ID)
    edge_data.pop(constants.GLOBAL_TYPE_EID)
    edge_data.pop(constants.ETYPE_ID)
    edge_data.pop(constants.GLOBAL_EID)

    return edge_data

    Exchange edge_data among processes in the world.
    Prepare list of sliced data targeting each process and trigger
    alltoallv_cpu to trigger messaging api

    Parameters:
    -----------
    rank : int
        rank of the process
    world_size : int
        total no. of processes
    edge_data : dictionary
        edge information, as a dicitonary which stores column names as keys and values
        as column data. This information is read from the edges.txt file.
    id_lookup : DistLookupService instance
        this object will be used to retrieve ownership information of nodes

    Returns:
    --------
    dictionary :
        the input argument, edge_data, is updated with the edge data received by other processes
        in the world.
    
tools/distpartitioning/data_shuffle.py