[Graphbolt] Combine data block classes to a single file. (#6233)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>

python/dgl/graphbolt/__init__.py

@@ -14,10 +14,8 @@ from .feature_fetcher import *
 from .feature_store import *
 from .impl import *
 from .itemset import *
-from .link_prediction_block import *
 from .minibatch_sampler import *
 from .negative_sampler import *
-from .node_classification_block import *
 from .sampled_subgraph import *
 from .subgraph_sampler import *
 from .utils import unique_and_compact, unique_and_compact_node_pairs

python/dgl/graphbolt/data_block.py

@@ -7,7 +7,7 @@ import torch
 
 from .sampled_subgraph import SampledSubgraph
 
-__all__ = ["DataBlock"]
+__all__ = ["DataBlock", "NodeClassificationBlock", "LinkPredictionBlock"]
 
 
 @dataclass
@@ -44,3 +44,94 @@ class DataBlock:
     - If `input_nodes` is a dictionary: The keys should be node type and the
       value should be corresponding heterogeneous node id.
     """
+
+
+@dataclass
+class NodeClassificationBlock(DataBlock):
+    r"""A subclass of 'UnifiedDataStruct', specialized for handling node level
+    tasks."""
+
+    seed_node: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Representation of seed nodes used for sampling in the graph.
+    - If `seed_node` is a tensor: It indicates the graph is homogeneous.
+    - If `seed_node` is a dictionary: The keys should be node type and the
+      value should be corresponding heterogeneous node ids.
+    """
+
+    label: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Labels associated with seed nodes in the graph.
+    - If `label` is a tensor: It indicates the graph is homogeneous.
+    - If `label` is a dictionary: The keys should be node type and the
+      value should be corresponding node labels to given 'seed_node'.
+    """
+
+
+@dataclass
+class LinkPredictionBlock(DataBlock):
+    r"""A subclass of 'UnifiedDataStruct', specialized for handling edge level
+    tasks."""
+
+    node_pair: Union[
+        Tuple[torch.Tensor, torch.Tensor],
+        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
+    ] = None
+    """
+    Representation of seed node pairs utilized in link prediction tasks.
+    - If `node_pair` is a tuple: It indicates a homogeneous graph where each
+      tuple contains two tensors representing source-destination node pairs.
+    - If `node_pair` is a dictionary: The keys should be edge type, and the
+      value should be a tuple of tensors representing node pairs of the given
+      type.
+    """
+
+    label: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Labels associated with the link prediction task.
+    - If `label` is a tensor: It indicates a homogeneous graph. The value are
+      edge labels corresponding to given 'node_pair'.
+    - If `label` is a dictionary: The keys should be edge type, and the value
+      should correspond to given 'node_pair'.
+    """
+
+    negative_head: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Representation of negative samples for the head nodes in the link
+    prediction task.
+    - If `negative_head` is a tensor: It indicates a homogeneous graph.
+    - If `negative_head` is a dictionary: The key should be edge type, and the
+      value should correspond to the negative samples for head nodes of the
+      given type.
+    """
+
+    negative_tail: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Representation of negative samples for the tail nodes in the link
+    prediction task.
+    - If `negative_tail` is a tensor: It indicates a homogeneous graph.
+    - If `negative_tail` is a dictionary: The key should be edge type, and the
+      value should correspond to the negative samples for head nodes of the
+      given type.
+    """
+
+    compacted_node_pair: Union[
+        Tuple[torch.Tensor, torch.Tensor],
+        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
+    ] = None
+    """
+    Representation of compacted node pairs corresponding to 'node_pair', where
+    all node ids inside are compacted.
+    """
+
+    compacted_negative_head: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Representation of compacted nodes corresponding to 'negative_head', where
+    all node ids inside are compacted.
+    """
+
+    compacted_negative_tail: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
+    """
+    Representation of compacted nodes corresponding to 'negative_tail', where
+    all node ids inside are compacted.
+    """

python/dgl/graphbolt/link_prediction_block.py

-"""Unified data structure for input and ouput of all the stages in loading
-process, especially for edge level task."""
-
-from dataclasses import dataclass
-from typing import Dict, Tuple, Union
-
-import torch
-
-from .data_block import DataBlock
-
-
-@dataclass
-class LinkPredictionBlock(DataBlock):
-    r"""A subclass of 'UnifiedDataStruct', specialized for handling edge level
-    tasks."""
-
-    node_pair: Union[
-        Tuple[torch.Tensor, torch.Tensor],
-        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
-    ] = None
-    """
-    Representation of seed node pairs utilized in link prediction tasks.
-    - If `node_pair` is a tuple: It indicates a homogeneous graph where each
-      tuple contains two tensors representing source-destination node pairs.
-    - If `node_pair` is a dictionary: The keys should be edge type, and the
-      value should be a tuple of tensors representing node pairs of the given
-      type.
-    """
-
-    label: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Labels associated with the link prediction task.
-    - If `label` is a tensor: It indicates a homogeneous graph. The value are
-      edge labels corresponding to given 'node_pair'.
-    - If `label` is a dictionary: The keys should be edge type, and the value
-      should correspond to given 'node_pair'.
-    """
-
-    negative_head: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Representation of negative samples for the head nodes in the link
-    prediction task.
-    - If `negative_head` is a tensor: It indicates a homogeneous graph.
-    - If `negative_head` is a dictionary: The key should be edge type, and the
-      value should correspond to the negative samples for head nodes of the
-      given type.
-    """
-
-    negative_tail: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Representation of negative samples for the tail nodes in the link
-    prediction task.
-    - If `negative_tail` is a tensor: It indicates a homogeneous graph.
-    - If `negative_tail` is a dictionary: The key should be edge type, and the
-      value should correspond to the negative samples for head nodes of the
-      given type.
-    """
-
-    compacted_node_pair: Union[
-        Tuple[torch.Tensor, torch.Tensor],
-        Dict[str, Tuple[torch.Tensor, torch.Tensor]],
-    ] = None
-    """
-    Representation of compacted node pairs corresponding to 'node_pair', where
-    all node ids inside are compacted.
-    """
-
-    compacted_negative_head: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Representation of compacted nodes corresponding to 'negative_head', where
-    all node ids inside are compacted.
-    """
-
-    compacted_negative_tail: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Representation of compacted nodes corresponding to 'negative_tail', where
-    all node ids inside are compacted.
-    """

python/dgl/graphbolt/node_classification_block.py

-"""Unified data structure for input and ouput of all the stages in loading
-process, especially for node level task."""
-
-from dataclasses import dataclass
-from typing import Dict, Union
-
-import torch
-
-from .data_block import DataBlock
-
-
-@dataclass
-class NodeClassificationBlock(DataBlock):
-    r"""A subclass of 'UnifiedDataStruct', specialized for handling node level
-    tasks."""
-
-    seed_node: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Representation of seed nodes used for sampling in the graph.
-    - If `seed_node` is a tensor: It indicates the graph is homogeneous.
-    - If `seed_node` is a dictionary: The keys should be node type and the
-      value should be corresponding heterogeneous node ids.
-    """
-
-    label: Union[torch.Tensor, Dict[str, torch.Tensor]] = None
-    """
-    Labels associated with seed nodes in the graph.
-    - If `label` is a tensor: It indicates the graph is homogeneous.
-    - If `label` is a dictionary: The keys should be node type and the
-      value should be corresponding node labels to given 'seed_node'.
-    """

python/dgl/graphbolt/subgraph_sampler.py

@@ -5,8 +5,7 @@ from typing import Dict
 
 from torchdata.datapipes.iter import Mapper
 
-from .link_prediction_block import LinkPredictionBlock
-from .node_classification_block import NodeClassificationBlock
+from .data_block import LinkPredictionBlock, NodeClassificationBlock
 from .utils import unique_and_compact