code: 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
  
file_path: dgl/tools/distpartitioning/globalids.py
imports: import itertools
  import operator
  import constants
  import numpy
  import torch
  from dist_lookup import DistLookupService
  from gloo_wrapper import allgather_sizes, alltoallv_cpu
  from utils import memory_snapshot
label_desc: Use image node_id 3 for calling a global function with example usage: assign_shuffle_global_nids_nodes(rank, world_size, num_parts, node_data) without return types