code: def _prepare_test_data(
      partitions_dir, ntypes, gid_ranges, world_size
  ):
      # read node-id to partition-id mappings from disk
      ntype_partids = []
      for ntype_id, ntype in enumerate(ntypes):
          filename = f"{ntype}.txt"
          assert os.path.isfile(os.path.join(partitions_dir, filename))
  
          read_options = pyarrow.csv.ReadOptions(
              use_threads=True,
              block_size=4096,
              autogenerate_column_names=True,
          )
          parse_options = pyarrow.csv.ParseOptions(delimiter=" ")
  
          with pyarrow.csv.open_csv(
              os.path.join(partitions_dir, "{}.txt".format(ntype)),
              read_options=read_options,
              parse_options=parse_options,
          ) as reader:
              for next_chunk in reader:
                  if next_chunk is None:
                      break
                  next_table = pyarrow.Table.from_batches([next_chunk])
                  ntype_partids.append(next_table["f0"].to_numpy())
  
      # prepare test data for each rank here
      # key = f'rank-{rank}'
      # value is a list of tuple [(method-name, request, response)]
      test_data = {}
      for rank in range(world_size):
          ntype_id = np.random.randint(0, len(ntypes) - 1)
          ntype = ntypes[ntype_id]
          request = (
              np.arange(len(ntype_partids[ntype_id]))
              + gid_ranges[ntypes[ntype_id]][0, 0]
          )
          response = ntype_partids[ntype_id]
  
          test_data[f"rank-{rank}"] = [
              ("getpartitionids", request, response)
          ]
  
      # randomly shuffle the global-nids and retrieve their partition-ids.
      for rank in range(world_size):
          ntype_id = np.random.randint(0, len(ntypes) - 1)
          ntype = ntypes[ntype_id]
          idx = np.arange(len(ntype_partids[ntype_id]))
          request = idx + gid_ranges[ntypes[ntype_id]][0, 0]
  
          np.random.shuffle(idx)
          request = request[idx]
          response = ntype_partids[ntype_id][idx]
  
          test_data[f"rank-{rank}"] = [
              ("getpartitionids", request, response)
          ]
  
      # one final test
      # mix all the ntypes and shuffle randomly
      request = []
      response = []
      for idx in range(len(ntype_partids)):
          request.append(
              np.arange(len(ntype_partids[idx]))
              + gid_ranges[ntypes[idx]][0, 0]
          )
          response.append(ntype_partids[idx])
  
      request = np.concatenate(request)
      response = np.concatenate(response)
  
      idx = np.arange(len(request))
      np.random.shuffle(idx)
      request = request[idx]
      response = response[idx]
      for idx in range(world_size):
          test_data[f"rank-{idx}"] = [
              ("getpartitionids", request, response)
          ]
  
      return test_data
  
imports: import logging
  import os
  import platform
  import tempfile
  from datetime import timedelta
  import dgl
  import numpy
  import pyarrow
  import pytest
  import torch.distributed
  import torch.multiprocessing
  from pytest_utils import create_chunked_dataset
  from tools.distpartitioning import constants, dist_lookup
  from tools.distpartitioning.gloo_wrapper import allgather_sizes
  from tools.distpartitioning.utils import get_idranges, get_ntype_counts_map, read_json
label_desc: Use image node_id 5 for calling a global function with example usage: _prepare_test_data(partitions_dir, ntypes, gid_ranges, world_size) and returns: test_data