Repository: pyg-team/pytorch_geometric
Pull request #9293 (2024-05-05 08:57:10)

Integrate `Index` into common data utils


Patch:
diff --git a/torch_geometric/data/collate.py b/torch_geometric/data/collate.py
index c8a8f5e31..a36d3e95f 100644
--- a/torch_geometric/data/collate.py
+++ b/torch_geometric/data/collate.py
@@ -16,7 +16,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage, NodeStorage
 from torch_geometric.edge_index import SortOrder
@@ -185,7 +185,7 @@ def _collate(
 
         out = None
         if (torch.utils.data.get_worker_info() is not None
-                and not isinstance(elem, EdgeIndex)):
+                and not isinstance(elem, (Index, EdgeIndex))):
             # Write directly into shared memory to avoid an extra copy:
             numel = sum(value.numel() for value in values)
             if torch_geometric.typing.WITH_PT20:
@@ -204,6 +204,11 @@ def _collate(
 
         value = torch.cat(values, dim=cat_dim or 0, out=out)
 
+        if increment and isinstance(value, Index) and values[0].is_sorted:
+            # Check whether the whole `Index` is sorted:
+            if (value.diff() >= 0).all():
+                value._is_sorted = True
+
         if increment and isinstance(value, EdgeIndex) and values[0].is_sorted:
             # Check whether the whole `EdgeIndex` is sorted by row:
             if values[0].is_sorted_by_row and (value[0].diff() >= 0).all():
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 8c8c6e528..5d413e971 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -31,6 +31,7 @@ from torch_geometric.data.storage import (
     NodeStorage,
 )
 from torch_geometric.deprecation import deprecated
+from torch_geometric.index import Index
 from torch_geometric.typing import (
     EdgeTensorType,
     EdgeType,
@@ -653,6 +654,8 @@ class Data(BaseData, FeatureStore, GraphStore):
 
     def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
         if 'batch' in key and isinstance(value, Tensor):
+            if isinstance(value, Index):
+                return value.get_dim_size()
             return int(value.max()) + 1
         elif 'index' in key or key == 'face':
             return self.num_nodes
diff --git a/torch_geometric/data/database.py b/torch_geometric/data/database.py
index 5f93a322d..9f84d6ee6 100644
--- a/torch_geometric/data/database.py
+++ b/torch_geometric/data/database.py
@@ -9,7 +9,7 @@ import torch
 from torch import Tensor
 from tqdm import tqdm
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.edge_index import SortOrder
 from torch_geometric.utils.mixin import CastMixin
 
@@ -18,9 +18,17 @@ from torch_geometric.utils.mixin import CastMixin
 class TensorInfo(CastMixin):
     dtype: torch.dtype
     size: Tuple[int, ...] = (-1, )
+    is_index: bool = False
     is_edge_index: bool = False
 
     def __post_init__(self) -> None:
+        if self.is_index and self.is_edge_index:
+            raise ValueError("Tensor cannot be a 'Index' and 'EdgeIndex' "
+                             "tensor at the same time")
+
+        if self.is_index:
+            self.size = (-1, )
+
         if self.is_edge_index:
             self.size = (2, -1)
 
@@ -32,7 +40,8 @@ def maybe_cast_to_tensor_info(value: Any) -> Union[Any, TensorInfo]:
         return value
     if 'dtype' not in value:
         return value
-    if len(set(value.keys()) | {'dtype', 'size', 'is_edge_index'}) != 3:
+    valid_keys = {'dtype', 'size', 'is_index', 'is_edge_index'}
+    if len(set(value.keys()) | valid_keys) != len(valid_keys):
         return value
     return TensorInfo.cast(value)
 
@@ -452,10 +461,22 @@ class SQLiteDatabase(Database):
             if isinstance(col, Tensor) and not isinstance(schema, TensorInfo):
                 self.schema[key] = schema = TensorInfo(
                     col.dtype,
+                    is_index=isinstance(col, Index),
                     is_edge_index=isinstance(col, EdgeIndex),
                 )
 
-            if isinstance(schema, TensorInfo) and schema.is_edge_index:
+            if isinstance(schema, TensorInfo) and schema.is_index:
+                assert isinstance(col, Index)
+
+                meta = torch.tensor([
+                    col.dim_size if col.dim_size is not None else -1,
+                    col.is_sorted,
+                ], dtype=torch.long)
+
+                out.append(meta.numpy().tobytes() +
+                           col.as_tensor().numpy().tobytes())
+
+            elif isinstance(schema, TensorInfo) and schema.is_edge_index:
                 assert isinstance(col, EdgeIndex)
 
                 num_rows, num_cols = col.sparse_size()
@@ -491,7 +512,23 @@ class SQLiteDatabase(Database):
         for i, (key, schema) in enumerate(self.schema.items()):
             value = row[i]
 
-            if isinstance(schema, TensorInfo) and schema.is_edge_index:
+            if isinstance(schema, TensorInfo) and schema.is_index:
+                meta = torch.frombuffer(value[:16], dtype=torch.long).tolist()
+                dim_size = meta[0] if meta[0] >= 0 else None
+                is_sorted = meta[1] > 0
+
+                if len(value) > 16:
+                    tensor = torch.frombuffer(value[16:], dtype=schema.dtype)
+                else:
+                    tensor = torch.empty(0, dtype=schema.dtype)
+
+                out_dict[key] = Index(
+                    tensor.view(*schema.size),
+                    dim_size=dim_size,
+                    is_sorted=is_sorted,
+                )
+
+            elif isinstance(schema, TensorInfo) and schema.is_edge_index:
                 meta = torch.frombuffer(value[:32], dtype=torch.long).tolist()
                 num_rows = meta[0] if meta[0] >= 0 else None
                 num_cols = meta[1] if meta[1] >= 0 else None
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index 3f54fdca6..f0e0b242a 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -10,6 +10,7 @@ import torch
 from torch import Tensor
 from typing_extensions import Self
 
+from torch_geometric import Index
 from torch_geometric.data import EdgeAttr, FeatureStore, GraphStore, TensorAttr
 from torch_geometric.data.data import BaseData, Data, size_repr, warn_or_raise
 from torch_geometric.data.graph_store import EdgeLayout
@@ -346,6 +347,8 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                 store: Optional[NodeOrEdgeStorage] = None, *args,
                 **kwargs) -> Any:
         if 'batch' in key and isinstance(value, Tensor):
+            if isinstance(value, Index):
+                return value.get_dim_size()
             return int(value.max()) + 1
         elif isinstance(store, EdgeStorage) and 'index' in key:
             return torch.tensor(store.size()).view(2, 1)
diff --git a/torch_geometric/data/separate.py b/torch_geometric/data/separate.py
index d08454e23..2910b6679 100644
--- a/torch_geometric/data/separate.py
+++ b/torch_geometric/data/separate.py
@@ -3,7 +3,7 @@ from typing import Any, Type, TypeVar
 
 from torch import Tensor
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage
 from torch_geometric.typing import SparseTensor, TensorFrame
@@ -76,6 +76,11 @@ def _separate(
         value = narrow(values, cat_dim or 0, start, end - start)
         value = value.squeeze(0) if cat_dim is None else value
 
+        if isinstance(values, Index) and values._cat_metadata is not None:
+            # Reconstruct original `Index` metadata:
+            value._dim_size = values._cat_metadata.dim_size[idx]
+            value._is_sorted = values._cat_metadata.is_sorted[idx]
+
         if isinstance(values, EdgeIndex) and values._cat_metadata is not None:
             # Reconstruct original `EdgeIndex` metadata:
             value._sparse_size = values._cat_metadata.sparse_size[idx]


Test patch:
diff --git a/test/data/test_batch.py b/test/data/test_batch.py
index e23478ad5..128536e6e 100644
--- a/test/data/test_batch.py
+++ b/test/data/test_batch.py
@@ -5,7 +5,7 @@ import pytest
 import torch
 
 import torch_geometric
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data import Batch, Data, HeteroData
 from torch_geometric.testing import get_random_edge_index, withPackage
 from torch_geometric.typing import SparseTensor
@@ -108,6 +108,31 @@ def test_batch_basic():
     torch_geometric.set_debug(True)
 
 
+def test_index():
+    index1 = Index([0, 1, 1, 2], dim_size=3, is_sorted=True)
+    index2 = Index([0, 1, 1, 2, 2, 3], dim_size=4, is_sorted=True)
+
+    data1 = Data(index=index1, num_nodes=3)
+    data2 = Data(index=index2, num_nodes=4)
+
+    batch = Batch.from_data_list([data1, data2])
+
+    assert len(batch) == 2
+    assert batch.batch.equal(torch.tensor([0, 0, 0, 1, 1, 1, 1]))
+    assert batch.ptr.equal(torch.tensor([0, 3, 7]))
+    assert isinstance(batch.index, Index)
+    assert batch.index.equal(torch.tensor([0, 1, 1, 2, 3, 4, 4, 5, 5, 6]))
+    assert batch.index.dim_size == 7
+    assert batch.index.is_sorted
+
+    for i, index in enumerate([index1, index2]):
+        data = batch[i]
+        assert isinstance(data.index, Index)
+        assert data.index.equal(index)
+        assert data.index.dim_size == index.dim_size
+        assert data.index.is_sorted == index.is_sorted
+
+
 def test_edge_index():
     edge_index1 = EdgeIndex(
         [[0, 1, 1, 2], [1, 0, 2, 1]],
diff --git a/test/data/test_database.py b/test/data/test_database.py
index 4cb13ecce..cb009da04 100644
--- a/test/data/test_database.py
+++ b/test/data/test_database.py
@@ -4,7 +4,7 @@ import os.path as osp
 import pytest
 import torch
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data import Data, RocksDatabase, SQLiteDatabase
 from torch_geometric.data.database import TensorInfo
 from torch_geometric.profile import benchmark
@@ -149,6 +149,36 @@ def test_database_schema(tmp_path, Database):
     db.close()
 
 
+@pytest.mark.parametrize('Database', AVAILABLE_DATABASES)
+def test_index(tmp_path, Database):
+    kwargs = dict(name='test_table') if Database == SQLiteDatabase else {}
+
+    path = osp.join(tmp_path, 'tuple_storage.db')
+    schema = dict(dtype=torch.long, is_index=True)
+    db = Database(path, schema=schema, **kwargs)
+    assert db.schema == {
+        0: TensorInfo(dtype=torch.long, is_index=True),
+    }
+
+    index1 = Index([0, 1, 1, 2], dim_size=3, is_sorted=True)
+    index2 = Index([0, 1, 1, 2, 2, 3], dim_size=None, is_sorted=True)
+    index3 = Index([], dtype=torch.long)
+
+    db.insert(0, index1)
+    db.multi_insert([1, 2], [index2, index3])
+
+    out1 = db.get(0)
+    out2, out3 = db.multi_get([1, 2])
+
+    for out, index in zip([out1, out2, out3], [index1, index2, index3]):
+        assert index.equal(out)
+        assert index.dtype == out.dtype
+        assert index.dim_size == out.dim_size
+        assert index.is_sorted == out.is_sorted
+
+    db.close()
+
+
 @pytest.mark.parametrize('Database', AVAILABLE_DATABASES)
 def test_edge_index(tmp_path, Database):
     kwargs = dict(name='test_table') if Database == SQLiteDatabase else {}
diff --git a/test/data/test_dataset.py b/test/data/test_dataset.py
index a7bbe048f..5b23412a3 100644
--- a/test/data/test_dataset.py
+++ b/test/data/test_dataset.py
@@ -3,7 +3,7 @@ import copy
 import pytest
 import torch
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data import Data, HeteroData, InMemoryDataset
 from torch_geometric.datasets import KarateClub
 from torch_geometric.testing import withPackage
@@ -122,6 +122,30 @@ def test_stored_hetero_in_memory_dataset(tmp_path):
     assert dataset[1]['paper'].num_nodes == 4
 
 
+def test_index(tmp_path):
+    index1 = Index([0, 1, 1, 2], dim_size=3, is_sorted=True)
+    index2 = Index([0, 1, 1, 2, 2, 3], dim_size=4, is_sorted=True)
+
+    data1 = Data(batch=index1)
+    data2 = Data(batch=index2)
+
+    dataset = MyTestDataset([data1, data2])
+    assert len(dataset) == 2
+    for data, index in zip(dataset, [index1, index2]):
+        assert isinstance(data.batch, Index)
+        assert data.batch.equal(index)
+        assert data.batch.dim_size == index.dim_size
+        assert data.batch.is_sorted == index.is_sorted
+
+    dataset = MyStoredTestDataset(tmp_path, [data1, data2])
+    assert len(dataset) == 2
+    for data, index in zip(dataset, [index1, index2]):
+        assert isinstance(data.batch, Index)
+        assert data.batch.equal(index)
+        assert data.batch.dim_size == index.dim_size
+        assert data.batch.is_sorted == index.is_sorted
+
+
 def test_edge_index(tmp_path):
     edge_index1 = EdgeIndex(
         [[0, 1, 1, 2], [1, 0, 2, 1]],
diff --git a/test/loader/test_dataloader.py b/test/loader/test_dataloader.py
index cb359001a..2464c32a1 100644
--- a/test/loader/test_dataloader.py
+++ b/test/loader/test_dataloader.py
@@ -5,7 +5,7 @@ from collections import namedtuple
 import pytest
 import torch
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data import Data, HeteroData, OnDiskDataset
 from torch_geometric.loader import DataLoader
 from torch_geometric.testing import (
@@ -189,6 +189,27 @@ def test_heterogeneous_dataloader(num_workers):
             assert id(batch) == id(store._parent())
 
 
+@pytest.mark.parametrize('num_workers', num_workers_list)
+def test_index_dataloader(num_workers):
+    index1 = Index([0, 1, 1, 2], dim_size=3, is_sorted=True)
+    index2 = Index([0, 1, 1, 2, 2, 3], dim_size=4, is_sorted=True)
+
+    data1 = Data(index=index1, num_nodes=3)
+    data2 = Data(index=index2, num_nodes=4)
+
+    loader = DataLoader(
+        [data1, data2, data1, data2],
+        batch_size=2,
+        num_workers=num_workers,
+    )
+    assert len(loader) == 2
+
+    for batch in loader:
+        assert isinstance(batch.index, Index)
+        assert batch.index.dim_size == 7
+        assert batch.index.is_sorted
+
+
 @pytest.mark.parametrize('num_workers', num_workers_list)
 @pytest.mark.parametrize('sort_order', [None, 'row', 'col'])
 def test_edge_index_dataloader(num_workers, sort_order):