Repository: pyg-team/pytorch_geometric
Pull request #8110 (2023-10-02 05:54:54)

Add TensorFrame support for`Data`
Supports `TensorFrame` attributes in `Data`.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 08e028a28..bdf1eae21 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `torch-frame` support ([#8110](https://github.com/pyg-team/pytorch_geometric/pull/8110))
 - Added the `DistLoader` base class ([#8079](https://github.com/pyg-team/pytorch_geometric/pull/8079))
 - Added `HyperGraphData` to support hypergraphs ([#7611](https://github.com/pyg-team/pytorch_geometric/pull/7611))
 - Added the `PCQM4Mv2` dataset as a reference implementation for `OnDiskDataset` ([#8102](https://github.com/pyg-team/pytorch_geometric/pull/8102))
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 935c3b43a..1a855bf28 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -438,9 +438,15 @@ class Data(BaseData, FeatureStore, GraphStore):
             :obj:`[num_nodes, num_dimensions]`. (default: :obj:`None`)
         **kwargs (optional): Additional attributes.
     """
-    def __init__(self, x: OptTensor = None, edge_index: OptTensor = None,
-                 edge_attr: OptTensor = None, y: OptTensor = None,
-                 pos: OptTensor = None, **kwargs):
+    def __init__(
+        self,
+        x: Optional[Tensor] = None,
+        edge_index: OptTensor = None,
+        edge_attr: OptTensor = None,
+        y: OptTensor = None,
+        pos: OptTensor = None,
+        **kwargs,
+    ):
         # `Data` doesn't support group_name, so we need to adjust `TensorAttr`
         # accordingly here to avoid requiring `group_name` to be set:
         super().__init__(tensor_attr_cls=DataTensorAttr)
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index cc6dda142..9860165af 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -22,7 +22,12 @@ import torch
 from torch import Tensor
 
 from torch_geometric.data.view import ItemsView, KeysView, ValuesView
-from torch_geometric.typing import EdgeType, NodeType, SparseTensor
+from torch_geometric.typing import (
+    EdgeType,
+    NodeType,
+    SparseTensor,
+    TensorFrame,
+)
 from torch_geometric.utils import (
     coalesce,
     contains_isolated_nodes,
@@ -31,7 +36,7 @@ from torch_geometric.utils import (
     sort_edge_index,
 )
 
-N_KEYS = {'x', 'feat', 'pos', 'batch', 'node_type', 'n_id'}
+N_KEYS = {'x', 'feat', 'pos', 'batch', 'node_type', 'n_id', 'tf'}
 E_KEYS = {'edge_index', 'edge_weight', 'edge_attr', 'edge_type', 'e_id'}
 
 
@@ -304,6 +309,8 @@ class NodeStorage(BaseStorage):
             if isinstance(value, np.ndarray) and key in N_KEYS:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
                 return value.shape[cat_dim]
+            if isinstance(value, TensorFrame) and key in N_KEYS:
+                return value.num_rows
         for key, value in self.items():
             if isinstance(value, Tensor) and 'node' in key:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
@@ -311,6 +318,8 @@ class NodeStorage(BaseStorage):
             if isinstance(value, np.ndarray) and 'node' in key:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
                 return value.shape[cat_dim]
+            if isinstance(value, TensorFrame) and 'node' in key:
+                return value.num_rows
         if 'adj' in self and isinstance(self.adj, SparseTensor):
             return self.adj.size(0)
         if 'adj_t' in self and isinstance(self.adj_t, SparseTensor):
@@ -335,10 +344,14 @@ class NodeStorage(BaseStorage):
 
     @property
     def num_node_features(self) -> int:
-        if 'x' in self and isinstance(self.x, (Tensor, np.ndarray)):
+        if 'x' in self and isinstance(self.x, Tensor):
+            return 1 if self.x.dim() == 1 else self.x.size(-1)
+        if 'x' in self and isinstance(self.x, np.ndarray):
             return 1 if self.x.ndim == 1 else self.x.shape[-1]
         if 'x' in self and isinstance(self.x, SparseTensor):
             return 1 if self.x.dim() == 1 else self.x.size(-1)
+        if 'x' in self and isinstance(self.x, TensorFrame):
+            return self.x.num_cols
         return 0
 
     @property
@@ -356,7 +369,8 @@ class NodeStorage(BaseStorage):
 
         value = self[key]
 
-        if isinstance(value, (list, tuple)) and len(value) == self.num_nodes:
+        if (isinstance(value, (list, tuple, TensorFrame))
+                and len(value) == self.num_nodes):
             self._cached_attr[AttrType.NODE].add(key)
             return True
 
@@ -423,13 +437,23 @@ class EdgeStorage(BaseStorage):
         if 'num_edges' in self:
             return self['num_edges']
         for key, value in self.items():
-            if isinstance(value, (Tensor, np.ndarray)) and key in E_KEYS:
+            if isinstance(value, Tensor) and key in E_KEYS:
+                cat_dim = self._parent().__cat_dim__(key, value, self)
+                return value.size(cat_dim)
+            if isinstance(value, Tensor) and key in E_KEYS:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
                 return value.shape[cat_dim]
+            if isinstance(value, TensorFrame) and key in E_KEYS:
+                return value.num_rows
         for key, value in self.items():
-            if isinstance(value, (Tensor, np.ndarray)) and 'edge' in key:
+            if isinstance(value, Tensor) and 'edge' in key:
+                cat_dim = self._parent().__cat_dim__(key, value, self)
+                return value.size(cat_dim)
+            if isinstance(value, np.ndarray) and 'edge' in key:
                 cat_dim = self._parent().__cat_dim__(key, value, self)
                 return value.shape[cat_dim]
+            if isinstance(value, TensorFrame) and 'edge' in key:
+                return value.num_rows
         for value in self.values('adj', 'adj_t'):
             if isinstance(value, SparseTensor):
                 return value.nnz()
@@ -439,8 +463,9 @@ class EdgeStorage(BaseStorage):
 
     @property
     def num_edge_features(self) -> int:
-        if ('edge_attr' in self and isinstance(self.edge_attr,
-                                               (Tensor, np.ndarray))):
+        if 'edge_attr' in self and isinstance(self.edge_attr, Tensor):
+            return 1 if self.edge_attr.dim() == 1 else self.edge_attr.size(-1)
+        if 'edge_attr' in self and isinstance(self.edge_attr, np.ndarray):
             return 1 if self.edge_attr.ndim == 1 else self.edge_attr.shape[-1]
         return 0
 
@@ -475,7 +500,8 @@ class EdgeStorage(BaseStorage):
 
         value = self[key]
 
-        if isinstance(value, (list, tuple)) and len(value) == self.num_edges:
+        if (isinstance(value, (list, tuple, TensorFrame))
+                and len(value) == self.num_edges):
             self._cached_attr[AttrType.EDGE].add(key)
             return True
 
@@ -609,7 +635,8 @@ class GlobalStorage(NodeStorage, EdgeStorage):
 
         value = self[key]
 
-        if isinstance(value, (list, tuple)) and len(value) == self.num_nodes:
+        if (isinstance(value, (list, tuple, TensorFrame))
+                and len(value) == self.num_nodes):
             self._cached_attr[AttrType.NODE].add(key)
             return True
 
@@ -655,7 +682,8 @@ class GlobalStorage(NodeStorage, EdgeStorage):
 
         value = self[key]
 
-        if isinstance(value, (list, tuple)) and len(value) == self.num_edges:
+        if (isinstance(value, (list, tuple, TensorFrame))
+                and len(value) == self.num_edges):
             self._cached_attr[AttrType.EDGE].add(key)
             return True
 
diff --git a/torch_geometric/typing.py b/torch_geometric/typing.py
index 190e15f50..bdfe404eb 100644
--- a/torch_geometric/typing.py
+++ b/torch_geometric/typing.py
@@ -211,6 +211,17 @@ except Exception as e:
             raise ImportError("'masked_select_nnz' requires 'torch-sparse'")
 
 
+try:
+    import torch_frame  # noqa
+    WITH_TORCH_FRAME = True
+    from torch_frame import TensorFrame
+except Exception:
+    WITH_TORCH_FRAME = False
+
+    class TensorFrame:
+        pass
+
+
 try:
     import intel_extension_for_pytorch  # noqa
     WITH_IPEX = True
diff --git a/torch_geometric/utils/mask.py b/torch_geometric/utils/mask.py
index 95ab71ce7..d9b9228bf 100644
--- a/torch_geometric/utils/mask.py
+++ b/torch_geometric/utils/mask.py
@@ -3,6 +3,8 @@ from typing import Optional
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import TensorFrame
+
 
 def mask_select(src: Tensor, dim: int, mask: Tensor) -> Tensor:
     r"""Returns a new tensor which masks the :obj:`src` tensor along the
@@ -15,6 +17,12 @@ def mask_select(src: Tensor, dim: int, mask: Tensor) -> Tensor:
             index with.
     """
     assert mask.dim() == 1
+
+    if not torch.jit.is_scripting():
+        if isinstance(src, TensorFrame):
+            assert dim == 0 and src.num_rows == mask.numel()
+            return src[mask]
+
     assert src.size(dim) == mask.numel()
     dim = dim + src.dim() if dim < 0 else dim
     assert dim >= 0 and dim < src.dim()
diff --git a/torch_geometric/utils/select.py b/torch_geometric/utils/select.py
index 54d0a3961..fede7fb2f 100644
--- a/torch_geometric/utils/select.py
+++ b/torch_geometric/utils/select.py
@@ -3,12 +3,16 @@ from typing import Any, List, Union
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import TensorFrame
 from torch_geometric.utils.mask import mask_select
 from torch_geometric.utils.sparse import is_torch_sparse_tensor
 
 
-def select(src: Union[Tensor, List[Any]], index_or_mask: Tensor,
-           dim: int) -> Union[Tensor, List[Any]]:
+def select(
+    src: Union[Tensor, List[Any], TensorFrame],
+    index_or_mask: Tensor,
+    dim: int,
+) -> Union[Tensor, List[Any]]:
     r"""Selects the input tensor or input list according to a given index or
     mask vector.
 
@@ -29,6 +33,12 @@ def select(src: Union[Tensor, List[Any]], index_or_mask: Tensor,
             return [src[i] for i, m in enumerate(index_or_mask) if m]
         return [src[i] for i in index_or_mask]
 
+    if isinstance(src, TensorFrame):
+        assert dim == 0
+        if index_or_mask.dtype == torch.bool:
+            return mask_select(src, dim, index_or_mask)
+        return src[index_or_mask]
+
     raise ValueError(f"Encountered invalid input type (got '{type(src)}')")
 
 


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index e52f5f8c6..e762b179d 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -8,6 +8,7 @@ import torch_geometric
 from torch_geometric.data import Data
 from torch_geometric.data.storage import AttrType
 from torch_geometric.testing import withPackage
+from torch_geometric.typing import TensorFrame
 
 
 def test_data():
@@ -481,3 +482,49 @@ def test_data_generate_ids():
     assert len(data) == 4
     assert data.n_id.tolist() == [0, 1, 2]
     assert data.e_id.tolist() == [0, 1, 2, 3, 4]
+
+
+def get_fake_tensor_frame(num_rows: int) -> TensorFrame:
+    import torch_frame
+
+    feat_dict = {
+        torch_frame.categorical: torch.randint(0, 3, size=(num_rows, 3)),
+        torch_frame.numerical: torch.randn(size=(num_rows, 2)),
+    }
+    col_names_dict = {
+        torch_frame.categorical: ['a', 'b', 'c'],
+        torch_frame.numerical: ['x', 'y'],
+    }
+    y = torch.randn(num_rows)
+
+    return TensorFrame(
+        feat_dict=feat_dict,
+        col_names_dict=col_names_dict,
+        y=y,
+    )
+
+
+@withPackage('torch_frame')
+def test_data_with_tensor_frame():
+    tf = get_fake_tensor_frame(num_rows=10)
+    data = Data(tf=tf, edge_index=torch.randint(0, 10, size=(2, 20)))
+
+    # Test basic attributes:
+    assert data.is_node_attr('x')
+    assert data.num_nodes == tf.num_rows
+    assert data.num_edges == 20
+    assert data.num_node_features == tf.num_cols
+
+    # Test subgraph:
+    index = torch.tensor([1, 2, 3])
+    sub_data = data.subgraph(index)
+    assert sub_data.num_nodes == 3
+    for key, value in sub_data.tf.feat_dict.items():
+        assert torch.allclose(value, tf.feat_dict[key][index])
+
+    mask = torch.tensor(
+        [False, True, True, True, False, False, False, False, False, False])
+    data_sub = data.subgraph(mask)
+    assert data_sub.num_nodes == 3
+    for key, value in sub_data.tf.feat_dict.items():
+        assert torch.allclose(value, tf.feat_dict[key][mask])