def assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data):
    """
    Utility function to assign shuffle global ids to nodes at a given rank
    node_data gets converted from [ntype, global_type_nid, global_nid]
    to [shuffle_global_nid, ntype, global_type_nid, global_nid, part_local_type_nid]
    where shuffle_global_nid : global id of the node after data shuffle
            ntype : node-type as read from xxx_nodes.txt
            global_type_nid : node-type-id as read from xxx_nodes.txt
            global_nid : node-id as read from xxx_nodes.txt, implicitly
                            this is the line no. in the file
            part_local_type_nid : type_nid assigned by the current rank within its scope

    Parameters:
    -----------
    rank : integer
        rank of the process
    world_size : integer
        total number of processes used in the process group
    num_parts : integer
        total number of output graph partitions
    node_data : dictionary
        node_data is a dictionary with keys as column names and values as numpy arrays
    """
    # Compute prefix sum to determine node-id offsets
    local_row_counts = []
    for local_part_id in range(num_parts // world_size):
        local_row_counts.append(
            node_data[constants.GLOBAL_NID + "/" + str(local_part_id)].shape[0]
        )

    # Perform allgather to compute the local offsets.
    prefix_sum_nodes = allgather_sizes(local_row_counts, world_size, num_parts)

    for local_part_id in range(num_parts // world_size):
        shuffle_global_nid_start = prefix_sum_nodes[
            rank + (local_part_id * world_size)
        ]
        shuffle_global_nid_end = prefix_sum_nodes[
            rank + 1 + (local_part_id * world_size)
        ]
        shuffle_global_nids = np.arange(
            shuffle_global_nid_start, shuffle_global_nid_end, dtype=np.int64
        )
        node_data[
            constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)
        ] = shuffle_global_nids

    Utility function to assign shuffle global ids to nodes at a given rank
    node_data gets converted from [ntype, global_type_nid, global_nid]
    to [shuffle_global_nid, ntype, global_type_nid, global_nid, part_local_type_nid]
    where shuffle_global_nid : global id of the node after data shuffle
            ntype : node-type as read from xxx_nodes.txt
            global_type_nid : node-type-id as read from xxx_nodes.txt
            global_nid : node-id as read from xxx_nodes.txt, implicitly
                            this is the line no. in the file
            part_local_type_nid : type_nid assigned by the current rank within its scope

    Parameters:
    -----------
    rank : integer
        rank of the process
    world_size : integer
        total number of processes used in the process group
    num_parts : integer
        total number of output graph partitions
    node_data : dictionary
        node_data is a dictionary with keys as column names and values as numpy arrays
    
tools/distpartitioning/globalids.py