import os import pickle import random import re import tempfile import unittest import warnings import numpy as np import pandas as pd import pydantic import pytest import torch import yaml from dgl import graphbolt as gb from dgl.graphbolt import GBWarning from .. import gb_test_utils as gbt def write_yaml_file(yaml_content, dir): os.makedirs(os.path.join(dir, "preprocessed"), exist_ok=True) yaml_file = os.path.join(dir, "preprocessed/metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) def load_dataset(dataset): with warnings.catch_warnings(): warnings.simplefilter("ignore", category=UserWarning) return dataset.load() def write_yaml_and_load_dataset(yaml_content, dir, force_preprocess=False): write_yaml_file(yaml_content, dir) return load_dataset( gb.OnDiskDataset(dir, force_preprocess=force_preprocess) ) def load_sampling_graph(test_dir, processed_dataset): return torch.load( os.path.join(test_dir, processed_dataset["graph_topology"]["path"]), weights_only=False, ) def test_OnDiskDataset_TVTSet_exceptions(): """Test excpetions thrown when parsing TVTSet.""" with tempfile.TemporaryDirectory() as test_dir: # Case 1: ``format`` is invalid. yaml_content = """ tasks: - name: node_classification train_set: - type: paper data: - format: torch_invalid path: set/paper-train.pt """ write_yaml_file(yaml_content, test_dir) with pytest.raises(pydantic.ValidationError): _ = gb.OnDiskDataset(test_dir, force_preprocess=False).load() # Case 2: ``type`` is not specified while multiple TVT sets are # specified. yaml_content = """ tasks: - name: node_classification train_set: - type: null data: - format: numpy path: set/train.npy - type: null data: - format: numpy path: set/train.npy """ write_yaml_file(yaml_content, test_dir) with pytest.raises( AssertionError, match=r"Only one TVT set is allowed if type is not specified.", ): _ = gb.OnDiskDataset(test_dir, force_preprocess=False).load() def test_OnDiskDataset_multiple_tasks(): """Teset multiple tasks are supported.""" with tempfile.TemporaryDirectory() as test_dir: train_ids = np.arange(1000) train_ids_path = os.path.join(test_dir, "train_ids.npy") np.save(train_ids_path, train_ids) 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) yaml_content = f""" tasks: - name: node_classification_1 num_classes: 10 train_set: - type: null data: - name: seeds format: numpy in_memory: true path: {train_ids_path} - name: labels format: numpy in_memory: true path: {train_labels_path} - format: numpy in_memory: true path: {train_labels_path} - name: node_classification_2 num_classes: 10 train_set: - type: null data: - name: seeds format: numpy in_memory: true path: {train_ids_path} - name: labels format: numpy in_memory: true path: {train_labels_path} - format: numpy in_memory: true path: {train_labels_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) assert len(dataset.tasks) == 2 for task_id in range(2): assert ( dataset.tasks[task_id].metadata["name"] == f"node_classification_{task_id + 1}" ) assert dataset.tasks[task_id].metadata["num_classes"] == 10 # Verify train set. train_set = dataset.tasks[task_id].train_set assert len(train_set) == 1000 assert isinstance(train_set, gb.ItemSet) for i, (id, label, _) in enumerate(train_set): assert id == train_ids[i] assert label == train_labels[i] assert train_set.names == ("seeds", "labels", None) train_set = None dataset = None def test_OnDiskDataset_TVTSet_ItemSet_names(): """Test TVTSet which returns ItemSet with IDs, labels and corresponding names.""" with tempfile.TemporaryDirectory() as test_dir: train_ids = np.arange(1000) train_ids_path = os.path.join(test_dir, "train_ids.npy") np.save(train_ids_path, train_ids) 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) yaml_content = f""" tasks: - name: node_classification num_classes: 10 train_set: - type: null data: - name: seeds format: numpy in_memory: true path: {train_ids_path} - name: labels format: numpy in_memory: true path: {train_labels_path} - format: numpy in_memory: true path: {train_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) == 1000 assert isinstance(train_set, gb.ItemSet) for i, (id, label, _) in enumerate(train_set): assert id == train_ids[i] assert label == train_labels[i] assert train_set.names == ("seeds", "labels", None) train_set = None def test_OnDiskDataset_TVTSet_HeteroItemSet_names(): """Test TVTSet which returns ItemSet with IDs, labels and corresponding names.""" with tempfile.TemporaryDirectory() as test_dir: train_ids = np.arange(1000) train_ids_path = os.path.join(test_dir, "train_ids.npy") np.save(train_ids_path, train_ids) 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) yaml_content = f""" tasks: - name: node_classification num_classes: 10 train_set: - type: "author:writes:paper" data: - name: seeds format: numpy in_memory: true path: {train_ids_path} - name: labels format: numpy in_memory: true path: {train_labels_path} - format: numpy in_memory: true path: {train_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) == 1000 assert isinstance(train_set, gb.HeteroItemSet) for i, item in enumerate(train_set): assert isinstance(item, dict) assert "author:writes:paper" in item id, label, _ = item["author:writes:paper"] assert id == train_ids[i] assert label == train_labels[i] assert train_set.names == ("seeds", "labels", None) train_set = None def test_OnDiskDataset_TVTSet_ItemSet_id_label(): """Test TVTSet which returns ItemSet with IDs and labels.""" with tempfile.TemporaryDirectory() as test_dir: train_ids = np.arange(1000) train_ids_path = os.path.join(test_dir, "train_ids.npy") np.save(train_ids_path, train_ids) 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_ids = np.arange(1000, 2000) validation_ids_path = os.path.join(test_dir, "validation_ids.npy") np.save(validation_ids_path, validation_ids) 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_ids = np.arange(2000, 3000) test_ids_path = os.path.join(test_dir, "test_ids.npy") np.save(test_ids_path, test_ids) 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) # Case 1: # all TVT sets are specified. # ``type`` is not specified or specified as ``null``. # ``in_memory`` could be ``true`` and ``false``. yaml_content = f""" tasks: - name: node_classification num_classes: 10 train_set: - type: null data: - name: seeds format: numpy in_memory: true path: {train_ids_path} - name: labels format: numpy in_memory: true path: {train_labels_path} validation_set: - data: - name: seeds format: numpy in_memory: true path: {validation_ids_path} - name: labels format: numpy in_memory: true path: {validation_labels_path} test_set: - type: null data: - name: seeds format: numpy in_memory: true path: {test_ids_path} - name: labels format: numpy in_memory: true path: {test_labels_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) # Verify tasks. assert len(dataset.tasks) == 1 assert dataset.tasks[0].metadata["name"] == "node_classification" assert dataset.tasks[0].metadata["num_classes"] == 10 # Verify train set. train_set = dataset.tasks[0].train_set assert len(train_set) == 1000 assert isinstance(train_set, gb.ItemSet) for i, (id, label) in enumerate(train_set): assert id == train_ids[i] assert label == train_labels[i] assert train_set.names == ("seeds", "labels") train_set = None # Verify validation set. validation_set = dataset.tasks[0].validation_set assert len(validation_set) == 1000 assert isinstance(validation_set, gb.ItemSet) for i, (id, label) in enumerate(validation_set): assert id == validation_ids[i] assert label == validation_labels[i] assert validation_set.names == ("seeds", "labels") validation_set = None # Verify test set. test_set = dataset.tasks[0].test_set assert len(test_set) == 1000 assert isinstance(test_set, gb.ItemSet) for i, (id, label) in enumerate(test_set): assert id == test_ids[i] assert label == test_labels[i] assert test_set.names == ("seeds", "labels") test_set = None dataset = None # Case 2: Some TVT sets are None. yaml_content = f""" tasks: - name: node_classification train_set: - type: null data: - format: numpy path: {train_ids_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) assert dataset.tasks[0].train_set is not None assert dataset.tasks[0].validation_set is None assert dataset.tasks[0].test_set is None dataset = None def test_OnDiskDataset_TVTSet_ItemSet_node_pairs_labels(): """Test TVTSet which returns ItemSet 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: link_prediction train_set: - type: null data: - name: seeds format: numpy in_memory: true path: {train_seeds_path} - name: labels format: numpy in_memory: true path: {train_labels_path} validation_set: - data: - name: seeds format: numpy in_memory: true path: {validation_seeds_path} - name: labels format: numpy in_memory: true path: {validation_labels_path} test_set: - type: null 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) == 1000 assert isinstance(train_set, gb.ItemSet) for i, (node_pair, label) in enumerate(train_set): assert node_pair[0] == train_seeds[i][0] assert node_pair[1] == train_seeds[i][1] assert label == train_labels[i] assert train_set.names == ("seeds", "labels") train_set = None # Verify validation set. validation_set = dataset.tasks[0].validation_set assert len(validation_set) == 1000 assert isinstance(validation_set, gb.ItemSet) for i, (node_pair, label) in enumerate(validation_set): assert node_pair[0] == validation_seeds[i][0] assert node_pair[1] == validation_seeds[i][1] assert label == validation_labels[i] assert validation_set.names == ("seeds", "labels") validation_set = None # Verify test set. test_set = dataset.tasks[0].test_set assert len(test_set) == 1000 assert isinstance(test_set, gb.ItemSet) for i, (node_pair, label) in enumerate(test_set): assert node_pair[0] == test_seeds[i][0] assert node_pair[1] == test_seeds[i][1] assert label == test_labels[i] assert test_set.names == ("seeds", "labels") test_set = None dataset = None def test_OnDiskDataset_TVTSet_ItemSet_node_pairs_labels_indexes(): """Test TVTSet which returns ItemSet with node pairs and negative ones.""" with tempfile.TemporaryDirectory() as test_dir: train_seeds = np.arange(2000).reshape(1000, 2) train_neg_dst = np.random.choice(1000 * 10, size=1000 * 10) train_neg_src = train_seeds[:, 0].repeat(10) train_neg_seeds = ( np.concatenate((train_neg_dst, train_neg_src)).reshape(2, -1).T ) train_seeds = np.concatenate((train_seeds, train_neg_seeds)) train_seeds_path = os.path.join(test_dir, "train_seeds.npy") np.save(train_seeds_path, train_seeds) train_labels = torch.empty(1000 * 11) train_labels[:1000] = 1 train_labels[1000:] = 0 train_labels_path = os.path.join(test_dir, "train_labels.pt") torch.save(train_labels, train_labels_path) train_indexes = torch.arange(0, 1000) train_indexes = np.concatenate( (train_indexes, train_indexes.repeat_interleave(10)) ) train_indexes_path = os.path.join(test_dir, "train_indexes.pt") torch.save(train_indexes, train_indexes_path) validation_seeds = np.arange(2000, 4000).reshape(1000, 2) validation_neg_seeds = train_neg_seeds + 1 validation_seeds = np.concatenate( (validation_seeds, validation_neg_seeds) ) validation_seeds_path = os.path.join(test_dir, "validation_seeds.npy") np.save(validation_seeds_path, validation_seeds) validation_labels = train_labels validation_labels_path = os.path.join(test_dir, "validation_labels.pt") torch.save(validation_labels, validation_labels_path) validation_indexes = train_indexes validation_indexes_path = os.path.join( test_dir, "validation_indexes.pt" ) torch.save(validation_indexes, validation_indexes_path) test_seeds = np.arange(4000, 6000).reshape(1000, 2) test_neg_seeds = train_neg_seeds + 2 test_seeds = np.concatenate((test_seeds, test_neg_seeds)) test_seeds_path = os.path.join(test_dir, "test_seeds.npy") np.save(test_seeds_path, test_seeds) test_labels = train_labels test_labels_path = os.path.join(test_dir, "test_labels.pt") torch.save(test_labels, test_labels_path) test_indexes = train_indexes test_indexes_path = os.path.join(test_dir, "test_indexes.pt") torch.save(test_indexes, test_indexes_path) yaml_content = f""" tasks: - name: link_prediction train_set: - type: null data: - name: seeds format: numpy in_memory: true path: {train_seeds_path} - name: labels format: torch in_memory: true path: {train_labels_path} - name: indexes format: torch in_memory: true path: {train_indexes_path} validation_set: - data: - name: seeds format: numpy in_memory: true path: {validation_seeds_path} - name: labels format: torch in_memory: true path: {validation_labels_path} - name: indexes format: torch in_memory: true path: {validation_indexes_path} test_set: - type: null data: - name: seeds format: numpy in_memory: true path: {test_seeds_path} - name: labels format: torch in_memory: true path: {test_labels_path} - name: indexes format: torch in_memory: true path: {test_indexes_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) == 1000 * 11 assert isinstance(train_set, gb.ItemSet) for i, (node_pair, label, index) in enumerate(train_set): assert node_pair[0] == train_seeds[i][0] assert node_pair[1] == train_seeds[i][1] assert label == train_labels[i] assert index == train_indexes[i] assert train_set.names == ("seeds", "labels", "indexes") train_set = None # Verify validation set. validation_set = dataset.tasks[0].validation_set assert len(validation_set) == 1000 * 11 assert isinstance(validation_set, gb.ItemSet) for i, (node_pair, label, index) in enumerate(validation_set): assert node_pair[0] == validation_seeds[i][0] assert node_pair[1] == validation_seeds[i][1] assert label == validation_labels[i] assert index == validation_indexes[i] assert validation_set.names == ("seeds", "labels", "indexes") validation_set = None # Verify test set. test_set = dataset.tasks[0].test_set assert len(test_set) == 1000 * 11 assert isinstance(test_set, gb.ItemSet) for i, (node_pair, label, index) in enumerate(test_set): assert node_pair[0] == test_seeds[i][0] assert label == test_labels[i] assert index == test_indexes[i] assert test_set.names == ("seeds", "labels", "indexes") test_set = None dataset = None def test_OnDiskDataset_TVTSet_HeteroItemSet_id_label(): """Test TVTSet which returns HeteroItemSet with IDs and labels.""" with tempfile.TemporaryDirectory() as test_dir: train_ids = np.arange(1000) train_labels = np.random.randint(0, 10, size=1000) train_data = np.vstack([train_ids, train_labels]).T train_path = os.path.join(test_dir, "train.npy") np.save(train_path, train_data) validation_ids = np.arange(1000, 2000) validation_labels = np.random.randint(0, 10, size=1000) validation_data = np.vstack([validation_ids, validation_labels]).T validation_path = os.path.join(test_dir, "validation.npy") np.save(validation_path, validation_data) test_ids = np.arange(2000, 3000) test_labels = np.random.randint(0, 10, size=1000) test_data = np.vstack([test_ids, test_labels]).T test_path = os.path.join(test_dir, "test.npy") np.save(test_path, test_data) yaml_content = f""" tasks: - name: node_classification train_set: - type: paper data: - name: seeds format: numpy in_memory: true path: {train_path} - type: author data: - name: seeds format: numpy path: {train_path} validation_set: - type: paper data: - name: seeds format: numpy path: {validation_path} - type: author data: - name: seeds format: numpy path: {validation_path} test_set: - type: paper data: - name: seeds format: numpy in_memory: false path: {test_path} - type: author data: - name: seeds format: numpy path: {test_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", "author"] id, label = item[key] assert id == train_ids[i % 1000] assert label == train_labels[i % 1000] assert train_set.names == ("seeds",) 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", "author"] id, label = item[key] assert id == validation_ids[i % 1000] assert label == validation_labels[i % 1000] assert validation_set.names == ("seeds",) 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", "author"] id, label = item[key] assert id == test_ids[i % 1000] assert label == test_labels[i % 1000] assert test_set.names == ("seeds",) test_set = None dataset = None 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 def test_OnDiskDataset_Feature_heterograph(): """Test Feature storage.""" with tempfile.TemporaryDirectory() as test_dir: # Generate node data. node_data_paper = np.random.rand(1000, 10) node_data_paper_path = os.path.join(test_dir, "node_data_paper.npy") np.save(node_data_paper_path, node_data_paper) node_data_label = torch.tensor( [[random.randint(0, 10)] for _ in range(1000)] ) node_data_label_path = os.path.join(test_dir, "node_data_label.npy") np.save(node_data_label_path, node_data_label) # Generate edge data. edge_data_writes = np.random.rand(1000, 10) edge_data_writes_path = os.path.join(test_dir, "edge_writes_paper.npy") np.save(edge_data_writes_path, edge_data_writes) edge_data_label = torch.tensor( [[random.randint(0, 10)] for _ in range(1000)] ) edge_data_label_path = os.path.join(test_dir, "edge_data_label.npy") np.save(edge_data_label_path, edge_data_label) # Generate YAML. yaml_content = f""" feature_data: - domain: node type: paper name: feat format: numpy in_memory: false path: {node_data_paper_path} num_categories: 10 - domain: node type: paper name: labels format: numpy in_memory: true path: {node_data_label_path} - domain: edge type: "author:writes:paper" name: feat format: numpy in_memory: false path: {edge_data_writes_path} num_categories: 10 - domain: edge type: "author:writes:paper" name: labels format: numpy in_memory: true path: {edge_data_label_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) # Verify feature data storage. feature_data = dataset.feature assert len(feature_data) == 4 # Verify node feature data. assert torch.equal( feature_data.read("node", "paper", "feat"), torch.tensor(node_data_paper), ) assert ( feature_data.metadata("node", "paper", "feat")["num_categories"] == 10 ) assert torch.equal( feature_data.read("node", "paper", "labels"), node_data_label.clone().detach(), ) assert len(feature_data.metadata("node", "paper", "labels")) == 0 # Verify edge feature data. assert torch.equal( feature_data.read("edge", "author:writes:paper", "feat"), torch.tensor(edge_data_writes), ) assert ( feature_data.metadata("edge", "author:writes:paper", "feat")[ "num_categories" ] == 10 ) assert torch.equal( feature_data.read("edge", "author:writes:paper", "labels"), edge_data_label.clone().detach(), ) assert ( len(feature_data.metadata("edge", "author:writes:paper", "labels")) == 0 ) feature_data = None dataset = None def test_OnDiskDataset_Feature_homograph(): """Test Feature storage.""" with tempfile.TemporaryDirectory() as test_dir: # Generate node data. node_data_feat = np.random.rand(1000, 10) node_data_feat_path = os.path.join(test_dir, "node_data_feat.npy") np.save(node_data_feat_path, node_data_feat) node_data_label = torch.tensor( [[random.randint(0, 10)] for _ in range(1000)] ) node_data_label_path = os.path.join(test_dir, "node_data_label.npy") np.save(node_data_label_path, node_data_label) # Generate edge data. edge_data_feat = np.random.rand(1000, 10) edge_data_feat_path = os.path.join(test_dir, "edge_data_feat.npy") np.save(edge_data_feat_path, edge_data_feat) edge_data_label = torch.tensor( [[random.randint(0, 10)] for _ in range(1000)] ) edge_data_label_path = os.path.join(test_dir, "edge_data_label.npy") np.save(edge_data_label_path, edge_data_label) # Generate YAML. # ``type`` is not specified in the YAML. yaml_content = f""" feature_data: - domain: node name: feat format: numpy in_memory: false path: {node_data_feat_path} num_categories: 10 - domain: node name: labels format: numpy in_memory: true path: {node_data_label_path} - domain: edge name: feat format: numpy in_memory: false path: {edge_data_feat_path} num_categories: 10 - domain: edge name: labels format: numpy in_memory: true path: {edge_data_label_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) # Verify feature data storage. feature_data = dataset.feature assert len(feature_data) == 4 # Verify node feature data. assert torch.equal( feature_data.read("node", None, "feat"), torch.tensor(node_data_feat), ) assert ( feature_data.metadata("node", None, "feat")["num_categories"] == 10 ) assert torch.equal( feature_data.read("node", None, "labels"), node_data_label.clone().detach(), ) assert len(feature_data.metadata("node", None, "labels")) == 0 # Verify edge feature data. assert torch.equal( feature_data.read("edge", None, "feat"), torch.tensor(edge_data_feat), ) assert ( feature_data.metadata("edge", None, "feat")["num_categories"] == 10 ) assert torch.equal( feature_data.read("edge", None, "labels"), edge_data_label.clone().detach(), ) assert len(feature_data.metadata("edge", None, "labels")) == 0 feature_data = None dataset = None def test_OnDiskDataset_Graph_Exceptions(): """Test exceptions in parsing graph topology.""" with tempfile.TemporaryDirectory() as test_dir: # Invalid graph type. yaml_content = """ graph_topology: type: CSRSamplingGraph path: /path/to/graph """ write_yaml_file(yaml_content, test_dir) with pytest.raises( pydantic.ValidationError, match="1 validation error for OnDiskMetaData", ): _ = gb.OnDiskDataset(test_dir, force_preprocess=False).load() def test_OnDiskDataset_Graph_homogeneous(): """Test homogeneous graph topology.""" csc_indptr, indices = gbt.random_homo_graph(1000, 10 * 1000) graph = gb.fused_csc_sampling_graph(csc_indptr, indices) with tempfile.TemporaryDirectory() as test_dir: graph_path = os.path.join(test_dir, "fused_csc_sampling_graph.pt") torch.save(graph, graph_path) yaml_content = f""" graph_topology: type: FusedCSCSamplingGraph path: {graph_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) graph2 = dataset.graph assert graph.total_num_nodes == graph2.total_num_nodes assert graph.total_num_edges == graph2.total_num_edges assert torch.equal(graph.csc_indptr, graph2.csc_indptr) assert torch.equal(graph.indices, graph2.indices) assert ( graph.node_type_offset is None and graph2.node_type_offset is None ) assert graph.type_per_edge is None and graph2.type_per_edge is None assert graph.node_type_to_id is None and graph2.node_type_to_id is None assert graph.edge_type_to_id is None and graph2.edge_type_to_id is None def test_OnDiskDataset_Graph_heterogeneous(): """Test heterogeneous graph topology.""" ( csc_indptr, indices, node_type_offset, type_per_edge, node_type_to_id, edge_type_to_id, ) = gbt.random_hetero_graph(1000, 10 * 1000, 3, 4) graph = gb.fused_csc_sampling_graph( csc_indptr, indices, node_type_offset=node_type_offset, type_per_edge=type_per_edge, node_type_to_id=node_type_to_id, edge_type_to_id=edge_type_to_id, ) with tempfile.TemporaryDirectory() as test_dir: graph_path = os.path.join(test_dir, "fused_csc_sampling_graph.pt") torch.save(graph, graph_path) yaml_content = f""" graph_topology: type: FusedCSCSamplingGraph path: {graph_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) graph2 = dataset.graph assert graph.total_num_nodes == graph2.total_num_nodes assert graph.total_num_edges == graph2.total_num_edges assert torch.equal(graph.csc_indptr, graph2.csc_indptr) assert torch.equal(graph.indices, graph2.indices) assert torch.equal(graph.node_type_offset, graph2.node_type_offset) assert torch.equal(graph.type_per_edge, graph2.type_per_edge) assert graph.node_type_to_id == graph2.node_type_to_id assert graph.edge_type_to_id == graph2.edge_type_to_id def test_OnDiskDataset_Metadata(): """Test metadata of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" yaml_content = f""" dataset_name: {dataset_name} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) assert dataset.dataset_name == dataset_name # Only dataset_name is specified. yaml_content = f""" dataset_name: {dataset_name} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) assert dataset.dataset_name == dataset_name @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) def test_OnDiskDataset_preprocess_homogeneous(edge_fmt): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) output_file = gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False ) with open(output_file, "rb") as f: processed_dataset = yaml.load(f, Loader=yaml.Loader) assert processed_dataset["dataset_name"] == dataset_name assert processed_dataset["tasks"][0]["num_classes"] == num_classes assert "graph" not in processed_dataset assert "graph_topology" in processed_dataset fused_csc_sampling_graph = load_sampling_graph( test_dir, processed_dataset ) assert fused_csc_sampling_graph.total_num_nodes == num_nodes assert fused_csc_sampling_graph.total_num_edges == num_edges assert ( fused_csc_sampling_graph.node_attributes is not None and "feat" in fused_csc_sampling_graph.node_attributes ) assert ( fused_csc_sampling_graph.edge_attributes is not None and gb.ORIGINAL_EDGE_ID not in fused_csc_sampling_graph.edge_attributes and "feat" in fused_csc_sampling_graph.edge_attributes ) num_samples = 100 fanout = 1 subgraph = fused_csc_sampling_graph.sample_neighbors( torch.arange( 0, num_samples, dtype=fused_csc_sampling_graph.indices.dtype, ), torch.tensor([fanout]), ) assert len(subgraph.sampled_csc.indices) <= num_samples with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Test generating original_edge_id. output_file = gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=True ) with open(output_file, "rb") as f: processed_dataset = yaml.load(f, Loader=yaml.Loader) fused_csc_sampling_graph = load_sampling_graph( test_dir, processed_dataset ) assert ( fused_csc_sampling_graph.edge_attributes is not None and gb.ORIGINAL_EDGE_ID in fused_csc_sampling_graph.edge_attributes ) fused_csc_sampling_graph = None @pytest.mark.parametrize("auto_cast", [False, True]) def test_OnDiskDataset_preprocess_homogeneous_hardcode( auto_cast, edge_fmt="numpy" ): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: """Original graph in COO: 0 1 1 0 0 0 0 1 1 0 0 0 0 1 1 1 0 0 0 1 1 1 0 0 0 node_feats: [0.0, 1.9, 2.8, 3.7, 4.6] edge_feats: [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9] """ dataset_name = "graphbolt_test" num_nodes = 5 num_edges = 10 num_classes = 1 # Generate edges. edges = np.array( [[0, 0, 1, 1, 2, 2, 3, 3, 4, 4], [1, 2, 2, 3, 3, 4, 4, 0, 0, 1]], dtype=np.int64, ).T os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True) edges = edges.T edge_path = os.path.join("edges", "edge.npy") np.save(os.path.join(test_dir, edge_path), edges) # Generate graph edge-feats. edge_feats = np.array( [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=np.float64, ) os.makedirs(os.path.join(test_dir, "data"), exist_ok=True) edge_feat_path = os.path.join("data", "edge-feat.npy") np.save(os.path.join(test_dir, edge_feat_path), edge_feats) # Generate node-feats. node_feats = np.array( [0.0, 1.9, 2.8, 3.7, 4.6], dtype=np.float64, ) node_feat_path = os.path.join("data", "node-feat.npy") np.save(os.path.join(test_dir, node_feat_path), node_feats) # Generate train/test/valid set. os.makedirs(os.path.join(test_dir, "set"), exist_ok=True) train_data = np.array([0, 1, 2, 3, 4]) train_path = os.path.join("set", "train.npy") np.save(os.path.join(test_dir, train_path), train_data) valid_data = np.array([0, 1, 2, 3, 4]) valid_path = os.path.join("set", "valid.npy") np.save(os.path.join(test_dir, valid_path), valid_data) test_data = np.array([0, 1, 2, 3, 4]) test_path = os.path.join("set", "test.npy") np.save(os.path.join(test_dir, test_path), test_data) yaml_content = ( f"dataset_name: {dataset_name}\n" f"graph:\n" f" nodes:\n" f" - num: {num_nodes}\n" f" edges:\n" f" - format: {edge_fmt}\n" f" path: {edge_path}\n" f" feature_data:\n" f" - domain: node\n" f" type: null\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {node_feat_path}\n" f" - domain: edge\n" f" type: null\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {edge_feat_path}\n" f"feature_data:\n" f" - domain: node\n" f" type: null\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {node_feat_path}\n" f" - domain: edge\n" f" type: null\n" f" name: feat\n" f" format: numpy\n" f" path: {edge_feat_path}\n" f"tasks:\n" f" - name: node_classification\n" f" num_classes: {num_classes}\n" f" train_set:\n" f" - type: null\n" f" data:\n" f" - name: seeds\n" f" format: numpy\n" f" in_memory: true\n" f" path: {train_path}\n" f" validation_set:\n" f" - type: null\n" f" data:\n" f" - name: seeds\n" f" format: numpy\n" f" in_memory: true\n" f" path: {valid_path}\n" f" test_set:\n" f" - type: null\n" f" data:\n" f" - name: seeds\n" f" format: numpy\n" f" in_memory: true\n" f" path: {test_path}\n" ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) output_file = gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=True, auto_cast_to_optimal_dtype=auto_cast, ) with open(output_file, "rb") as f: processed_dataset = yaml.load(f, Loader=yaml.Loader) assert processed_dataset["dataset_name"] == dataset_name assert processed_dataset["tasks"][0]["num_classes"] == num_classes assert "graph" not in processed_dataset assert "graph_topology" in processed_dataset fused_csc_sampling_graph = load_sampling_graph( test_dir, processed_dataset ) assert fused_csc_sampling_graph.total_num_nodes == num_nodes assert fused_csc_sampling_graph.total_num_edges == num_edges assert torch.equal( fused_csc_sampling_graph.csc_indptr, torch.tensor([0, 2, 4, 6, 8, 10]), ) assert torch.equal( fused_csc_sampling_graph.indices, torch.tensor([3, 4, 0, 4, 0, 1, 1, 2, 2, 3]), ) assert torch.equal( fused_csc_sampling_graph.node_attributes["feat"], torch.tensor([0.0, 1.9, 2.8, 3.7, 4.6], dtype=torch.float64), ) assert torch.equal( fused_csc_sampling_graph.edge_attributes["feat"], torch.tensor( [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9], dtype=torch.float64, ), ) assert torch.equal( fused_csc_sampling_graph.edge_attributes[gb.ORIGINAL_EDGE_ID], torch.tensor([7, 8, 0, 9, 1, 2, 3, 4, 5, 6]), ) expected_dtype = torch.int32 if auto_cast else torch.int64 assert fused_csc_sampling_graph.csc_indptr.dtype == expected_dtype assert fused_csc_sampling_graph.indices.dtype == expected_dtype assert ( fused_csc_sampling_graph.edge_attributes[gb.ORIGINAL_EDGE_ID].dtype == expected_dtype ) num_samples = 5 fanout = 1 subgraph = fused_csc_sampling_graph.sample_neighbors( torch.arange( 0, num_samples, dtype=fused_csc_sampling_graph.indices.dtype, ), torch.tensor([fanout]), ) assert len(subgraph.sampled_csc.indices) <= num_samples @pytest.mark.parametrize("auto_cast", [False, True]) def test_OnDiskDataset_preprocess_heterogeneous_hardcode( auto_cast, edge_fmt="numpy" ): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: """Original graph in COO: 0 1 1 0 0 0 0 1 1 0 0 0 0 1 1 1 0 0 0 1 1 1 0 0 0 node_type_0: [0, 1] node_type_1: [2, 3, 4] edge_type_0: node_type_0 -> node_type_0 edge_type_1: node_type_0 -> node_type_1 edge_type_2: node_type_1 -> node_type_1 edge_type_3: node_type_1 -> node_type_0 node_feats: [0.0, 1.9, 2.8, 3.7, 4.6] edge_feats: [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9] """ dataset_name = "graphbolt_test" num_nodes = { "A": 2, "B": 3, } num_edges = { ("A", "a_a", "A"): 1, ("A", "a_b", "B"): 3, ("B", "b_b", "A"): 3, ("B", "b_a", "B"): 3, } num_classes = 1 # Generate edges. os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True) np.save( os.path.join(test_dir, "edges", "a_a.npy"), np.array([[0], [1]], dtype=np.int64), ) np.save( os.path.join(test_dir, "edges", "a_b.npy"), np.array([[0, 1, 1], [0, 0, 1]], dtype=np.int64), ) np.save( os.path.join(test_dir, "edges", "b_b.npy"), np.array([[0, 0, 1], [1, 2, 2]], dtype=np.int64), ) np.save( os.path.join(test_dir, "edges", "b_a.npy"), np.array([[1, 2, 2], [0, 0, 1]], dtype=np.int64), ) # Generate node features. os.makedirs(os.path.join(test_dir, "data"), exist_ok=True) np.save( os.path.join(test_dir, "data", "A-feat.npy"), np.array([0.0, 1.9], dtype=np.float64), ) np.save( os.path.join(test_dir, "data", "B-feat.npy"), np.array([2.8, 3.7, 4.6], dtype=np.float64), ) # Generate edge features. os.makedirs(os.path.join(test_dir, "data"), exist_ok=True) np.save( os.path.join(test_dir, "data", "a_a-feat.npy"), np.array([0.0], dtype=np.float64), ) np.save( os.path.join(test_dir, "data", "a_b-feat.npy"), np.array([1.1, 2.2, 3.3], dtype=np.float64), ) np.save( os.path.join(test_dir, "data", "b_b-feat.npy"), np.array([4.4, 5.5, 6.6], dtype=np.float64), ) np.save( os.path.join(test_dir, "data", "b_a-feat.npy"), np.array([7.7, 8.8, 9.9], dtype=np.float64), ) yaml_content = ( f"dataset_name: {dataset_name}\n" f"graph:\n" f" nodes:\n" f" - type: A\n" f" num: 2\n" f" - type: B\n" f" num: 3\n" f" edges:\n" f" - type: A:a_a:A\n" f" format: {edge_fmt}\n" f" path: {os.path.join('edges', 'a_a.npy')}\n" f" - type: A:a_b:B\n" f" format: {edge_fmt}\n" f" path: {os.path.join('edges', 'a_b.npy')}\n" f" - type: B:b_b:B\n" f" format: {edge_fmt}\n" f" path: {os.path.join('edges', 'b_b.npy')}\n" f" - type: B:b_a:A\n" f" format: {edge_fmt}\n" f" path: {os.path.join('edges', 'b_a.npy')}\n" f" feature_data:\n" f" - domain: node\n" f" type: A\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {os.path.join(test_dir, 'data', 'A-feat.npy')}\n" f" - domain: node\n" f" type: B\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {os.path.join(test_dir, 'data', 'B-feat.npy')}\n" f" - domain: edge\n" f" type: A:a_a:A\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {os.path.join(test_dir, 'data', 'a_a-feat.npy')}\n" f" - domain: edge\n" f" type: A:a_b:B\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {os.path.join(test_dir, 'data', 'a_b-feat.npy')}\n" f" - domain: edge\n" f" type: B:b_b:B\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {os.path.join(test_dir, 'data', 'b_b-feat.npy')}\n" f" - domain: edge\n" f" type: B:b_a:A\n" f" name: feat\n" f" format: numpy\n" f" in_memory: true\n" f" path: {os.path.join(test_dir, 'data', 'b_a-feat.npy')}\n" ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) output_file = gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=True, auto_cast_to_optimal_dtype=auto_cast, ) with open(output_file, "rb") as f: processed_dataset = yaml.load(f, Loader=yaml.Loader) assert processed_dataset["dataset_name"] == dataset_name assert "graph" not in processed_dataset assert "graph_topology" in processed_dataset fused_csc_sampling_graph = load_sampling_graph( test_dir, processed_dataset ) assert fused_csc_sampling_graph.total_num_nodes == 5 assert fused_csc_sampling_graph.total_num_edges == 10 assert torch.equal( fused_csc_sampling_graph.csc_indptr, torch.tensor([0, 2, 4, 6, 8, 10]), ) assert torch.equal( fused_csc_sampling_graph.indices, torch.tensor([3, 4, 0, 4, 0, 1, 1, 2, 2, 3]), ) assert torch.equal( fused_csc_sampling_graph.node_attributes["feat"], torch.tensor([0.0, 1.9, 2.8, 3.7, 4.6], dtype=torch.float64), ) assert torch.equal( fused_csc_sampling_graph.edge_attributes["feat"], torch.tensor( [0.0, 1.1, 2.2, 3.3, 7.7, 8.8, 9.9, 4.4, 5.5, 6.6], dtype=torch.float64, ), ) assert torch.equal( fused_csc_sampling_graph.type_per_edge, torch.tensor([2, 2, 0, 2, 1, 1, 1, 3, 3, 3]), ) assert torch.equal( fused_csc_sampling_graph.edge_attributes[gb.ORIGINAL_EDGE_ID], torch.tensor([0, 1, 0, 2, 0, 1, 2, 0, 1, 2]), ) expected_dtype = torch.int32 if auto_cast else torch.int64 assert fused_csc_sampling_graph.csc_indptr.dtype == expected_dtype assert fused_csc_sampling_graph.indices.dtype == expected_dtype assert ( fused_csc_sampling_graph.edge_attributes[gb.ORIGINAL_EDGE_ID].dtype == expected_dtype ) assert fused_csc_sampling_graph.node_type_offset.dtype == expected_dtype expected_etype_dtype = torch.uint8 if auto_cast else torch.int64 assert ( fused_csc_sampling_graph.type_per_edge.dtype == expected_etype_dtype ) def test_OnDiskDataset_preprocess_path(): """Test if the preprocess function can catch the path error.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" yaml_content = f""" dataset_name: {dataset_name} """ yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Case1. Test the passed in is the yaml file path. with pytest.raises( RuntimeError, match="The dataset must be a directory. " rf"But got {re.escape(yaml_file)}", ): _ = gb.OnDiskDataset(yaml_file) # Case2. Test the passed in is a fake directory. fake_dir = os.path.join(test_dir, "fake_dir") with pytest.raises( RuntimeError, match=rf"Invalid dataset path: {re.escape(fake_dir)}", ): _ = gb.OnDiskDataset(fake_dir) # Case3. Test the passed in is the dataset directory. # But the metadata.yaml is not in the directory. os.makedirs(os.path.join(test_dir, "fake_dir"), exist_ok=True) with pytest.raises( RuntimeError, match=r"metadata.yaml does not exist.", ): _ = gb.OnDiskDataset(fake_dir) def test_OnDiskDataset_preprocess_yaml_content(): """Test if the preprocessed metadata.yaml is correct.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random edges. nodes = np.repeat(np.arange(num_nodes), 5) neighbors = np.random.randint(0, num_nodes, size=(num_edges)) edges = np.stack([nodes, neighbors], axis=1) # Write into edges/edge.csv os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True) edges = pd.DataFrame(edges, columns=["src", "dst"]) edge_path = os.path.join("edges", "edge.csv") edges.to_csv( os.path.join(test_dir, edge_path), index=False, header=False, ) # Generate random graph edge-feats. edge_feats = np.random.rand(num_edges, 5) os.makedirs(os.path.join(test_dir, "data"), exist_ok=True) feature_edge = os.path.join("data", "edge-feat.npy") np.save(os.path.join(test_dir, feature_edge), edge_feats) # Generate random node-feats. node_feats = np.random.rand(num_nodes, 10) feature_node = os.path.join("data", "node-feat.npy") np.save(os.path.join(test_dir, feature_node), node_feats) # Generate train/test/valid set. os.makedirs(os.path.join(test_dir, "set"), exist_ok=True) train_pairs = (np.arange(1000), np.arange(1000, 2000)) train_labels = np.random.randint(0, 10, size=1000) train_data = np.vstack([train_pairs, train_labels]).T train_path = os.path.join("set", "train.npy") np.save(os.path.join(test_dir, train_path), train_data) validation_pairs = (np.arange(1000, 2000), np.arange(2000, 3000)) validation_labels = np.random.randint(0, 10, size=1000) validation_data = np.vstack([validation_pairs, validation_labels]).T validation_path = os.path.join("set", "validation.npy") np.save(os.path.join(test_dir, validation_path), validation_data) test_pairs = (np.arange(2000, 3000), np.arange(3000, 4000)) test_labels = np.random.randint(0, 10, size=1000) test_data = np.vstack([test_pairs, test_labels]).T test_path = os.path.join("set", "test.npy") np.save(os.path.join(test_dir, test_path), test_data) yaml_content = f""" dataset_name: {dataset_name} graph: # graph structure and required attributes. nodes: - num: {num_nodes} edges: - format: csv path: {edge_path} feature_data: - domain: edge type: null name: feat format: numpy in_memory: true path: {feature_edge} feature_data: - domain: node type: null name: feat format: numpy in_memory: false path: {feature_node} tasks: - name: node_classification num_classes: {num_classes} train_set: - type: null data: - format: numpy path: {train_path} validation_set: - type: null data: - format: numpy path: {validation_path} test_set: - type: null data: - format: numpy path: {test_path} """ yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) preprocessed_metadata_path = gb.preprocess_ondisk_dataset(test_dir) with open(preprocessed_metadata_path, "r") as f: yaml_data = yaml.safe_load(f) topo_path = os.path.join("preprocessed", "fused_csc_sampling_graph.pt") target_yaml_content = f""" dataset_name: {dataset_name} graph_topology: type: FusedCSCSamplingGraph path: {topo_path} feature_data: - domain: node type: null name: feat format: numpy in_memory: false path: {os.path.join("preprocessed", feature_node)} tasks: - name: node_classification num_classes: {num_classes} train_set: - type: null data: - format: numpy path: {os.path.join("preprocessed", train_path)} validation_set: - type: null data: - format: numpy path: {os.path.join("preprocessed", validation_path)} test_set: - type: null data: - format: numpy path: {os.path.join("preprocessed", test_path)} include_original_edge_id: False """ target_yaml_data = yaml.safe_load(target_yaml_content) # Check yaml content. assert ( yaml_data == target_yaml_data ), "The preprocessed metadata.yaml is not correct." # Check file existence. assert os.path.exists( os.path.join(test_dir, yaml_data["graph_topology"]["path"]) ) assert os.path.exists( os.path.join(test_dir, yaml_data["feature_data"][0]["path"]) ) for set_name in ["train_set", "validation_set", "test_set"]: assert os.path.exists( os.path.join( test_dir, yaml_data["tasks"][0][set_name][0]["data"][0]["path"], ) ) def test_OnDiskDataset_preprocess_force_preprocess(capsys): """Test force preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # First preprocess on-disk dataset. preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False, force_preprocess=False ) ) captured = capsys.readouterr().out.split("\n") assert captured == [ "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["tasks"][0]["name"] == "link_prediction" # Change yaml_data, but do not force preprocess on-disk dataset. with open(yaml_file, "r") as f: yaml_data = yaml.safe_load(f) yaml_data["tasks"][0]["name"] = "fake_name" with open(yaml_file, "w") as f: yaml.dump(yaml_data, f) preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False, force_preprocess=False ) ) captured = capsys.readouterr().out.split("\n") assert captured == ["The dataset is already preprocessed.", ""] with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["tasks"][0]["name"] == "link_prediction" # Force preprocess on-disk dataset. preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False, force_preprocess=True ) ) captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["tasks"][0]["name"] == "fake_name" def test_OnDiskDataset_preprocess_auto_force_preprocess(capsys): """Test force preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # First preprocess on-disk dataset. preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False ) ) captured = capsys.readouterr().out.split("\n") assert captured == [ "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["tasks"][0]["name"] == "link_prediction" # 1. Change yaml_data. with open(yaml_file, "r") as f: yaml_data = yaml.safe_load(f) yaml_data["tasks"][0]["name"] = "fake_name" with open(yaml_file, "w") as f: yaml.dump(yaml_data, f) preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False ) ) captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["tasks"][0]["name"] == "fake_name" # 2. Change edge feature. edge_feats = np.random.rand(num_edges, num_classes) edge_feat_path = os.path.join("data", "edge-feat.npy") np.save(os.path.join(test_dir, edge_feat_path), edge_feats) preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False ) ) captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] preprocessed_edge_feat = np.load( os.path.join(test_dir, "preprocessed", edge_feat_path) ) assert preprocessed_edge_feat.all() == edge_feats.all() with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["include_original_edge_id"] == False # 3. Change include_original_edge_id. preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=True ) ) captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] with open(preprocessed_metadata_path, "r") as f: target_yaml_data = yaml.safe_load(f) assert target_yaml_data["include_original_edge_id"] == True # 4. Change nothing. preprocessed_metadata_path = ( gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=True ) ) captured = capsys.readouterr().out.split("\n") assert captured == ["The dataset is already preprocessed.", ""] def test_OnDiskDataset_preprocess_not_include_eids(): with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) with pytest.warns( GBWarning, match="Edge feature is stored, but edge IDs are not saved.", ): gb.ondisk_dataset.preprocess_ondisk_dataset( test_dir, include_original_edge_id=False ) @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) def test_OnDiskDataset_load_name(edge_fmt): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Check modify `dataset_name` field. dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["dataset_name"] = "fake_name" dataset.load() assert dataset.dataset_name == "fake_name" dataset = None @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) def test_OnDiskDataset_load_feature(edge_fmt): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Case1. Test modify the `in_memory` field. dataset = gb.OnDiskDataset(test_dir).load() original_feature_data = dataset.feature dataset.yaml_data["feature_data"][0]["in_memory"] = True load_dataset(dataset) modify_feature_data = dataset.feature # After modify the `in_memory` field, the feature data should be # equal. assert torch.equal( original_feature_data.read("node", None, "feat"), modify_feature_data.read("node", None, "feat"), ) # Case2. Test modify the `format` field. dataset = gb.OnDiskDataset(test_dir) # If `format` is torch and `in_memory` is False, it will # raise an AssertionError. dataset.yaml_data["feature_data"][0]["in_memory"] = False dataset.yaml_data["feature_data"][0]["format"] = "torch" with pytest.raises( AssertionError, match="^Pytorch tensor can only be loaded in memory,", ): load_dataset(dataset) dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["feature_data"][0]["in_memory"] = True dataset.yaml_data["feature_data"][0]["format"] = "torch" # If `format` is torch and `in_memory` is True, it will # raise an UnpicklingError. with pytest.raises(pickle.UnpicklingError): load_dataset(dataset) # Case3. Test modify the `path` field. dataset = gb.OnDiskDataset(test_dir) # Use invalid path will raise an FileNotFoundError. dataset.yaml_data["feature_data"][0]["path"] = "fake_path" with pytest.raises( FileNotFoundError, match=r"\[Errno 2\] No such file or directory:", ): load_dataset(dataset) # Modifying the `path` field to an absolute path should work. # In os.path.join, if a segment is an absolute path (which # on Windows requires both a drive and a root), then all # previous segments are ignored and joining continues from # the absolute path segment. dataset = load_dataset(gb.OnDiskDataset(test_dir)) original_feature_data = dataset.feature dataset.yaml_data["feature_data"][0]["path"] = os.path.join( test_dir, dataset.yaml_data["feature_data"][0]["path"] ) load_dataset(dataset) modify_feature_data = dataset.feature assert torch.equal( original_feature_data.read("node", None, "feat"), modify_feature_data.read("node", None, "feat"), ) original_feature_data = None modify_feature_data = None dataset = None @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) def test_OnDiskDataset_load_graph(edge_fmt): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Check the different original_edge_id option to load edge_attributes. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=True ).load() assert ( dataset.graph.edge_attributes is not None and gb.ORIGINAL_EDGE_ID in dataset.graph.edge_attributes ) # Case1. Test modify the `type` field. dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["graph_topology"]["type"] = "fake_type" with pytest.raises( pydantic.ValidationError, # As error message diffs in pydantic 1.x and 2.x, we just match # keyword only. match="'FusedCSCSamplingGraph'", ): dataset.load() # Case2. Test modify the `path` field. dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["graph_topology"]["path"] = "fake_path" with pytest.raises( FileNotFoundError, match=r"\[Errno 2\] No such file or directory:", ): dataset.load() # Modifying the `path` field to an absolute path should work. # In os.path.join, if a segment is an absolute path (which # on Windows requires both a drive and a root), then all # previous segments are ignored and joining continues from # the absolute path segment. dataset = gb.OnDiskDataset(test_dir).load() original_graph = dataset.graph dataset.yaml_data["graph_topology"]["path"] = os.path.join( test_dir, dataset.yaml_data["graph_topology"]["path"] ) dataset.load() modify_graph = dataset.graph assert torch.equal( original_graph.csc_indptr, modify_graph.csc_indptr, ) original_graph = None modify_graph = None dataset = None with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Test do not generate original_edge_id. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False ).load() assert ( dataset.graph.edge_attributes is None or gb.ORIGINAL_EDGE_ID not in dataset.graph.edge_attributes ) dataset = None @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) def test_OnDiskDataset_load_tasks(edge_fmt): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Case1. Test modify the `name` field. dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["tasks"][0]["name"] = "fake_name" dataset.load() assert dataset.tasks[0].metadata["name"] == "fake_name" # Case2. Test modify the `num_classes` field. dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["tasks"][0]["num_classes"] = 100 dataset.load() assert dataset.tasks[0].metadata["num_classes"] == 100 # Case3. Test modify the `format` field. dataset = gb.OnDiskDataset(test_dir) # Change the `format` field to torch. dataset.yaml_data["tasks"][0]["train_set"][0]["data"][0][ "format" ] = "torch" with pytest.raises(pickle.UnpicklingError): dataset.load() dataset = gb.OnDiskDataset(test_dir) dataset.yaml_data["tasks"][0]["train_set"][0]["data"][0][ "format" ] = "torch" # Change the `in_memory` field to False will also raise an # UnpicklingError. Unlike the case of testing `feature_data`. dataset.yaml_data["tasks"][0]["train_set"][0]["data"][0][ "in_memory" ] = False with pytest.raises(pickle.UnpicklingError): dataset.load() # Case4. Test modify the `path` field. dataset = gb.OnDiskDataset(test_dir) # Use invalid path will raise an FileNotFoundError. dataset.yaml_data["tasks"][0]["train_set"][0]["data"][0][ "path" ] = "fake_path" with pytest.raises( FileNotFoundError, match=r"\[Errno 2\] No such file or directory:", ): dataset.load() # Modifying the `path` field to an absolute path should work. # In os.path.join, if a segment is an absolute path (which # on Windows requires both a drive and a root), then all # previous segments are ignored and joining continues from # the absolute path segment. dataset = gb.OnDiskDataset(test_dir).load() original_train_set = dataset.tasks[0].train_set._items dataset.yaml_data["tasks"][0]["train_set"][0]["data"][0][ "path" ] = os.path.join( test_dir, dataset.yaml_data["tasks"][0]["train_set"][0]["data"][0]["path"], ) dataset.load() modify_train_set = dataset.tasks[0].train_set._items assert torch.equal( original_train_set[0], modify_train_set[0], ) original_train_set = None modify_train_set = None dataset = None def test_OnDiskDataset_all_nodes_set_homo(): """Test homograph's all nodes set of OnDiskDataset.""" csc_indptr, indices = gbt.random_homo_graph(1000, 10 * 1000) graph = gb.fused_csc_sampling_graph(csc_indptr, indices) with tempfile.TemporaryDirectory() as test_dir: graph_path = os.path.join(test_dir, "fused_csc_sampling_graph.pt") torch.save(graph, graph_path) yaml_content = f""" graph_topology: type: FusedCSCSamplingGraph path: {graph_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) all_nodes_set = dataset.all_nodes_set assert isinstance(all_nodes_set, gb.ItemSet) assert all_nodes_set.names == ("seeds",) for i, item in enumerate(all_nodes_set): assert i == item dataset = None def test_OnDiskDataset_all_nodes_set_hetero(): """Test heterograph's all nodes set of OnDiskDataset.""" ( csc_indptr, indices, node_type_offset, type_per_edge, node_type_to_id, edge_type_to_id, ) = gbt.random_hetero_graph(1000, 10 * 1000, 3, 4) graph = gb.fused_csc_sampling_graph( csc_indptr, indices, node_type_offset=node_type_offset, type_per_edge=type_per_edge, node_type_to_id=node_type_to_id, edge_type_to_id=edge_type_to_id, edge_attributes=None, ) with tempfile.TemporaryDirectory() as test_dir: graph_path = os.path.join(test_dir, "fused_csc_sampling_graph.pt") torch.save(graph, graph_path) yaml_content = f""" graph_topology: type: FusedCSCSamplingGraph path: {graph_path} """ dataset = write_yaml_and_load_dataset(yaml_content, test_dir) all_nodes_set = dataset.all_nodes_set assert isinstance(all_nodes_set, gb.HeteroItemSet) assert all_nodes_set.names == ("seeds",) for i, item in enumerate(all_nodes_set): assert len(item) == 1 assert isinstance(item, dict) dataset = None @pytest.mark.parametrize("fmt", ["numpy", "torch"]) def test_OnDiskDataset_load_1D_feature(fmt): with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4 num_edges = 20 num_classes = 1 type_name = "npy" if fmt == "numpy" else "pt" # Generate random edges. nodes = np.repeat(np.arange(num_nodes), 5) neighbors = np.random.randint(0, num_nodes, size=(num_edges)) edges = np.stack([nodes, neighbors], axis=1) # Write into edges/edge.csv os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True) edges = pd.DataFrame(edges, columns=["src", "dst"]) edge_path = os.path.join("edges", "edge.csv") edges.to_csv( os.path.join(test_dir, edge_path), index=False, header=False, ) # Generate random graph edge-feats. edge_feats = np.random.rand(num_edges, 5) os.makedirs(os.path.join(test_dir, "data"), exist_ok=True) edge_feat_path = os.path.join("data", f"edge-feat.{type_name}") # Generate random 1-D node-feats. node_feats = np.random.rand(num_nodes) node_feat_path = os.path.join("data", f"node-feat.{type_name}") assert node_feats.ndim == 1 # Generate 1-D train set. os.makedirs(os.path.join(test_dir, "set"), exist_ok=True) train_path = os.path.join("set", f"train.{type_name}") if fmt == "numpy": np.save(os.path.join(test_dir, edge_feat_path), edge_feats) np.save(os.path.join(test_dir, node_feat_path), node_feats) np.save(os.path.join(test_dir, train_path), np.array([0, 1, 0])) else: torch.save( torch.from_numpy(edge_feats), os.path.join(test_dir, edge_feat_path), ) torch.save( torch.from_numpy(node_feats), os.path.join(test_dir, node_feat_path), ) torch.save( torch.tensor([0, 1, 0]), os.path.join(test_dir, train_path) ) yaml_content = f""" dataset_name: {dataset_name} graph: # graph structure and required attributes. nodes: - num: {num_nodes} edges: - format: csv path: {edge_path} feature_data: - domain: edge type: null name: feat format: {fmt} in_memory: true path: {edge_feat_path} feature_data: - domain: node type: null name: feat format: {fmt} in_memory: false path: {node_feat_path} tasks: - name: node_classification num_classes: {num_classes} train_set: - type: null data: - format: {fmt} path: {train_path} """ yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) dataset = gb.OnDiskDataset(test_dir).load() feature = dataset.feature.read("node", None, "feat") # Test whether feature has changed. assert torch.equal(torch.from_numpy(node_feats.reshape(-1, 1)), feature) # Test whether itemsets keep same. assert torch.equal( dataset.tasks[0].train_set._items[0], torch.tensor([0, 1, 0]) ) dataset = None node_feats = None feature = None def test_BuiltinDataset(): """Test BuiltinDataset.""" with tempfile.TemporaryDirectory() as test_dir: # Case 1: download from DGL S3 storage. dataset_name = "test-dataset-231207" # Add dataset to the builtin dataset list for testing only. Due to we # add `seeds` suffix to datasets when downloading, so we append # dataset name with `-seeds` suffix here. gb.BuiltinDataset._all_datasets.append(dataset_name + "-seeds") dataset = gb.BuiltinDataset(name=dataset_name, root=test_dir).load() assert dataset.graph is not None assert dataset.feature is not None assert dataset.tasks is not None assert dataset.dataset_name == dataset_name # Case 2: dataset is already downloaded. dataset = gb.BuiltinDataset(name=dataset_name, root=test_dir).load() assert dataset.graph is not None assert dataset.feature is not None assert dataset.tasks is not None assert dataset.dataset_name == dataset_name dataset = None # Case 3: dataset is not available. dataset_name = "fake_name-seeds" with pytest.raises( RuntimeError, match=rf"Dataset {dataset_name} is not available.*", ): _ = gb.BuiltinDataset(name=dataset_name, root=test_dir).load() @pytest.mark.parametrize("auto_cast", [True, False]) @pytest.mark.parametrize("include_original_edge_id", [True, False]) @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) def test_OnDiskDataset_homogeneous( auto_cast, include_original_edge_id, edge_fmt ): """Preprocess and instantiate OnDiskDataset for homogeneous graph.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt=edge_fmt, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) 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 == num_nodes assert graph.total_num_edges == num_edges 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.ItemSet) assert isinstance(tasks[0].validation_set, gb.ItemSet) assert isinstance(tasks[0].test_set, gb.ItemSet) assert tasks[0].train_set._items[0].dtype == graph.indices.dtype assert tasks[0].validation_set._items[0].dtype == graph.indices.dtype assert tasks[0].test_set._items[0].dtype == graph.indices.dtype assert dataset.all_nodes_set._items.dtype == graph.indices.dtype assert tasks[0].metadata["num_classes"] == num_classes assert tasks[0].metadata["name"] == "link_prediction" assert dataset.feature.size("node", None, "feat")[0] == num_classes assert dataset.feature.size("edge", None, "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=["feat"] ) dataloader = gb.DataLoader(datapipe) for _ in dataloader: pass graph = None tasks = None dataset = None @pytest.mark.parametrize("auto_cast", [True, False]) @pytest.mark.parametrize("include_original_edge_id", [True, False]) @pytest.mark.parametrize("edge_fmt", ["csv", "numpy"]) 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 def test_OnDiskDataset_force_preprocess(capsys): """Test force preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # First preprocess on-disk dataset. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False, force_preprocess=False ).load() captured = capsys.readouterr().out.split("\n") assert captured == [ "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] tasks = dataset.tasks assert tasks[0].metadata["name"] == "link_prediction" # Change yaml_data, but do not force preprocess on-disk dataset. with open(yaml_file, "r") as f: yaml_data = yaml.safe_load(f) yaml_data["tasks"][0]["name"] = "fake_name" with open(yaml_file, "w") as f: yaml.dump(yaml_data, f) dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False, force_preprocess=False ).load() captured = capsys.readouterr().out.split("\n") assert captured == ["The dataset is already preprocessed.", ""] tasks = dataset.tasks assert tasks[0].metadata["name"] == "link_prediction" # Force preprocess on-disk dataset. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False, force_preprocess=True ).load() captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] tasks = dataset.tasks assert tasks[0].metadata["name"] == "fake_name" tasks = None dataset = None def test_OnDiskDataset_auto_force_preprocess(capsys): """Test force preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # First preprocess on-disk dataset. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False ).load() captured = capsys.readouterr().out.split("\n") assert captured == [ "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] tasks = dataset.tasks assert tasks[0].metadata["name"] == "link_prediction" # 1. Change yaml_data. with open(yaml_file, "r") as f: yaml_data = yaml.safe_load(f) yaml_data["tasks"][0]["name"] = "fake_name" with open(yaml_file, "w") as f: yaml.dump(yaml_data, f) dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False ).load() captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] tasks = dataset.tasks assert tasks[0].metadata["name"] == "fake_name" # 2. Change edge feature. edge_feats = np.random.rand(num_edges, num_classes) edge_feat_path = os.path.join("data", "edge-feat.npy") np.save(os.path.join(test_dir, edge_feat_path), edge_feats) dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=False ).load() captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] assert torch.equal( dataset.feature.read("edge", None, "feat"), torch.from_numpy(edge_feats), ) graph = dataset.graph assert gb.ORIGINAL_EDGE_ID not in graph.edge_attributes # 3. Change include_original_edge_id. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=True ).load() captured = capsys.readouterr().out.split("\n") assert captured == [ "The on-disk dataset is re-preprocessing, so the existing " + "preprocessed dataset has been removed.", "Start to preprocess the on-disk dataset.", "Finish preprocessing the on-disk dataset.", "", ] graph = dataset.graph assert gb.ORIGINAL_EDGE_ID in graph.edge_attributes # 4. Change Nothing. dataset = gb.OnDiskDataset( test_dir, include_original_edge_id=True ).load() captured = capsys.readouterr().out.split("\n") assert captured == ["The dataset is already preprocessed.", ""] graph = None tasks = None dataset = None def test_OnDiskTask_repr_homogeneous(): item_set = gb.ItemSet( (torch.arange(0, 5), torch.arange(5, 10)), names=("seeds", "labels"), ) metadata = {"name": "node_classification"} task = gb.OnDiskTask(metadata, item_set, item_set, item_set) expected_str = ( "OnDiskTask(validation_set=ItemSet(\n" " items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n" " names=('seeds', 'labels'),\n" " ),\n" " train_set=ItemSet(\n" " items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n" " names=('seeds', 'labels'),\n" " ),\n" " test_set=ItemSet(\n" " items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n" " names=('seeds', 'labels'),\n" " ),\n" " metadata={'name': 'node_classification'},)" ) assert repr(task) == expected_str, task def test_OnDiskDataset_not_include_eids(): with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) with pytest.warns( GBWarning, match="Edge feature is stored, but edge IDs are not saved.", ): gb.OnDiskDataset(test_dir, include_original_edge_id=False) def test_OnDiskTask_repr_heterogeneous(): item_set = gb.HeteroItemSet( { "user": gb.ItemSet(torch.arange(0, 5), names="seeds"), "item": gb.ItemSet(torch.arange(5, 10), names="seeds"), } ) metadata = {"name": "node_classification"} task = gb.OnDiskTask(metadata, item_set, item_set, item_set) expected_str = ( "OnDiskTask(validation_set=HeteroItemSet(\n" " itemsets={'user': ItemSet(\n" " items=(tensor([0, 1, 2, 3, 4]),),\n" " names=('seeds',),\n" " ), 'item': ItemSet(\n" " items=(tensor([5, 6, 7, 8, 9]),),\n" " names=('seeds',),\n" " )},\n" " names=('seeds',),\n" " ),\n" " train_set=HeteroItemSet(\n" " itemsets={'user': ItemSet(\n" " items=(tensor([0, 1, 2, 3, 4]),),\n" " names=('seeds',),\n" " ), 'item': ItemSet(\n" " items=(tensor([5, 6, 7, 8, 9]),),\n" " names=('seeds',),\n" " )},\n" " names=('seeds',),\n" " ),\n" " test_set=HeteroItemSet(\n" " itemsets={'user': ItemSet(\n" " items=(tensor([0, 1, 2, 3, 4]),),\n" " names=('seeds',),\n" " ), 'item': ItemSet(\n" " items=(tensor([5, 6, 7, 8, 9]),),\n" " names=('seeds',),\n" " )},\n" " names=('seeds',),\n" " ),\n" " metadata={'name': 'node_classification'},)" ) assert repr(task) == expected_str, task def test_OnDiskDataset_load_tasks_selectively(): """Test preprocess of OnDiskDataset.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 num_classes = 10 # Generate random graph. yaml_content = gbt.random_homo_graphbolt_graph( test_dir, dataset_name, num_nodes, num_edges, num_classes, ) train_path = os.path.join("set", "train.npy") yaml_content += f""" - name: node_classification num_classes: {num_classes} train_set: - type: null data: - format: numpy path: {train_path} """ yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) # Case1. Test load all tasks. dataset = gb.OnDiskDataset(test_dir).load() assert len(dataset.tasks) == 2 # Case2. Test load tasks selectively. dataset = gb.OnDiskDataset(test_dir).load(tasks="link_prediction") assert len(dataset.tasks) == 1 assert dataset.tasks[0].metadata["name"] == "link_prediction" dataset = gb.OnDiskDataset(test_dir).load(tasks=["link_prediction"]) assert len(dataset.tasks) == 1 assert dataset.tasks[0].metadata["name"] == "link_prediction" # Case3. Test load tasks with non-existent task name. with pytest.warns( GBWarning, match="Below tasks are not found in YAML: {'fake-name'}. Skipped.", ): dataset = gb.OnDiskDataset(test_dir).load(tasks=["fake-name"]) assert len(dataset.tasks) == 0 # Case4. Test load tasks selectively with incorrect task type. with pytest.raises(TypeError): dataset = gb.OnDiskDataset(test_dir).load(tasks=2) dataset = None def test_OnDiskDataset_preprocess_graph_with_single_type(): """Test for graph with single node/edge type.""" with tempfile.TemporaryDirectory() as test_dir: # All metadata fields are specified. dataset_name = "graphbolt_test" num_nodes = 4000 num_edges = 20000 # Generate random edges. nodes = np.repeat(np.arange(num_nodes), 5) neighbors = np.random.randint(0, num_nodes, size=(num_edges)) edges = np.stack([nodes, neighbors], axis=1) # Write into edges/edge.csv os.makedirs(os.path.join(test_dir, "edges/"), exist_ok=True) edges = pd.DataFrame(edges, columns=["src", "dst"]) edges.to_csv( os.path.join(test_dir, "edges/edge.csv"), index=False, header=False, ) # Generate random graph edge-feats. edge_feats = np.random.rand(num_edges, 5) os.makedirs(os.path.join(test_dir, "data/"), exist_ok=True) np.save(os.path.join(test_dir, "data/edge-feat.npy"), edge_feats) # Generate random node-feats. node_feats = np.random.rand(num_nodes, 10) np.save(os.path.join(test_dir, "data/node-feat.npy"), node_feats) yaml_content = f""" dataset_name: {dataset_name} graph: # graph structure and required attributes. nodes: - num: {num_nodes} type: author edges: - type: author:collab:author format: csv path: edges/edge.csv feature_data: - domain: edge type: author:collab:author name: feat format: numpy path: data/edge-feat.npy - domain: node type: author name: feat format: numpy path: data/node-feat.npy """ yaml_file = os.path.join(test_dir, "metadata.yaml") with open(yaml_file, "w") as f: f.write(yaml_content) dataset = gb.OnDiskDataset(test_dir).load() assert dataset.dataset_name == dataset_name graph = dataset.graph assert isinstance(graph, gb.FusedCSCSamplingGraph) assert graph.total_num_nodes == num_nodes assert graph.total_num_edges == num_edges 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 torch.equal(graph.node_type_offset, torch.tensor([0, num_nodes])) assert torch.equal( graph.type_per_edge, torch.zeros(num_edges), ) assert graph.edge_type_to_id == {"author:collab:author": 0} assert graph.node_type_to_id == {"author": 0}