def assign_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data):
    """
    Utility function to assign shuffle_global_eids to edges
    edge_data gets converted from [global_src_nid, global_dst_nid, global_type_eid, etype]
    to [shuffle_global_src_nid, shuffle_global_dst_nid, global_src_nid, global_dst_nid, global_type_eid, etype]

    Parameters:
    -----------
    rank : integer
        rank of the current process
    world_size : integer
        total count of processes in execution
    num_parts : integer
        total number of output graph partitions
    edge_data : numpy ndarray
        edge data as read from xxx_edges.txt file

    Returns:
    --------
    integer
        shuffle_global_eid_start, which indicates the starting value from which shuffle_global-ids are assigned to edges
        on this rank
    """
    # get prefix sum of edge counts per rank to locate the starting point
    # from which global-ids to edges are assigned in the current rank
    local_row_counts = []
    for local_part_id in range(num_parts // world_size):
        local_row_counts.append(
            edge_data[constants.GLOBAL_SRC_ID + "/" + str(local_part_id)].shape[
                0
            ]
        )

    shuffle_global_eid_offset = []
    prefix_sum_edges = allgather_sizes(local_row_counts, world_size, num_parts)
    for local_part_id in range(num_parts // world_size):
        shuffle_global_eid_start = prefix_sum_edges[
            rank + (local_part_id * world_size)
        ]
        shuffle_global_eid_end = prefix_sum_edges[
            rank + 1 + (local_part_id * world_size)
        ]
        shuffle_global_eids = np.arange(
            shuffle_global_eid_start, shuffle_global_eid_end, dtype=np.int64
        )
        edge_data[
            constants.SHUFFLE_GLOBAL_EID + "/" + str(local_part_id)
        ] = shuffle_global_eids
        shuffle_global_eid_offset.append(shuffle_global_eid_start)

    return shuffle_global_eid_offset

    Utility function to assign shuffle_global_eids to edges
    edge_data gets converted from [global_src_nid, global_dst_nid, global_type_eid, etype]
    to [shuffle_global_src_nid, shuffle_global_dst_nid, global_src_nid, global_dst_nid, global_type_eid, etype]

    Parameters:
    -----------
    rank : integer
        rank of the current process
    world_size : integer
        total count of processes in execution
    num_parts : integer
        total number of output graph partitions
    edge_data : numpy ndarray
        edge data as read from xxx_edges.txt file

    Returns:
    --------
    integer
        shuffle_global_eid_start, which indicates the starting value from which shuffle_global-ids are assigned to edges
        on this rank
    
tools/distpartitioning/globalids.py