[GraphBolt] change TVT format of OnDiskDataset (#6076)

python/dgl/graphbolt/impl/ondisk_dataset.py

@@ -15,7 +15,7 @@ import dgl
 
 from ..dataset import Dataset
 from ..itemset import ItemSet, ItemSetDict
-from ..utils import read_data, save_data, tensor_to_tuple
+from ..utils import read_data, save_data
 
 from .csc_sampling_graph import (
     CSCSamplingGraph,
@@ -174,34 +174,36 @@ def preprocess_ondisk_dataset(input_config_path: str) -> str:
             for input_set_per_type, output_set_per_type in zip(
                 intput_set_split, output_set_split
             ):
-                # Always save the feature in numpy format.
-                output_set_per_type["format"] = "numpy"
-                output_set_per_type["path"] = str(
-                    processed_dir_prefix
-                    / input_set_per_type["path"].replace("pt", "npy")
-                )
-                if input_set_per_type["format"] == "numpy":
-                    # If the original format is numpy, just copy the file.
-                    os.makedirs(
-                        dataset_path
-                        / os.path.dirname(output_set_per_type["path"]),
-                        exist_ok=True,
-                    )
-                    shutil.copy(
-                        dataset_path / input_set_per_type["path"],
-                        dataset_path / output_set_per_type["path"],
-                    )
-                else:
-                    # If the original format is not numpy, convert it to numpy.
-                    input_set = read_data(
-                        dataset_path / input_set_per_type["path"],
-                        input_set_per_type["format"],
-                    )
-                    save_data(
-                        input_set,
-                        dataset_path / output_set_per_type["path"],
-                        output_set_per_type["format"],
+                for input_data, output_data in zip(
+                    input_set_per_type["data"], output_set_per_type["data"]
+                ):
+                    # Always save the feature in numpy format.
+                    output_data["format"] = "numpy"
+                    output_data["path"] = str(
+                        processed_dir_prefix
+                        / input_data["path"].replace("pt", "npy")
                     )
+                    if input_data["format"] == "numpy":
+                        # If the original format is numpy, just copy the file.
+                        os.makedirs(
+                            dataset_path / os.path.dirname(output_data["path"]),
+                            exist_ok=True,
+                        )
+                        shutil.copy(
+                            dataset_path / input_data["path"],
+                            dataset_path / output_data["path"],
+                        )
+                    else:
+                        # If the original format is not numpy, convert it to numpy.
+                        input_set = read_data(
+                            dataset_path / input_data["path"],
+                            input_data["format"],
+                        )
+                        save_data(
+                            input_set,
+                            dataset_path / output_data["path"],
+                            output_set_per_type["format"],
+                        )
 
     # 8. Save the output_config.
     output_config_path = dataset_path / "output_config.yaml"
@@ -245,19 +247,25 @@ class OnDiskDataset(Dataset):
             path: edge_data/author-writes-paper-feat.npy
         train_sets:
           - - type: paper # could be null for homogeneous graph.
-              format: numpy
-              in_memory: true # If not specified, default to true.
-              path: set/paper-train.npy
+              data: # multiple data sources could be specified.
+                - format: numpy
+                  in_memory: true # If not specified, default to true.
+                  path: set/paper-train-src.npy
+                - format: numpy
+                  in_memory: false
+                  path: set/paper-train-dst.npy
         validation_sets:
           - - type: paper
-              format: numpy
-              in_memory: true
-              path: set/paper-validation.npy
+              data:
+                - format: numpy
+                  in_memory: true
+                  path: set/paper-validation.npy
         test_sets:
           - - type: paper
-              format: numpy
-              in_memory: true
-              path: set/paper-test.npy
+              data:
+                - format: numpy
+                  in_memory: true
+                  path: set/paper-test.npy
 
     Parameters
     ----------
@@ -347,16 +355,21 @@ class OnDiskDataset(Dataset):
                 assert (
                     len(tvt_set) == 1
                 ), "Only one TVT set is allowed if type is not specified."
-                data = read_data(
-                    tvt_set[0].path, tvt_set[0].format, tvt_set[0].in_memory
+                ret.append(
+                    ItemSet(
+                        tuple(
+                            read_data(data.path, data.format, data.in_memory)
+                            for data in tvt_set[0].data
+                        )
+                    )
                 )
-                ret.append(ItemSet(tensor_to_tuple(data)))
             else:
                 data = {}
                 for tvt in tvt_set:
                     data[tvt.type] = ItemSet(
-                        tensor_to_tuple(
-                            read_data(tvt.path, tvt.format, tvt.in_memory)
+                        tuple(
+                            read_data(data.path, data.format, data.in_memory)
+                            for data in tvt.data
                         )
                     )
                 ret.append(ItemSetDict(data))

python/dgl/graphbolt/impl/ondisk_metadata.py

@@ -8,6 +8,7 @@ import pydantic
 
 __all__ = [
     "OnDiskFeatureDataFormat",
+    "OnDiskTVTSetData",
     "OnDiskTVTSet",
     "OnDiskFeatureDataDomain",
     "OnDiskFeatureData",
@@ -24,15 +25,21 @@ class OnDiskFeatureDataFormat(str, Enum):
     NUMPY = "numpy"
 
 
-class OnDiskTVTSet(pydantic.BaseModel):
-    """Train-Validation-Test set."""
+class OnDiskTVTSetData(pydantic.BaseModel):
+    """Train-Validation-Test set data."""
 
-    type: Optional[str] = None
     format: OnDiskFeatureDataFormat
     in_memory: Optional[bool] = True
     path: str
 
 
+class OnDiskTVTSet(pydantic.BaseModel):
+    """Train-Validation-Test set."""
+
+    type: Optional[str] = None
+    data: List[OnDiskTVTSetData]
+
+
 class OnDiskFeatureDataDomain(str, Enum):
     """Enum of feature data domain."""
 

python/dgl/graphbolt/utils/internal.py

@@ -45,9 +45,3 @@ def save_data(data, path, fmt):
         np.save(path, data)
     elif fmt == "torch":
         torch.save(data, path)
-
-
-def tensor_to_tuple(data):
-    """Split a torch.Tensor in column-wise to a tuple."""
-    assert isinstance(data, torch.Tensor), "data must be a torch.Tensor"
-    return tuple(data.t())