code: def test_bulk_sampler_partitions(dask_client, scratch_dir, mg_input):
      el = (
          karate.get_edgelist()
          .reset_index()
          .rename(columns={"index": "eid"})
      )
      el["eid"] = el["eid"].astype("int32")
      el["etp"] = cupy.int32(0)
  
      G = cugraph.Graph(directed=True)
      G.from_dask_cudf_edgelist(
          dask_cudf.from_cudf(el, npartitions=2),
          source="src",
          destination="dst",
          edge_attr=["wgt", "eid", "etp"],
      )
  
      samples_path = os.path.join(
          scratch_dir, "test_bulk_sampler_partitions_mg"
      )
      if os.path.exists(samples_path):
          shutil.rmtree(samples_path)
      os.makedirs(samples_path)
  
      bs = BulkSampler(
          batch_size=3,
          output_path=samples_path,
          graph=G,
          fanout_vals=[2, 2],
          with_replacement=False,
          batches_per_partition=2,
          renumber=True,
      )
  
      batches = cudf.DataFrame(
          {
              "start": cudf.Series(
                  [0, 5, 6, 10, 15, 17, 18, 9, 23], dtype="int32"
              ),
              "batch": cudf.Series(
                  [0, 0, 0, 1, 1, 1, 2, 2, 2], dtype="int32"
              ),
          }
      )
  
      if mg_input:
          batches = dask_cudf.from_cudf(batches, npartitions=4)
  
      bs.add_batches(
          batches, start_col_name="start", batch_col_name="batch"
      )
      bs.flush()
  
      for file in os.listdir(samples_path):
          start_batch_id, end_batch_id = [
              int(x)
              for x in re.match(
                  r"batch=([0-9]+)-([0-9]+).parquet", file
              ).groups()
          ]
  
          recovered_samples = cudf.read_parquet(
              os.path.join(samples_path, file)
          )
          recovered_map = recovered_samples.map
          recovered_samples = recovered_samples.drop(
              "map", axis=1
          ).dropna()
  
          for current_batch_id in range(
              start_batch_id, end_batch_id + 1
          ):
              map_start_ix = recovered_map.iloc[
                  current_batch_id - start_batch_id
              ]
              map_end_ix = recovered_map.iloc[
                  current_batch_id - start_batch_id + 1
              ]
              map_current_batch = recovered_map.iloc[
                  map_start_ix:map_end_ix
              ]
              n_unique = cudf.concat(
                  [
                      recovered_samples[
                          recovered_samples.batch_id == current_batch_id
                      ].sources,
                      recovered_samples[
                          recovered_samples.batch_id == current_batch_id
                      ].destinations,
                  ]
              ).nunique()
              assert len(map_current_batch) == n_unique
  
file_path: cugraph/python/cugraph/cugraph/tests/sampling/test_bulk_sampler_mg.py
imports: import os
  import shutil
  import re
  import pytest
  import cudf
  import cupy
  import cugraph
  import dask_cudf
  from cugraph.datasets import karate, email_Eu_core
  from cugraph.gnn import BulkSampler
  from cugraph.utilities.utils import create_directory_with_overwrite
label_desc: Use image node_id 3 for calling a global function with example usage: test_bulk_sampler_partitions(dask_client, scratch_dir, mg_input) without return types