chore: import upstream snapshot with attribution

python/dgl/graphbolt/dataset.py

@@ -0,0 +1,95 @@
+"""GraphBolt Dataset."""
+
+from typing import Dict, List, Union
+
+from .feature_store import FeatureStore
+from .itemset import HeteroItemSet, ItemSet
+from .sampling_graph import SamplingGraph
+
+__all__ = [
+    "Task",
+    "Dataset",
+]
+
+
+class Task:
+    """An abstract task which consists of meta information and
+    Train/Validation/Test Set.
+
+    * meta information
+        The meta information of a task includes any kinds of data that are
+        defined by the user in YAML when instantiating the task.
+
+    * Train/Validation/Test Set
+        The train/validation/test (TVT) set which is used to train the neural
+        networks. We calculate the embeddings based on their respective features
+        and the graph structure, and then utilize the embeddings to optimize the
+        neural network parameters.
+    """
+
+    @property
+    def metadata(self) -> Dict:
+        """Return the task metadata."""
+        raise NotImplementedError
+
+    @property
+    def train_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the training set."""
+        raise NotImplementedError
+
+    @property
+    def validation_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the validation set."""
+        raise NotImplementedError
+
+    @property
+    def test_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the test set."""
+        raise NotImplementedError
+
+
+class Dataset:
+    """An abstract dataset which provides abstraction for accessing the data
+    required for training.
+
+    The data abstraction could be a native CPU memory block, a shared memory
+    block, a file handle of an opened file on disk, a service that provides
+    the API to access the data e.t.c. There are 3 primary components in the
+    dataset:
+
+    * Task
+        A task consists of several meta information and the
+        Train/Validation/Test Set. A dataset could have multiple tasks.
+
+    * Feature Storage
+        A key-value store which stores node/edge/graph features.
+
+    * Graph Topology
+        Graph topology is used by the subgraph sampling algorithm to generate
+        a subgraph.
+    """
+
+    @property
+    def tasks(self) -> List[Task]:
+        """Return the tasks."""
+        raise NotImplementedError
+
+    @property
+    def graph(self) -> SamplingGraph:
+        """Return the graph."""
+        raise NotImplementedError
+
+    @property
+    def feature(self) -> FeatureStore:
+        """Return the feature."""
+        raise NotImplementedError
+
+    @property
+    def dataset_name(self) -> str:
+        """Return the dataset name."""
+        raise NotImplementedError
+
+    @property
+    def all_nodes_set(self) -> Union[ItemSet, HeteroItemSet]:
+        """Return the itemset containing all nodes."""
+        raise NotImplementedError