def test_OnDiskDataset_TVTSet_HeteroItemSet_node_pairs_labels():
    """Test TVTSet which returns HeteroItemSet with node pairs and labels."""
    with tempfile.TemporaryDirectory() as test_dir:
        train_seeds = np.arange(2000).reshape(1000, 2)
        train_seeds_path = os.path.join(test_dir, "train_seeds.npy")
        np.save(train_seeds_path, train_seeds)
        train_labels = np.random.randint(0, 10, size=1000)
        train_labels_path = os.path.join(test_dir, "train_labels.npy")
        np.save(train_labels_path, train_labels)

        validation_seeds = np.arange(2000, 4000).reshape(1000, 2)
        validation_seeds_path = os.path.join(test_dir, "validation_seeds.npy")
        np.save(validation_seeds_path, validation_seeds)
        validation_labels = np.random.randint(0, 10, size=1000)
        validation_labels_path = os.path.join(test_dir, "validation_labels.npy")
        np.save(validation_labels_path, validation_labels)

        test_seeds = np.arange(4000, 6000).reshape(1000, 2)
        test_seeds_path = os.path.join(test_dir, "test_seeds.npy")
        np.save(test_seeds_path, test_seeds)
        test_labels = np.random.randint(0, 10, size=1000)
        test_labels_path = os.path.join(test_dir, "test_labels.npy")
        np.save(test_labels_path, test_labels)

        yaml_content = f"""
            tasks:
              - name: edge_classification
                train_set:
                  - type: paper:cites:paper
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {train_seeds_path}
                      - name: labels
                        format: numpy
                        in_memory: true
                        path: {train_labels_path}
                  - type: author:writes:paper
                    data:
                      - name: seeds
                        format: numpy
                        path: {train_seeds_path}
                      - name: labels
                        format: numpy
                        path: {train_labels_path}
                validation_set:
                  - type: paper:cites:paper
                    data:
                      - name: seeds
                        format: numpy
                        path: {validation_seeds_path}
                      - name: labels
                        format: numpy
                        path: {validation_labels_path}
                  - type: author:writes:paper
                    data:
                      - name: seeds
                        format: numpy
                        path: {validation_seeds_path}
                      - name: labels
                        format: numpy
                        path: {validation_labels_path}
                test_set:
                  - type: paper:cites:paper
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {test_seeds_path}
                      - name: labels
                        format: numpy
                        in_memory: true
                        path: {test_labels_path}
                  - type: author:writes:paper
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {test_seeds_path}
                      - name: labels
                        format: numpy
                        in_memory: true
                        path: {test_labels_path}
        """
        dataset = write_yaml_and_load_dataset(yaml_content, test_dir)

        # Verify train set.
        train_set = dataset.tasks[0].train_set
        assert len(train_set) == 2000
        assert isinstance(train_set, gb.HeteroItemSet)
        for i, item in enumerate(train_set):
            assert isinstance(item, dict)
            assert len(item) == 1
            key = list(item.keys())[0]
            assert key in ["paper:cites:paper", "author:writes:paper"]
            node_pair, label = item[key]
            assert node_pair[0] == train_seeds[i % 1000][0]
            assert node_pair[1] == train_seeds[i % 1000][1]
            assert label == train_labels[i % 1000]
        assert train_set.names == ("seeds", "labels")
        train_set = None

        # Verify validation set.
        validation_set = dataset.tasks[0].validation_set
        assert len(validation_set) == 2000
        assert isinstance(validation_set, gb.HeteroItemSet)
        for i, item in enumerate(validation_set):
            assert isinstance(item, dict)
            assert len(item) == 1
            key = list(item.keys())[0]
            assert key in ["paper:cites:paper", "author:writes:paper"]
            node_pair, label = item[key]
            assert node_pair[0] == validation_seeds[i % 1000][0]
            assert node_pair[1] == validation_seeds[i % 1000][1]
            assert label == validation_labels[i % 1000]
        assert validation_set.names == ("seeds", "labels")
        validation_set = None

        # Verify test set.
        test_set = dataset.tasks[0].test_set
        assert len(test_set) == 2000
        assert isinstance(test_set, gb.HeteroItemSet)
        for i, item in enumerate(test_set):
            assert isinstance(item, dict)
            assert len(item) == 1
            key = list(item.keys())[0]
            assert key in ["paper:cites:paper", "author:writes:paper"]
            node_pair, label = item[key]
            assert node_pair[0] == test_seeds[i % 1000][0]
            assert node_pair[1] == test_seeds[i % 1000][1]
            assert label == test_labels[i % 1000]
        assert test_set.names == ("seeds", "labels")
        test_set = None
        dataset = None
Test TVTSet which returns HeteroItemSet with node pairs and labels.
tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py