code: def test_OnDiskDataset_heterogeneous(
      include_original_edge_id, edge_fmt
  ):
      """Preprocess and instantiate OnDiskDataset for heterogeneous graph."""
      with tempfile.TemporaryDirectory() as test_dir:
          dataset_name = "OnDiskDataset_hetero"
          num_nodes = {
              "user": 1000,
              "item": 2000,
          }
          num_edges = {
              ("user", "follow", "user"): 10000,
              ("user", "click", "item"): 20000,
          }
          num_classes = 10
          gbt.genereate_raw_data_for_hetero_dataset(
              test_dir,
              dataset_name,
              num_nodes,
              num_edges,
              num_classes,
              edge_fmt=edge_fmt,
          )
  
          dataset = gb.OnDiskDataset(
              test_dir,
              include_original_edge_id=include_original_edge_id,
          ).load()
  
          assert dataset.dataset_name == dataset_name
  
          graph = dataset.graph
          assert isinstance(graph, gb.FusedCSCSamplingGraph)
          assert graph.total_num_nodes == sum(
              num_nodes for num_nodes in num_nodes.values()
          )
          assert graph.total_num_edges == sum(
              num_edge for num_edge in num_edges.values()
          )
          assert (
              graph.node_attributes is not None
              and "feat" in graph.node_attributes
          )
          assert (
              graph.edge_attributes is not None
              and "feat" in graph.edge_attributes
          )
          assert (
              not include_original_edge_id
          ) or gb.ORIGINAL_EDGE_ID in graph.edge_attributes
  
          tasks = dataset.tasks
          assert len(tasks) == 1
          assert isinstance(tasks[0].train_set, gb.ItemSetDict)
          assert isinstance(tasks[0].validation_set, gb.ItemSetDict)
          assert isinstance(tasks[0].test_set, gb.ItemSetDict)
          assert tasks[0].metadata["num_classes"] == num_classes
          assert tasks[0].metadata["name"] == "node_classification"
  
          assert (
              dataset.feature.size("node", "user", "feat")[0]
              == num_classes
          )
          assert (
              dataset.feature.size("node", "item", "feat")[0]
              == num_classes
          )
  
          for itemset in [
              tasks[0].train_set,
              tasks[0].validation_set,
              tasks[0].test_set,
          ]:
              datapipe = gb.ItemSampler(itemset, batch_size=10)
              datapipe = datapipe.sample_neighbor(graph, [-1])
              datapipe = datapipe.fetch_feature(
                  dataset.feature, node_feature_keys={"user": ["feat"]}
              )
              dataloader = gb.DataLoader(datapipe)
              for _ in dataloader:
                  pass
  
          graph = None
          tasks = None
          dataset = None
  
imports: import os
  import pickle
  import random
  import re
  import tempfile
  import unittest
  import warnings
  import numpy
  import pandas
  import pydantic
  import pytest
  import torch
  import yaml
  from dgl import graphbolt
  from dgl.base import DGLWarning
  from ..None import gb_test_utils
label_desc: Use image node_id 35 for calling a global function with example usage: test_OnDiskDataset_heterogeneous(include_original_edge_id, edge_fmt) without return types