code: def load_disk_dataset(
      dataset,
      dataset_dir=".",
      reverse_edges=True,
      replication_factor=1,
      add_edge_types=False,
  ):
      from pathlib import Path
  
      path = Path(dataset_dir) / dataset
      parquet_path = path / "parquet"
  
      n_workers = get_n_workers()
  
      with open(os.path.join(path, "meta.json")) as meta_file:
          meta = json.load(meta_file)
  
      (
          node_offsets,
          node_offsets_replicated,
          total_num_nodes,
      ) = assign_offsets_pyg(
          meta["num_nodes"], replication_factor=replication_factor
      )
  
      edge_index_dict = {}
      for edge_type in meta["num_edges"].keys():
          print(f"Loading edge index for edge type {edge_type}")
  
          can_edge_type = tuple(edge_type.split("__"))
          edge_index_dict[can_edge_type] = dask_cudf.read_parquet(
              Path(parquet_path) / edge_type / "edge_index.parquet"
          ).repartition(n_workers * 2)
  
          edge_index_dict[can_edge_type][
              "src"
          ] += node_offsets_replicated[can_edge_type[0]]
          edge_index_dict[can_edge_type][
              "dst"
          ] += node_offsets_replicated[can_edge_type[-1]]
  
          edge_index_dict[can_edge_type] = edge_index_dict[
              can_edge_type
          ]
  
          if replication_factor > 1:
              edge_index_dict[can_edge_type] = edge_index_dict[
                  can_edge_type
              ].map_partitions(
                  _replicate_df,
                  replication_factor,
                  {
                      "src": meta["num_nodes"][can_edge_type[0]],
                      "dst": meta["num_nodes"][can_edge_type[2]],
                  },
                  meta=cudf.DataFrame(
                      {
                          "src": cudf.Series(dtype="int64"),
                          "dst": cudf.Series(dtype="int64"),
                      }
                  ),
              )
  
          gc.collect()
  
          if reverse_edges:
              edge_index_dict[can_edge_type] = edge_index_dict[
                  can_edge_type
              ].rename(columns={"src": "dst", "dst": "src"})
  
      # Assign numeric edge type ids based on lexicographic order
      edge_offsets = {}
      edge_count = 0
      for num_edge_type, can_edge_type in enumerate(
          sorted(edge_index_dict.keys())
      ):
          if add_edge_types:
              edge_index_dict[can_edge_type]["etp"] = cupy.int32(
                  num_edge_type
              )
          edge_offsets[can_edge_type] = edge_count
          edge_count += len(edge_index_dict[can_edge_type])
  
      all_edges_df = dask_cudf.concat(list(edge_index_dict.values()))
  
      del edge_index_dict
      gc.collect()
  
      node_labels = {}
      for node_type, offset in node_offsets_replicated.items():
          print(
              f"Loading node labels for node type {node_type} (offset={offset})"
          )
          node_label_path = os.path.join(
              os.path.join(parquet_path, node_type),
              "node_label.parquet",
          )
          if os.path.exists(node_label_path):
              node_labels[node_type] = (
                  dask_cudf.read_parquet(node_label_path)
                  .repartition(n_workers)
                  .drop("label", axis=1)
                  .persist()
              )
              node_labels[node_type]["node"] += offset
              node_labels[node_type] = node_labels[node_type].persist()
  
              if replication_factor > 1:
                  node_labels[node_type] = node_labels[
                      node_type
                  ].map_partitions(
                      _replicate_df,
                      replication_factor,
                      {"node": meta["num_nodes"][node_type]},
                      meta=cudf.DataFrame(
                          {"node": cudf.Series(dtype="int64")}
                      ),
                  )
  
              gc.collect()
  
      node_labels_df = dask_cudf.concat(
          list(node_labels.values())
      ).reset_index(drop=True)
  
      del node_labels
      gc.collect()
  
      return (
          all_edges_df,
          node_labels_df,
          node_offsets_replicated,
          edge_offsets,
          total_num_nodes,
      )
  
imports: import logging
  import warnings
  import argparse
  import traceback
  from cugraph.testing.mg_utils import generate_edgelist_rmat, get_allocation_counts_dask_lazy, sizeof_fmt, get_peak_output_ratio_across_workers, start_dask_client, stop_dask_client, enable_spilling
  from cugraph.structure.symmetrize import symmetrize
  from cugraph.gnn import BulkSampler
  import cugraph
  import json
  import re
  import os
  import gc
  from time import sleep, perf_counter
  from math import ceil
  import pandas
  import numpy
  import cupy
  import cudf
  import dask_cudf
  import dask.dataframe
  from dask.distributed import default_client
  from cugraph.dask import get_n_workers
  from typing import Optional, Union, Dict
label_desc: Use image node_id 9 for calling a global function with example usage: load_disk_dataset(dataset, dataset_dir, reverse_edges, replication_factor, add_edge_types) and returns: all_edges_df,  node_labels_df,  node_offsets_replicated,  edge_offsets,  total_num_nodes