def test_OnDiskDataset_heterogeneous(
    auto_cast, 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.generate_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,
            auto_cast_to_optimal_dtype=auto_cast,
        ).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()
        )
        expected_dtype = torch.int32 if auto_cast else torch.int64
        assert graph.indices.dtype == expected_dtype
        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.HeteroItemSet)
        assert isinstance(tasks[0].validation_set, gb.HeteroItemSet)
        assert isinstance(tasks[0].test_set, gb.HeteroItemSet)
        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,
            dataset.all_nodes_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
Preprocess and instantiate OnDiskDataset for heterogeneous graph.
tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py