code: def lookup_shuffle_global_nids_edges(
      rank, world_size, num_parts, edge_data, id_lookup, node_data
  ):
      """
      This function is a helper function used to lookup shuffle-global-nids for a given set of
      global-nids using a distributed lookup service.
  
      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
      edge_data : dictionary
          edge_data is a dicitonary with keys as column names and values as numpy arrays representing
          all the edges present in the current graph partition
      id_lookup : instance of DistLookupService class
          instance of a distributed lookup service class which is used to retrieve partition-ids and
          shuffle-global-nids for any given set of global-nids
      node_data : dictionary
          node_data is a dictionary with keys as column names and values as numpy arrays representing
          all the nodes owned by the current process
  
      Returns:
      --------
      dictionary :
          dictionary where keys are column names and values are numpy arrays representing all the
          edges present in the current graph partition
      """
      # Make sure that the outgoing message size does not exceed 2GB in size.
      # Even though gloo can handle upto 10GB size of data in the outgoing messages,
      # it needs additional memory to store temporary information into the buffers which will increase
      # the memory needs of the process.
      MILLION = 1000 * 1000
      BATCH_SIZE = 250 * MILLION
      memory_snapshot("GlobalToShuffleIDMapBegin: ", rank)
  
      local_nids = []
      local_shuffle_nids = []
      for local_part_id in range(num_parts // world_size):
          local_nids.append(
              node_data[constants.GLOBAL_NID + "/" + str(local_part_id)]
          )
          local_shuffle_nids.append(
              node_data[
                  constants.SHUFFLE_GLOBAL_NID
                  + "/"
                  + str(local_part_id)
              ]
          )
  
      local_nids = np.concatenate(local_nids)
      local_shuffle_nids = np.concatenate(local_shuffle_nids)
  
      for local_part_id in range(num_parts // world_size):
          node_list = edge_data[
              constants.GLOBAL_SRC_ID + "/" + str(local_part_id)
          ]
  
          # Determine the no. of times each process has to send alltoall messages.
          all_sizes = allgather_sizes(
              [node_list.shape[0]],
              world_size,
              num_parts,
              return_sizes=True,
          )
          max_count = np.amax(all_sizes)
          num_splits = max_count // BATCH_SIZE + 1
  
          # Split the message into batches and send.
          splits = np.array_split(node_list, num_splits)
          shuffle_mappings = []
          for item in splits:
              shuffle_ids = id_lookup.get_shuffle_nids(
                  item, local_nids, local_shuffle_nids, world_size
              )
              shuffle_mappings.append(shuffle_ids)
  
          shuffle_ids = np.concatenate(shuffle_mappings)
          assert shuffle_ids.shape[0] == node_list.shape[0]
          edge_data[
              constants.SHUFFLE_GLOBAL_SRC_ID + "/" + str(local_part_id)
          ] = shuffle_ids
  
          # Destination end points of edges are owned by the current node and therefore
          # should have corresponding SHUFFLE_GLOBAL_NODE_IDs.
          # Here retrieve SHUFFLE_GLOBAL_NODE_IDs for the destination end points of local edges.
          uniq_ids, inverse_idx = np.unique(
              edge_data[
                  constants.GLOBAL_DST_ID + "/" + str(local_part_id)
              ],
              return_inverse=True,
          )
          common, idx1, idx2 = np.intersect1d(
              uniq_ids,
              node_data[
                  constants.GLOBAL_NID + "/" + str(local_part_id)
              ],
              assume_unique=True,
              return_indices=True,
          )
          assert len(common) == len(uniq_ids)
  
          edge_data[
              constants.SHUFFLE_GLOBAL_DST_ID + "/" + str(local_part_id)
          ] = node_data[
              constants.SHUFFLE_GLOBAL_NID + "/" + str(local_part_id)
          ][
              idx2
          ][
              inverse_idx
          ]
          assert len(
              edge_data[
                  constants.SHUFFLE_GLOBAL_DST_ID
                  + "/"
                  + str(local_part_id)
              ]
          ) == len(
              edge_data[
                  constants.GLOBAL_DST_ID + "/" + str(local_part_id)
              ]
          )
  
      memory_snapshot(
          "GlobalToShuffleIDMap_AfterLookupServiceCalls: ", rank
      )
      return edge_data
  
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 2 for calling a global function with example usage: lookup_shuffle_global_nids_edges(rank, world_size, num_parts, edge_data, id_lookup, node_data) and returns: edge_data