code: def augment_edge_data(
      edge_data, lookup_service, edge_tids, rank, world_size, num_parts
  ):
      """
      Add partition-id (rank which owns an edge) column to the edge_data.
  
      Parameters:
      -----------
      edge_data : numpy ndarray
          Edge information as read from the xxx_edges.txt file
      lookup_service : instance of class DistLookupService
         Distributed lookup service used to map global-nids to respective partition-ids and▒
         shuffle-global-nids
      edge_tids: dictionary
          dictionary where keys are canonical edge types and values are list of tuples
          which indicate the range of edges assigned to each of the partitions
      rank : integer
          rank of the current process
      world_size : integer
          total no. of process participating in the communication primitives
      num_parts : integer
          total no. of partitions requested for the input graph
  
      Returns:
      --------
      dictionary :
          dictionary with keys as column names and values as numpy arrays and this information is
          loaded from input dataset files. In addition to this we include additional columns which
          aid this pipelines computation, like constants.OWNER_PROCESS
      """
      # add global_nids to the node_data
      etype_offset = {}
      offset = 0
      for etype_name, tid_range in edge_tids.items():
          etype_offset[etype_name] = offset + int(tid_range[0][0])
          offset += int(tid_range[-1][1])
  
      global_eids = []
      for etype_name, tid_range in edge_tids.items():
          for idx in range(num_parts):
              if map_partid_rank(idx, world_size) == rank:
                  if len(tid_range) > idx:
                      global_eid_start = etype_offset[etype_name]
                      begin = global_eid_start + int(tid_range[idx][0])
                      end = global_eid_start + int(tid_range[idx][1])
                      global_eids.append(
                          np.arange(begin, end, dtype=np.int64)
                      )
  
      global_eids = (
          np.concatenate(global_eids)
          if len(global_eids) > 0
          else np.array([], dtype=np.int64)
      )
      assert (
          global_eids.shape[0] == edge_data[constants.ETYPE_ID].shape[0]
      )
      edge_data[constants.GLOBAL_EID] = global_eids
      return edge_data
  
imports: import json
  import logging
  import os
  from itertools import cycle
  import constants
  import dgl
  import numpy
  import psutil
  import pyarrow
  import torch
  from dgl.distributed.partition import _dump_part_config
  from pyarrow import csv
label_desc: Use image node_id 10 for calling a global function with example usage: augment_edge_data(edge_data, lookup_service, edge_tids, rank, world_size, num_parts) and returns: edge_data